import json, math
import numpy as np

SEED = 1675
A = np.array([[[1.18, 0.22], [0.04, 0.82]],
              [[0.86, 0.16], [-0.10, 1.26]]], dtype=float)


def transition_matrix(p):
    # symmetric two-state chain: persistence p; primitive for 0<p<1
    return np.array([[p, 1-p], [1-p, p]], dtype=float)


def simulate(p, T=12000, burn=1000, seed=0, target=None, lr=0.0, epochs=0):
    """Finite-time top exponent; optional scalar persistence controller."""
    rng = np.random.default_rng(seed)
    v = np.array([1.0, 0.37]); v /= np.linalg.norm(v)
    z = 0
    vals = []
    for _ in range(epochs):
        v = np.array([1.0, 0.37]); v /= np.linalg.norm(v); z = 0; s = 0.0
        for t in range(T):
            w = A[z] @ v
            n = np.linalg.norm(w)
            s += np.log(n)
            v = w / n
            z = int(rng.random() >= p) if z == 0 else int(rng.random() < p)
        lam = s / T
        vals.append(lam)
        if target is not None:
            # finite-difference d lambda/dp, then gradient descent on squared target error
            hp=0.01
            lp=lambda_estimate(min(.98,p+hp), T=max(3000,T//3), seed=700+_)
            lm=lambda_estimate(max(.02,p-hp), T=max(3000,T//3), seed=900+_)
            grad=(lp-lm)/(2*hp)
            p = float(np.clip(p - lr * 2.0 * (lam-target) * grad, 0.02, 0.98))
    # independent estimate, collecting state frequencies and transitions
    counts = np.zeros(2, dtype=int); trans = np.zeros((2,2), dtype=int)
    s = 0.0; v = np.array([1.0, 0.37]); v /= np.linalg.norm(v); z = 0
    for t in range(T + burn):
        old = z
        w = A[z] @ v; n = np.linalg.norm(w); v = w/n
        z = int(rng.random() >= p) if z == 0 else int(rng.random() < p)
        if t >= burn:
            s += np.log(n); counts[old] += 1; trans[old,z] += 1
    return dict(lam=float(s/T), p=float(p), counts=counts.tolist(), trans=trans.tolist(), controller_trace=vals)


def lambda_estimate(p, T=8000, seed=1):
    return simulate(p, T=T, burn=1000, seed=seed)['lam']


def common_random_lambda(p, uniforms, burn=1000):
    """Coupled estimate: identical random uniforms make FD noise small."""
    z=0; v=np.array([1.0,.37]); v/=np.linalg.norm(v); s=0.0
    for t,u in enumerate(uniforms):
        old=z
        w=A[z]@v; n=np.linalg.norm(w); v=w/n
        z = int(u >= p) if z == 0 else int(u < p)
        if t >= burn: s += np.log(n)
    return s/(len(uniforms)-burn)


def finite_difference_scaling(p=0.5):
    # At symmetry p=.5, lambda'(p)=0 is predicted by mode-label exchange.
    rng=np.random.default_rng(991)
    uniforms=rng.random(8000)
    rows=[]
    for h in [0.08,0.04,0.02,0.01,0.005]:
        d=(common_random_lambda(p+h,uniforms)-common_random_lambda(p-h,uniforms))/(2*h)
        rows.append({'h':h, 'derivative':float(d), 'abs_error_vs_predicted_zero':float(abs(d))})
    # Away from symmetry, check derivative consistency under h halving.
    p2=.35; u2=np.random.default_rng(992).random(8000)
    drows=[]
    for h in [.04,.02,.01,.005]:
        d=(common_random_lambda(p2+h,u2)-common_random_lambda(p2-h,u2))/(2*h)
        drows.append({'h':h,'derivative':float(d)})
    return {'symmetry_point':{'p':p,'predicted_derivative':0.0,'rows':rows},
            'interior_point':{'p':p2,'rows':drows}}

def stationary_check(p):
    x=simulate(p,T=8000,burn=2000,seed=31)
    freq=np.array(x['counts'],float)/sum(x['counts'])
    pi=np.array([.5,.5])
    # q_j mass equation reduces to pi=pi P; compare empirical frequencies.
    residual=float(np.max(np.abs(freq - pi @ transition_matrix(p))))
    return {'empirical_pi':freq.tolist(), 'theoretical_pi':[.5,.5], 'mass_equation_max_residual':residual}


def boundary_sweep():
    rows=[]
    for p in [0.05,0.10,0.25,0.50,0.75,0.90,0.95]:
        vals=[lambda_estimate(p,10000,s) for s in [2,3,4,5]]
        rows.append({'p':p,'mean_lambda':float(np.mean(vals)), 'std_over_seeds':float(np.std(vals,ddof=1))})
    return rows


def controller_demo():
    # Target is the exponent at p=.5; controller starts at p=.2.
    target=lambda_estimate(.5,20000,77)
    x=simulate(.2,T=5000,burn=500,seed=100,target=target,lr=.08,epochs=12)
    return {'target_lambda':target,'initial_p':.2,'final_p':x['p'],'trace':x['controller_trace']}


def main():
    result={
      'matrices':A.tolist(),
      'stationary_check_p05':stationary_check(.5),
      'stationary_check_p095':stationary_check(.95),
      'smoothness_fd':finite_difference_scaling(.5),
      'boundary_sweep':boundary_sweep(),
      'controller':controller_demo(),
      'predictions':[
        'For primitive p in (0,1), empirical stationary masses should approach pi=(.5,.5).',
        'In the smooth interior, central finite-difference derivative error should decrease approximately quadratically with h after accounting for Monte Carlo noise.',
        'Near p=0 or p=1, mixing slows and finite-time Lyapunov estimates should have larger seed variance than at p=.5.'
      ]
    }
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()