import json, math, os
import numpy as np

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


def corr_matrix(L, gamma):
    idx = np.arange(L)
    return (1.0 + np.abs(idx[:, None] - idx[None, :])) ** (-gamma)


def sums(L, gamma):
    t = np.arange(1, L + 1, dtype=float)
    c = t ** (-gamma)
    return float(np.sum(c)), float(np.sum(c * c))


def log_slope(xs, ys):
    return float(np.polyfit(np.log(xs), np.log(np.maximum(ys, 1e-30)), 1)[0])


def empirical_gaussian_check(L, gamma, n=12000, block=512):
    # For x ~ N(0,C), E[(sum x_l^2)^2] = tr(C)^2 + 2 tr(C^2).
    C = corr_matrix(L, gamma)
    chol = np.linalg.cholesky(C + 1e-10 * np.eye(L))
    vals = []
    for start in range(0, n, block):
        z = rng.normal(size=(min(block, n-start), L))
        x = z @ chol.T
        vals.append(np.sum(x*x, axis=1))
    q = np.concatenate(vals)
    empirical = float(np.mean(q*q))
    exact = float(np.trace(C)**2 + 2*np.trace(C@C))
    # Excess over independent diagonal baseline, divided by L, is governed by sum c_t^2.
    ind = float(L*L + 2*L)
    return empirical, exact, float((exact-ind)/(2*L))


def residual_jacobian_stats(L, gamma, width=32, trials=20, alpha=.8):
    # Small fixed residual MLP Jacobian: x_{l+1}=x_l + scale*tanh(W_l x_l).
    # W_l has prescribed cross-depth covariance coordinatewise.
    d = width
    scale = 0.35 / math.sqrt(L)
    C = corr_matrix(L, gamma)
    chol = np.linalg.cholesky(C + 1e-10*np.eye(L))
    singular = []
    grad4 = []
    for _ in range(trials):
        E = rng.normal(size=(L,d,d)) / math.sqrt(d)
        G = np.empty((L,d,d))
        for i in range(d):
            for j in range(d):
                G[:,i,j] = chol @ rng.normal(size=L) / math.sqrt(d)
        W = math.sqrt(alpha)*G + math.sqrt(1-alpha)*E
        J = np.eye(d)
        x = rng.normal(size=d)
        for l in range(L):
            a = W[l] @ x
            deriv = 1.0 - np.tanh(a)**2
            J = (np.eye(d) + scale*(deriv[:,None]*W[l])) @ J
            x = x + scale*np.tanh(a)
        sv = np.linalg.svd(J, compute_uv=False)
        singular.append([float(np.median(sv)), float(np.max(sv))])
        # row-wise Jacobian sensitivity proxy
        grad4.append(float(np.mean(np.sum(J*J,axis=1)**2) / (np.mean(np.sum(J*J,axis=1))**2)))
    return {'median_sv': float(np.mean(np.array(singular)[:,0])),
            'max_sv': float(np.mean(np.array(singular)[:,1])),
            'normalized_gradient_fourth': float(np.mean(grad4))}


def main():
    depths = np.array([32,64,128,256,512,1024,2048,4096,8192,16384,32768])
    gammas = [0.3, 0.5, 0.75, 1.0, 1.25]
    rows=[]
    for g in gammas:
        s1=np.array([sums(int(L),g)[0] for L in depths])
        s2=np.array([sums(int(L),g)[1] for L in depths])
        # Predicted asymptotic slopes: max(1-g,0), max(1-2g,0), with log boundary.
        expected_s1 = max(1-g,0.0)
        expected_s2 = max(1-2*g,0.0)
        empirical, exact, excess = empirical_gaussian_check(128,g)
        rows.append({'gamma':g, 's1_slope_observed':log_slope(depths,s1),
                     's1_slope_predicted':expected_s1,
                     's2_slope_observed':log_slope(depths,s2),
                     's2_slope_predicted':expected_s2,
                     's1_at_final':float(s1[-1]), 's2_at_final':float(s2[-1]),
                     's1_tail_slope_observed':log_slope(depths[-5:],s1[-5:]),
                     's2_tail_slope_observed':log_slope(depths[-5:],s2[-5:]),
                     's2_normalized_final':float(s2[-1]/(depths[-1]**max(1-2*g,0))) if g<.5 else None,
                     'gaussian_q2_empirical_L128':empirical,
                     'gaussian_q2_exact_L128':exact,
                     'gaussian_excess_per_2L':excess})
    # Check critical logarithms by fitting against log L; ratio should stabilize.
    critical=[]
    for g in [.5,1.0]:
        ss=np.array([sums(int(L),g)[1 if g==.5 else 0] for L in depths])
        critical.append({'gamma':g,'quantity':'S2' if g==.5 else 'S1',
                         'log_fit_slope':log_slope(depths,ss),
                         'endpoints': [float(ss[0]),float(ss[-1])],
                         'log_ratio_start':float(ss[0]/math.log(depths[0])),
                         'log_ratio_end':float(ss[-1]/math.log(depths[-1]))})
    # Network sanity comparison at moderate depth; initialization is the only changed factor.
    net=[]
    for L in [16,32,64,128]:
        iid = residual_jacobian_stats(L, 100.0, trials=8, alpha=0.0)
        corr = residual_jacobian_stats(L, .3, trials=8, alpha=1.0)
        net.append({'depth':L,'iid':iid,'powerlaw_gamma_0.3':corr})
    out={'seed':SEED,'depths':depths.tolist(),'formula_sweeps':rows,
         'critical_log_checks':critical,'network_jacobian_sanity':net,
         'notes':'S1=sum_{t=1}^L t^(-gamma), S2=sum c_t^2. Slopes use finite-depth log-log regression. Gaussian check uses exact Wick formula.'}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()