import json
import numpy as np
from pathlib import Path

SEED = 2137
rng = np.random.default_rng(SEED)


def H(z, delta, P, beta):
    # z, delta: [B,S]; P: [B,S,S]
    continuation = np.min(z + delta, axis=0)
    return delta + beta * np.einsum("bxy,y->bx", P, continuation)


def solve(delta, P, beta, tol=1e-13, max_iter=10000):
    z = np.zeros_like(delta)
    for k in range(max_iter):
        zn = H(z, delta, P, beta)
        if np.max(np.abs(zn-z)) < tol:
            return zn, k + 1
        z = zn
    raise RuntimeError("fixed point did not converge")


def quantile_width(xs, ps, lo=.1, hi=.9):
    # Width between interpolated crossings; NaN if sweep is too narrow.
    def cross(q):
        ind = np.where((ps[:-1]-q) * (ps[1:]-q) <= 0)[0]
        if len(ind) == 0:
            return np.nan
        i = ind[0]
        if ps[i+1] == ps[i]: return xs[i]
        return xs[i] + (q-ps[i])*(xs[i+1]-xs[i])/(ps[i+1]-ps[i])
    return abs(cross(hi)-cross(lo))


def contraction_check():
    B, S, beta = 4, 7, .73
    raw = rng.random((B,S,S)); P = raw/raw.sum(axis=2, keepdims=True)
    delta = rng.random((B,S))
    z = rng.normal(size=(B,S)); w = rng.normal(size=(B,S))
    ratio = np.max(np.abs(H(z,delta,P,beta)-H(w,delta,P,beta))) / np.max(np.abs(z-w))
    # Use a generic instance to avoid accidental finite-state exactness.
    zstar, _ = solve(delta,P,beta)
    zk = np.zeros_like(delta)
    errors = []
    bounds = []
    initial = np.max(np.abs(zk-zstar))
    for k in range(9):
        errors.append(float(np.max(np.abs(zk-zstar))))
        bounds.append(float(beta**k * initial))
        zk = H(zk,delta,P,beta)
    positive = [errors[i+1]/errors[i] for i in range(8) if errors[i] > 1e-14]
    beta_sweep=[]
    for bb in [.2, .5, .8, .95]:
        zz,_=solve(delta,P,bb); cur=np.zeros_like(delta); es=[]
        for _ in range(12):
            es.append(np.max(np.abs(cur-zz))); cur=H(cur,delta,P,bb)
        rr=es[-1]/es[-2] if es[-2] > 1e-14 else 0.0
        beta_sweep.append({"beta":bb,"late_error_ratio":float(rr),"predicted_upper":bb})
    return {"beta": beta, "operator_ratio": float(ratio),
            "max_error_over_bound": float(max(e/b for e,b in zip(errors,bounds))),
            "iteration_ratios": positive, "predicted_ratio_upper": beta,
            "beta_sweep":beta_sweep}


def tied_mdp():
    # State 0 is the decision state. Each branch reaches a different
    # continuation state; continuation states self-loop.
    B, S, beta = 3, 4, .8
    P = np.zeros((B,S,S))
    destinations = [1,2,3]
    for b in range(B):
        P[b,0,destinations[b]] = 1.
        for x in range(1,S): P[b,x,x] = 1.
    delta = np.zeros((B,S))
    future_cost = np.array([0., .20, .35, .27])
    delta[:,1:] = future_cost[1:]
    # Set root deficits so the three fixed-point scores tie, then add a
    # positive common offset so signed perturbations remain valid deficits.
    z,_ = solve(delta,P,beta)
    base_scores = z + delta
    target = np.max(base_scores[:,0])
    delta[:,0] += (target-base_scores[:,0])/2 + 1.0
    return P, delta, beta


def scaling_check():
    P, delta0, beta = tied_mdp()
    B = 3
    # Verify exact tie before perturbation.
    z0,_ = solve(delta0,P,beta); s0=z0+delta0
    tie_spread=float(np.ptp(s0[:,0]))
    Ns=[20,50,100,200]
    widths=[]; curves={}
    # kappa=d=1: predicted critical temperature tau ~ N^-1.
    for N in Ns:
        tau=.35/N
        ts=np.linspace(-8*tau,8*tau,501)
        ps=[]
        for t in ts:
            d=delta0.copy(); d[0,0]+=t
            z,_=solve(d,P,beta)
            scores=z[:,0]+d[:,0]
            a=np.exp(-(scores-scores.min())/tau); p=a/a.sum()
            ps.append(p[0])
        ps=np.asarray(ps)
        widths.append(float(quantile_width(ts,ps)))
        curves[N]=(ts/tau,ps)
    # Collapse error against N=100 on common normalized coordinates.
    grid=np.linspace(-8,8,1001)
    ref=np.interp(grid,*curves[100])
    collapse=max(float(np.max(np.abs(np.interp(grid,*curves[N])-ref))) for N in Ns)
    scaled_widths=[w*N for w,N in zip(widths,Ns)]

    # Ordinary noisy argmin: fixed critic noise does not inherit N^-1 scaling.
    noise_rng=np.random.default_rng(SEED+1)
    sigma=.08; widths_arg=[]
    ts=np.linspace(-.8,.8,301)
    noises=noise_rng.normal(0,sigma,size=(30000,3))
    for N in Ns:
        # N is only a label for the sampled-pool regime here; fixed noise
        # demonstrates the baseline's lack of critical N scaling.
        probs=[]
        for t in ts:
            scores=np.zeros((len(noises),3)); scores[:,0]=2*t
            probs.append(np.mean(np.argmin(scores+noises,axis=1)==0))
        widths_arg.append(float(quantile_width(ts,np.asarray(probs))))
    # Under critic noise, compare soft resolver probabilities and hard argmin
    # against the known lowest continuation-cost branch (branch 0 at root).
    rr=np.random.default_rng(SEED+9)
    M=50000; noise=rr.normal(0,.06,size=(M,3)); perturb=.03
    noisy=np.zeros((M,3))+noise; noisy[:,0]+=perturb
    hard=np.mean(np.argmin(noisy,axis=1)==0)
    tau=.03
    logits=-noisy/tau; logits-=logits.max(axis=1,keepdims=True)
    soft=np.exp(logits); soft/=soft.sum(axis=1,keepdims=True)
    resolver_prob=float(np.mean(soft[:,0]))
    return {"tie_spread":tie_spread, "Ns":Ns, "tau_rule":"0.35/N (kappa=d=1)",
            "resolver_widths":widths, "resolver_width_times_N":scaled_widths,
            "predicted_width_times_N_constant":True,
            "normalized_curve_max_deviation":collapse,
            "baseline_noisy_argmin_widths":widths_arg,
            "baseline_width_times_N":[w*N for w,N in zip(widths_arg,Ns)],
            "predicted_baseline_width_N_scaling":"constant width, hence width*N grows linearly",
            "direct_noisy_selection": {"perturbation":perturb,"noise_sigma":.06,
                "baseline_hard_best_branch_frequency":float(hard),
                "resolver_soft_probability_best_branch":resolver_prob}}


def main():
    out={"seed":SEED, "contraction":contraction_check(), "scaling":scaling_check()}
    Path("results.json").write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__ == "__main__": main()