import json, math, random
import numpy as np

# Toy conjugate expanding map F=h^{-1} o A o h, h(x)=x+eps*sin(x), A(z)=a*z.
# Exact identity: ell(x)=log(a)+u(F(x))-u(x), u(x)=-log(h'(x)).
def h(x, eps): return x + eps*np.sin(x)
def hp(x, eps): return 1.0 + eps*np.cos(x)
def hinv(z, eps):
    x=np.asarray(z,dtype=float).copy()
    for _ in range(40): x -= (h(x,eps)-z)/hp(x,eps)
    return x
def F(x,a,eps): return hinv(a*h(x,eps),eps)
def ell(x,a,eps):
    y=F(x,a,eps)
    return math.log(a)-math.log(hp(y,eps))+math.log(hp(x,eps))
def u(x,eps): return -np.log(hp(x,eps))

def mechanism_sweep():
    rng=np.random.default_rng(7); a=1.03; c=math.log(a)
    xs=rng.uniform(-math.pi,math.pi,1000)
    # P1: exact potential gives zero finite-horizon residual at every k.
    telescoping=[]
    for k in [1,2,4,8,16,32]:
        vals=[]
        for x0 in xs[:300]:
            x=x0; s=0.0
            for _ in range(k): s+=ell(x,a,.6); x=F(x,a,.6)
            vals.append(s-(u(x,.6)-u(x0,.6)+k*c))
        z=np.asarray(vals)
        telescoping.append({'k':k,'max_abs_R':float(np.max(np.abs(z))),
                            'std_R_over_k':float(np.std(z/k))})
    # P2: statewise variation grows quadratically for small eps.
    epsvals=np.array([.05,.1,.2,.4])
    variation=[]
    for e in [0,.05,.1,.2,.4,.6,.8]:
        z=np.array([ell(x,a,e)-c for x in xs])
        variation.append({'eps':e,'pointwise_var':float(np.var(z)),
                          'pointwise_max_abs':float(np.max(np.abs(z)))})
    log_slope=float(np.polyfit(np.log(epsvals),np.log([
        np.var([ell(x,a,e)-c for x in xs]) for e in epsvals]),1)[0])
    # P3: if u is scaled by alpha, R_k/k=(1-alpha)*(u(x0)-u(xk))/k.
    scaling=[]; e=.6
    for alpha in [0,.25,.5,.75,1.0]:
        row={'alpha':alpha}
        for k in [1,4,16,64]:
            vals=[]
            for x0 in xs[:200]:
                x=x0; s=0.
                for _ in range(k): s+=ell(x,a,e); x=F(x,a,e)
                vals.append(s-(alpha*(u(x,e)-u(x0,e))+k*c))
            row[str(k)]=float(np.std(np.asarray(vals)/k))
        scaling.append(row)
    rows=[r for r in scaling if r['alpha']<1]
    ratios=[]
    for r in rows:
        ratios.append({'alpha':r['alpha'],'ratio_to_alpha0_k1':r['1']/scaling[0]['1']})
    return {'telescoping':telescoping,'variation':variation,
            'small_eps_loglog_slope_variance_vs_eps':log_slope,
            'scaled_potential':scaling,
            'scaled_residual_ratio_prediction':ratios}

def learned_potential_demo():
    # Baseline: best constant log-Jacobian. Idea: fit a scalar potential and c.
    rng=np.random.default_rng(11); a=1.03; eps=.6
    x=rng.uniform(-math.pi,math.pi,1200); y=np.array([F(v,a,eps) for v in x])
    l=np.array([ell(v,a,eps) for v in x])
    cols=[]
    for n in range(1,9):
        cols += [np.sin(n*y)-np.sin(n*x), np.cos(n*y)-np.cos(n*x)]
    X=np.column_stack(cols+[np.ones(len(x))])
    coef=np.linalg.lstsq(X,l,rcond=None)[0]; r=l-X@coef
    point=np.mean((l-np.mean(l))**2)
    return {'cohomological_mse':float(np.mean(r*r)),
            'pointwise_const_mse':float(point),
            'fitted_c':float(coef[-1]), 'true_c':math.log(a),
            'mse_ratio_coh_over_point':float(np.mean(r*r)/point)}

def main():
    random.seed(0); np.random.seed(0)
    out={'predictions':mechanism_sweep(),'baseline_vs_idea':learned_potential_demo(),
         'prediction_statements':[
          'P1 exact conjugacy predicts R_k=0 for every horizon k.',
          'P2 pointwise variance is predicted to scale as eps^2 near eps=0.',
          'P3 imperfect-potential residual per step scales as (1-alpha) and has a 1/k boundary decay.'
         ]}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()