import json, math
import numpy as np
from pk_regularizer import poisson_density, poisson_nll, estimate_attractor, sample_poisson_phases


def hist_kl(phases, x, bins=48):
    h, _ = np.histogram(phases, bins=bins, range=(-np.pi, np.pi), density=False)
    p = (h + 1e-6) / (h.sum() + bins * 1e-6)
    centers = -np.pi + (np.arange(bins) + .5) * 2*np.pi/bins
    q = poisson_density(centers, x)
    q = q / q.sum()
    return float(np.sum(p * np.log(p/q)))


def main():
    rng=np.random.default_rng(965)
    out={'normalization':[],'contraction':[],'likelihood':[],'kl':[]}
    ph=np.linspace(0,2*np.pi,200000,endpoint=False)
    for r in [0,.2,.5,.8,.95]:
        x=r*np.exp(.73j); integ=float(np.mean(poisson_density(ph,x)))
        out['normalization'].append({'r':r,'integral':integ,'error':abs(integ-1)})
    x=.37+.19j; probes=np.array([-.7+.1j,.2-.5j,.6+.2j,-.1+.7j])
    for q in [.2,.5,.8]:
        ks=np.arange(1,13); es=[]
        for k in ks:
            es.append(abs(estimate_attractor([(q,(1-q)*x)]*int(k),probes)-x))
        slope=float(np.polyfit(ks,np.log(np.maximum(es,1e-16)),1)[0])
        out['contraction'].append({'q':q,'predicted_log_slope':math.log(q),'observed_log_slope':slope,'relative_error':abs(slope-math.log(q))/abs(math.log(q)),'error_K1':es[0],'error_K12':es[-1]})
    # Under normalized Lebesgue, uniform density is 1 and NLL is 0.
    for r in [0,.2,.5,.8]:
        x=r*np.exp(.41j); ps=sample_poisson_phases(x,200000,rng)
        pk=poisson_nll(ps,x); out['likelihood'].append({'r':r,'pk_nll':pk,'uniform_nll':0.0,'pk_minus_uniform':pk})
        out['kl'].append({'r':r,'empirical_to_pk':hist_kl(ps,x),'empirical_to_uniform':hist_kl(ps,0j)})
    out['summary']={'max_normalization_error':max(v['error'] for v in out['normalization']),'contraction_relative_errors':[v['relative_error'] for v in out['contraction']],'pk_nll_becomes_more_negative_with_radius':all(out['likelihood'][i]['pk_minus_uniform']<=out['likelihood'][i-1]['pk_minus_uniform'] for i in range(1,4))}
    with open('toy_results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()