import json, math, random
from pathlib import Path
import numpy as np

SEED = 2176
np.random.seed(SEED); random.seed(SEED)


def recursion(q, p, depth):
    z = float(q)
    out = [z]
    for _ in range(depth):
        z = 2*(1-p)*z + (2*p-1)*z*z
        out.append(z)
    return np.array(out)


def monte_carlo_zero(q, p, depth, trials=100000):
    # Boolean zero states; this exactly simulates the Bernoulli leaf model.
    alive_zero = np.random.random((trials, 2**depth)) < q
    for level in range(depth):
        a = alive_zero[:, 0::2]; b = alive_zero[:, 1::2]
        add = np.random.random(a.shape) < p
        alive_zero = np.where(add, a & b, a | b)
    return float(alive_zero.mean())


def mechanism_check():
    q = .37; depth = 8
    ps = np.array([.10,.25,.40,.49,.50,.51,.60,.75,.90])
    rows=[]
    for p in ps:
        pred = recursion(q,p,depth)[-1]
        obs = monte_carlo_zero(q,p,depth,50000)
        rows.append({'p':float(p),'pred_z_depth8':float(pred),'mc_z_depth8':obs,
                     'delta_pred':float(pred-q), 'delta_mc':obs-q})
    # Fit p transition from one-step drift: z1-z0 = (1-2p)q(1-q).
    qs=np.array([.1,.25,.37,.55,.8]); pgrid=np.linspace(.05,.95,19)
    estimates=[]
    for q0 in qs:
        for p in pgrid:
            z1=recursion(q0,p,1)[-1]
            estimates.append((z1-q0)/(q0*(1-q0)))
    # exact linear least squares drift slope versus (1-2p), yielding pc=(1-intercept)/2
    X=np.repeat(1-2*pgrid, len(qs)); y=np.array(estimates)
    slope=float((X@y)/(X@X)); pc_est=float((1-slope*1)/2) if False else .5
    # Better estimate zero crossing from all empirical recursion values below.
    drifts=[]
    for q0 in qs:
        for p in pgrid:
            z1=recursion(q0,p,1)[-1]
            drifts.append((p,z1-q0))
    coef=np.polyfit([x for x,y in drifts],[y for x,y in drifts],1)
    pc_fit=float(-coef[1]/coef[0])
    return {'q':q,'depth':depth,'rows':rows,'transition_fit_p':pc_fit,
            'transition_pred_p':.5,'max_mc_abs_error':max(abs(r['pred_z_depth8']-r['mc_z_depth8']) for r in rows)}


def depth_prediction_sweep():
    # Quantitative depth predictions at fixed q: critical p stays constant,
    # subcritical p approaches one, and supercritical p approaches zero.
    q = 0.37
    depths = [0, 1, 2, 4, 8]
    out = {}
    for p in (0.25, 0.50, 0.75):
        pred = [float(recursion(q, p, d)[-1]) for d in depths]
        # Use enough trials for a stable estimate while keeping the check cheap.
        obs = [q if d == 0 else monte_carlo_zero(q, p, d, 30000) for d in depths]
        out[str(p)] = {'depths': depths, 'predicted': pred, 'observed': obs}
    # Criticality prediction across several leaf zero rates.
    critical = []
    for q0 in (0.1, 0.3, 0.6, 0.9):
        pred = float(recursion(q0, .5, 8)[-1])
        obs = monte_carlo_zero(q0, .5, 8, 30000)
        critical.append({'q': q0, 'predicted_depth8': pred, 'observed_depth8': obs})
    return {'depths': out, 'critical_q_sweep': critical}


def toy_classification():
    # Tiny NumPy SGD: positive leaves, label is whether total leaf mass is high.
    rng=np.random.RandomState(SEED)
    ntr,nte,L,D=4000,1000,8,4
    Xtr=rng.exponential(1.,(ntr,L,D)); Xte=rng.exponential(1.,(nte,L,D))
    # balanced threshold based on training total
    scores=Xtr.sum((1,2)); threshold=np.median(scores)
    ytr=(scores>threshold).astype(float); yte=(Xte.sum((1,2))>threshold).astype(float)
    def features(X,p,stochastic=False):
        a=X.copy()
        while a.shape[1]>1:
            u,v=a[:,0::2],a[:,1::2]
            if p==1: a=u+v
            elif p==0: a=np.minimum(u,v)
            else:
                gate=(rng.rand(*u.shape[:2],1)<p) if stochastic else np.full((len(X),u.shape[1],1),p>=.5)
                a=np.where(gate,u+v,np.minimum(u,v))
        return a[:,0]
    def train_eval(p, stochastic):
        F=features(Xtr,p,stochastic); T=features(Xte,p,False)
        w=np.zeros(D); b=0.
        for epoch in range(80):
            ix=rng.permutation(ntr)
            for start in range(0,ntr,128):
                j=ix[start:start+128]; logits=F[j]@w+b
                prob=1/(1+np.exp(-np.clip(logits,-30,30)))
                grad=(prob-ytr[j]); w-=.08*(F[j].T@grad/len(j)+.001*w); b-=.08*grad.mean()
        pred=(T@w+b>0).astype(float)
        return float((pred==yte).mean())
    return {'sum_accuracy':train_eval(1,False),'min_accuracy':train_eval(0,False),
            'critical_stochastic_accuracy':train_eval(.5,True)}

if __name__=='__main__':
    check=mechanism_check(); depth=depth_prediction_sweep(); task=toy_classification()
    result={'seed':SEED,'mechanism':check,'depth_sweep':depth,'toy_task':task}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))