import math, random, time
import numpy as np
import torch

SEED = 2259
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def compositions(d, m):
    out=[]
    def rec(i, left, a):
        if i == d-1:
            out.append(tuple(a+[left])); return
        for k in range(left+1): rec(i+1, left-k, a+[k])
    rec(0,m,[])
    return out

def mult(alpha, m):
    z=math.factorial(m)
    for a in alpha: z//=math.factorial(a)
    return z

def features(x, alphas, weighted=True):
    # x: [batch,d]
    vals=[]
    for a in alphas:
        v=torch.ones(x.shape[0], device=x.device, dtype=x.dtype)
        for j,k in enumerate(a):
            if k: v=v*x[:,j]**k
        if weighted: v=v*math.sqrt(mult(a, sum(a)))
        vals.append(v)
    return torch.stack(vals,1)

def math_checks():
    print('MATH_CHECKS')
    # Prediction 1: exact parameter count and ratio.
    for d in (2,4,8,16):
        for m in (2,4,6):
            ordered=d**m; orbit=math.comb(d+m-1,m)
            print(f'COUNT d={d} m={m} observed_ratio={ordered/orbit:.6f} predicted={ordered/orbit:.6f} ordered={ordered} orbit={orbit}')
    # Prediction 2: diagonal identity, independently expanded ordered tensor.
    for d,m in [(3,2),(4,4),(5,6)]:
        aa=compositions(d,m); x=np.random.randn(d); vals=np.random.randn(len(aa))
        p=sum(mult(a,m)*v*np.prod(x**np.array(a)) for a,v in zip(aa,vals))
        ordered=0.0
        # enumerate ordered tuples and use histogram lookup
        from itertools import product
        lookup={a:v for a,v in zip(aa,vals)}
        for tup in product(range(d), repeat=m):
            hist=tuple(tup.count(j) for j in range(d))
            ordered += lookup[hist]*np.prod(x[list(tup)])
        print(f'IDENTITY d={d} m={m} abs_error={abs(p-ordered):.3e}')
    # Prediction 3: the explicit weighted feature has second moment exactly N times
    # the unweighted feature under iid standard normal inputs.
    for d,m in [(4,2),(4,4),(4,6),(8,4)]:
        aa=compositions(d,m); x=torch.randn(180000,d)
        f0=features(x,aa,False); f1=features(x,aa,True)
        # Select low, middle, max multiplicity orbit and average empirical ratios.
        ns=np.array([mult(a,m) for a in aa]); picks=[int(np.argmin(ns)), int(np.argsort(ns)[len(ns)//2]), int(np.argmax(ns))]
        for i in picks:
            ratio=(f1[:,i].square().mean()/f0[:,i].square().mean()).item()
            print(f'VAR d={d} m={m} N={ns[i]} observed={ratio:.4f} predicted={ns[i]} alpha={aa[i]}')
    # Prediction 4: repeating each orbit coefficient N times gives exactly the
    # ordered BH q-sum, for every q (including q=2m/(m+1)).
    for d,m in [(3,2),(4,4),(5,6)]:
        aa=compositions(d,m); q=2*m/(m+1); coeff=np.random.randn(len(aa))
        orbit=sum(mult(a,m)*abs(c)**q for a,c in zip(aa,coeff))
        repeated=sum(abs(coeff[aa.index(tuple(t.count(j) for j in range(d)))])**q
                     for t in __import__('itertools').product(range(d), repeat=m))
        print(f'BH d={d} m={m} q={q:.6f} abs_error={abs(orbit-repeated):.3e}')

def regression():
    # Same polynomial target, same orbit basis: compare exact multiplicity basis to unweighted compression.
    torch.manual_seed(SEED)
    d,m=8,4; aa=compositions(d,m); p=len(aa)
    ntr,nva=6000,1500
    xtr=torch.randn(ntr,d); xva=torch.randn(nva,d)
    target_theta=torch.randn(p)/math.sqrt(p)
    with torch.no_grad():
        ytr=features(xtr,aa,True)@target_theta; yva=features(xva,aa,True)@target_theta
    results={}
    for name,weighted in [('balanced',True),('unweighted',False)]:
        torch.manual_seed(SEED+int(weighted)); w=torch.nn.Parameter(torch.zeros(p))
        opt=torch.optim.Adam([w],lr=0.03)
        t=time.perf_counter(); losses=[]
        for step in range(500):
            ix=torch.randint(0,ntr,(128,)); pred=features(xtr[ix],aa,weighted)@w
            loss=((pred-ytr[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
            if step in (0,49,199,499): losses.append(float(loss))
        with torch.no_grad(): val=float(((features(xva,aa,weighted)@w-yva)**2).mean())
        results[name]=(val,time.perf_counter()-t,losses)
    print('REGRESSION d=8 m=4 params_each=%d' % p)
    for k,v in results.items(): print(f'{k} val_mse={v[0]:.6g} seconds={v[1]:.3f} checkpoints={v[2]}')
    print('DENSE_ORDERED_PARAMETER_COUNT',d**m,'ORBIT_PARAMETER_COUNT',p)

if __name__=='__main__':
    math_checks(); regression()