import itertools, math, json, random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED=17
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE='cuda' if torch.cuda.is_available() else 'cpu'

def supports(N,S):
    out=[]
    for s in range(1,S+1):
        out += [(tuple(c), tuple(k)) for c in itertools.combinations(range(N),s)
                for k in itertools.product((1,), repeat=s)]
    return out

def interaction_features(x,S):
    # x is {-1,+1}; q=2 Fourier characters, one per nonzero support
    cols=[x]
    N=x.shape[1]
    for s in range(2,S+1):
        for c in itertools.combinations(range(N),s): cols.append(torch.prod(x[:,c],dim=1,keepdim=True))
    return torch.cat(cols,1)

def bh_penalty(w,N,S,eps=1e-8):
    # w includes one coefficient per feature and one output bias
    pos=0; val=0.
    for s in range(1,S+1):
        n=math.comb(N,s); p=2*s/(s+1)
        a=w[pos:pos+n]; pos += n
        val=val+(a.square()+eps).pow(p/2).sum().pow(1/p)
    return val

class FourierModel(nn.Module):
    def __init__(self,N,S):
        super().__init__(); self.N=N; self.S=S
        self.w=nn.Parameter(torch.zeros(sum(math.comb(N,s) for s in range(1,S+1))))
        self.b=nn.Parameter(torch.zeros(1))
    def forward(self,x): return interaction_features(x,self.S)@self.w+self.b

class MLP(nn.Module):
    def __init__(self,N,h):
        super().__init__(); self.net=nn.Sequential(nn.Linear(N,h),nn.Tanh(),nn.Linear(h,1))
    def forward(self,x): return self.net(x).squeeze(1)

def train(model,xtr,ytr,xte,yte, penalty=False, epochs=350):
    opt=torch.optim.Adam(model.parameters(),lr=.03,weight_decay=0)
    for _ in range(epochs):
        opt.zero_grad(); z=model(xtr); loss=F.binary_cross_entropy_with_logits(z,ytr.float())
        if penalty: loss=loss+0.003*bh_penalty(model.w,model.N,model.S)
        loss.backward(); opt.step()
    with torch.no_grad():
        pred=(model(xte)>0).long(); acc=(pred==yte).float().mean().item()
        tr=(model(xtr)>0).long(); tracc=(tr==ytr).float().mean().item()
    return tracc,acc

def math_check():
    # Exact Fourier coefficients and infinity norm on C_2^N for homogeneous order s.
    rows=[]
    for N in [3,4,5,6,7,8]:
        s=2; terms=list(itertools.combinations(range(N),s)); X=torch.tensor(list(itertools.product([-1.,1.],repeat=N)))
        # same coefficients across dimensions: normalized random draws, repeated trials
        ratios=[]
        for t in range(30):
            a=torch.randn(len(terms)); f=torch.zeros(len(X))
            for j,c in enumerate(terms): f += a[j]*torch.prod(X[:,c],1)
            p=4/3; q=(a.abs().pow(p).sum())**(1/p); ratios.append((q/f.abs().max()).item())
        rows.append({'N':N,'max_ratio':max(ratios),'median_ratio':float(np.median(ratios))})
    return rows

def benchmark_one(N, S, order, ntrain, seed):
    X=np.array(list(itertools.product([-1.,1.],repeat=N)),dtype=np.float32)
    y=((np.prod(X[:,:order],axis=1))<0).astype(np.int64)
    rng=np.random.RandomState(seed); perm=rng.permutation(len(X)); tr,te=perm[:ntrain],perm[ntrain:]
    xtr=torch.tensor(X[tr],device=DEVICE); xte=torch.tensor(X[te],device=DEVICE)
    ytr=torch.tensor(y[tr],device=DEVICE); yte=torch.tensor(y[te],device=DEVICE)
    fm=FourierModel(N,S).to(DEVICE); fu=FourierModel(N,S).to(DEVICE)
    # width chosen to be close to the Fourier parameter count.
    target=sum(math.comb(N,s) for s in range(1,S+1))+1
    h=max(1,round((target-1-N-1)/(N+1)))
    mm=MLP(N,h).to(DEVICE)
    rp=train(fm,xtr,ytr,xte,yte,True,epochs=300)
    ru=train(fu,xtr,ytr,xte,yte,False,epochs=300)
    rm=train(mm,xtr,ytr,xte,yte,False,epochs=300)
    with torch.no_grad():
        w=fm.w.detach(); pos=0; norms={}
        for s in range(1,S+1):
            n=math.comb(N,s); p=2*s/(s+1)
            norms[str(s)]=float((w[pos:pos+n].abs().pow(p).sum().pow(1/p)).cpu()); pos+=n
    return {'order':order,'N':N,'ntrain':ntrain,'fourier_budget_params':sum(p.numel() for p in fm.parameters()),
            'mlp_params':sum(p.numel() for p in mm.parameters()),'mlp_width':h,
            'budgeted_fourier_test':rp[1],'unpenalized_fourier_test':ru[1],'mlp_test':rm[1],
            'budgeted_fourier_train':rp[0],'unpenalized_fourier_train':ru[0], 'mlp_train':rm[0],
            'budgeted_quasinorm_by_order':norms}

def benchmark():
    return {'task_2way_low_order': [benchmark_one(8,2,2,64,SEED+i) for i in range(3)],
            'task_3way_higher_order': [benchmark_one(8,3,3,64,SEED+10+i) for i in range(3)]}

if __name__=='__main__':
    # CUDA can fail on shared hardware; retry whole benchmark on CPU.
    try:
        result={'math_check':math_check(),'benchmark':benchmark()}
    except Exception as e:
        DEVICE='cpu'; result={'error_fallback':repr(e),'math_check':math_check(),'benchmark':benchmark()}
    print(json.dumps(result,indent=2))