import json, math
from dataclasses import dataclass
import numpy as np

@dataclass
class Node:
    left: object = None
    right: object = None
    leaf: int = None
    def is_leaf(self): return self.leaf is not None

def leaf(i): return Node(leaf=i)

def depths(t, d=0, out=None):
    if out is None: out={}
    if t.is_leaf(): out[t.leaf]=d
    else:
        depths(t.left,d+1,out); depths(t.right,d+1,out)
    return out

def internal_sizes(t):
    if t.is_leaf(): return 1, []
    na,sa=internal_sizes(t.left); nb,sb=internal_sizes(t.right)
    return na+nb, sa+sb+[na+nb]

def lambdas(t,k):
    ds=depths(t)
    _, internal=internal_sizes(t)
    return float(sum(ds.values())), float(sum(x*x for x in internal))

def Kmat(t,k):
    K=np.zeros((k,k))
    def walk(n, leaves):
        if n.is_leaf(): return [n.leaf]
        a=walk(n.left,leaves); b=walk(n.right,leaves); s=a+b
        for i in s:
            for j in s: K[i,j]+=1
        return s
    walk(t,[]); return K

def balanced(ids):
    ids=list(ids)
    if len(ids)==1: return leaf(ids[0])
    m=len(ids)//2
    return Node(balanced(ids[:m]),balanced(ids[m:]))

def chain(ids):
    ids=list(ids); t=leaf(ids[0])
    for i in ids[1:]: t=Node(t,leaf(i))
    return t

def huffman(weights):
    # Min weighted external path length; children are immaterial to score.
    heap=[(float(w), i, leaf(i)) for i,w in enumerate(weights)]
    serial=len(heap)
    import heapq
    heapq.heapify(heap)
    while len(heap)>1:
        wa,_,a=heapq.heappop(heap); wb,_,b=heapq.heappop(heap)
        heapq.heappush(heap,(wa+wb,serial,Node(a,b))); serial+=1
    return heap[0][2]

def proxy(t, variances, mean2=0.0, lam=0.0):
    d,l2=lambdas(t,len(variances))
    return float(np.dot(variances,[depths(t)[i] for i in range(len(variances))]) + lam*mean2*l2)

def exact_identity_check(rng,k=8,n=120000):
    t=balanced(range(k)); K=Kmat(t,k); d,l2=lambdas(t,k)
    rows=[]
    for mu in [0.0,0.25,1.0,3.0]:
        for tau in [0.1,0.5,1.5]:
            p=rng.normal(mu,tau,size=(n,k))
            empirical=np.mean(np.einsum('bi,ij,bj->b',p,K,p))
            predicted=tau*tau*d+mu*mu*l2
            rows.append({'mu':mu,'tau':tau,'predicted':predicted,'observed':float(empirical),
                         'relative_error':float(abs(empirical-predicted)/predicted)})
    return {'tree':'balanced_8','Lambda1':d,'Lambda2':l2,'rows':rows,
            'max_relative_error':max(x['relative_error'] for x in rows)}

def roundoff_scaling(rng,k=8,n=80000):
    # Conditional independent additive roundoff at each internal node:
    # Var(error | x)=u^2 x^2. This is precisely the paper's local model with nu=1.
    p=rng.normal(0.2,0.7,size=(n,k)); t=balanced(range(k)); K=Kmat(t,k)
    q=np.einsum('bi,ij,bj->b',p,K,p)
    rows=[]
    for u in [2**-5,2**-7,2**-9,2**-11]:
        # independently inject Gaussian error at each node, aggregate final error
        errs=np.zeros(n)
        def rec(node):
            if node.is_leaf(): return p[:,node.leaf]
            x=rec(node.left)+rec(node.right)
            e=rng.normal(size=n)*u*np.abs(x)
            nonlocal errs
            errs += e
            return x+e
        rec(t)
        observed=float(np.mean(errs**2)); predicted=float(u*u*np.mean(q))
        rows.append({'u':u,'observed_mse':observed,'predicted_mse':predicted,
                     'ratio_obs_over_u2':observed/(u*u),
                     'relative_error':abs(observed-predicted)/predicted})
    # slope in log-log should be 2
    slope=float(np.polyfit(np.log([x['u'] for x in rows]),np.log([x['observed_mse'] for x in rows]),1)[0])
    return {'rows':rows,'loglog_slope':slope,'predicted_slope':2.0}

def heterogeneous_monte_carlo(rng, n=120000):
    k=8; variances=np.array([16.,1.,1.,1.,1.,1.,1.,1.])
    trees=[('balanced',balanced(range(k))),('huffman',huffman(variances))]
    rows=[]
    for name,t in trees:
        K=Kmat(t,k)
        p=rng.normal(size=(n,k))*np.sqrt(variances)[None,:]
        observed=float(np.mean(np.einsum('bi,ij,bj->b',p,K,p)))
        predicted=float(np.dot(variances,[depths(t)[i] for i in range(k)]))
        rows.append({'tree':name,'predicted_cost':predicted,'observed_cost':observed,
                     'relative_error':abs(observed-predicted)/predicted,
                     'depths':depths(t)})
    return rows

def topology_sweep():
    k=8; bal=balanced(range(k)); rows=[]
    for ratio in [1,2,4,8,16,32,64]:
        w=np.ones(k); w[0]=ratio
        # Put the high-variance chunk at the shallowest available Huffman leaf.
        ht=huffman(w); hd=depths(ht); order=sorted(hd,key=hd.get)
        # huffman construction may assign leaf 0 shallow naturally; relabel if needed
        # evaluate optimal assignment for the fixed shape by placing largest weights shallow.
        ds=sorted(depths(ht).values()); weighted_h=float(np.dot(sorted(w,reverse=True),ds))
        weighted_bal=float(np.dot(sorted(w,reverse=True),sorted(depths(bal).values())))
        rows.append({'ratio':ratio,'balanced_cost':weighted_bal,'huffman_cost':weighted_h,
                     'predicted_gain_fraction':1-weighted_h/weighted_bal,
                     'huffman_depth_high_variance':min(ds)})
    return rows

def main():
    rng=np.random.default_rng(1119)
    identity=exact_identity_check(rng)
    rounding=roundoff_scaling(rng)
    topology=topology_sweep()
    heterogeneous=heterogeneous_monte_carlo(rng)
    out={'identity_check':identity,'roundoff_scaling':rounding,'variance_topology_sweep':topology,
         'heterogeneous_monte_carlo':heterogeneous,
         'notes':['All costs omit common nu*u^2 factors. Huffman minimizes sum sigma_i^2 depth_i among binary trees.',
                  'The roundoff simulation uses the stated conditional Gaussian local-noise approximation.']}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()