import itertools, json, math, random
import numpy as np
import torch
from torch import nn

SEED = 3020

def set_seed(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)

# Six edges of K4. A 3-edge subset is a spanning tree iff it is acyclic.
EDGES = [(0,1),(0,2),(0,3),(1,2),(1,3),(2,3)]

def is_forest(subset):
    parent = list(range(4))
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    for ei in subset:
        a,b = EDGES[ei]; ra,rb = find(a),find(b)
        if ra == rb: return False
        parent[ra] = rb
    return True

def enumerate_trees(weights=None):
    weights = np.ones(6) if weights is None else np.asarray(weights, float)
    masks=[]; ws=[]
    for comb in itertools.combinations(range(6), 3):
        if is_forest(comb):
            m=np.zeros(6,dtype=np.int64); m[list(comb)]=1
            masks.append(m); ws.append(float(np.prod(weights[list(comb)])))
    ws=np.asarray(ws); return np.asarray(masks), ws/ws.sum()

def math_check():
    masks, prob = enumerate_trees()
    inc = prob @ masks
    joint = (prob[:,None,None] * masks[:,:,None] * masks[:,None,:]).sum(0)
    cov = joint - inc[:,None]*inc[None,:]
    # For increasing coordinate events {i selected}, negative association is pairwise visible.
    pairwise_ok = bool(np.max(np.triu(cov,1)) <= 1e-12)
    # Direct multiaffinity: every enumerated monomial has exponents 0/1.
    multiaffine = bool(np.all((masks*masks)==masks))
    # Check the stated log-submodular inequality for Z(A)=sum weights of trees contained in A.
    z={}
    for bits in range(1<<6):
        A={i for i in range(6) if bits>>i & 1}
        z[bits]=float(sum(w for m,w in zip(masks, np.ones(len(masks))) if set(np.flatnonzero(m)).issubset(A)))
    min_slack=1e9; violations=0
    for a in range(1<<6):
        for b in range(1<<6):
            union=a|b; inter=a&b
            slack=z[a]*z[b]-z[union]*z[inter]
            min_slack=min(min_slack,slack)
            violations += slack < -1e-10
    return {
        'num_trees': int(len(masks)), 'inclusion_probs': inc.tolist(),
        'max_pair_covariance': float(np.max(np.triu(cov,1))),
        'min_pair_covariance': float(np.min(np.triu(cov,1)[np.triu_indices(6,1)])),
        'negative_pairwise_dependence': pairwise_ok,
        'multiaffine': multiaffine,
        'log_submodular_min_slack': float(min_slack),
        'log_submodular_violations': int(violations)
    }

def make_data(n, seed):
    rng=np.random.default_rng(seed)
    # Three latent factors generate correlated candidate routes; target uses all factors.
    latent=rng.normal(size=(n,3)); x=np.empty((n,6))
    x[:,0]=latent[:,0]+.18*rng.normal(size=n); x[:,1]=latent[:,0]+.18*rng.normal(size=n)
    x[:,2]=latent[:,1]+.18*rng.normal(size=n); x[:,3]=latent[:,1]+.18*rng.normal(size=n)
    x[:,4]=latent[:,2]+.18*rng.normal(size=n); x[:,5]=latent[:,2]+.18*rng.normal(size=n)
    y=latent @ np.array([1.0,-0.8,0.6]) + .15*rng.normal(size=n)
    return torch.tensor(x,dtype=torch.float32), torch.tensor(y[:,None],dtype=torch.float32)

class SmallNet(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(6,12),nn.Tanh(),nn.Linear(12,1))
    def forward(self,x): return self.net(x)

def train_once(kind, seed, steps=500):
    set_seed(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        tr_x,tr_y=make_data(1600,seed+100); va_x,va_y=make_data(800,seed+200)
        tr_x,tr_y,va_x,va_y=[v.to(device) for v in (tr_x,tr_y,va_x,va_y)]
        model=SmallNet().to(device); opt=torch.optim.Adam(model.parameters(),lr=0.012)
        forest_masks, forest_p=enumerate_trees()
        rng=np.random.default_rng(seed+500)
        for _ in range(steps):
            idx=torch.randint(0,len(tr_x),(64,),device=device); xb=tr_x[idx]; yb=tr_y[idx]
            if kind=='forest':
                mi=rng.choice(len(forest_masks),size=len(idx),p=forest_p)
                mask=torch.tensor(forest_masks[mi],dtype=torch.float32,device=device)
            else:
                mask=(torch.rand((len(idx),6),device=device)<0.5).float()
            xb=xb*mask/0.5
            loss=((model(xb)-yb)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            val=((model(va_x)-va_y)**2).mean().item()
        return val
    except Exception:
        # GPU failures (including shared-memory allocation issues) are retried on CPU.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        torch.set_default_device('cpu')
        return train_once_cpu(kind,seed,steps)

def train_once_cpu(kind, seed, steps=500):
    set_seed(seed); tr_x,tr_y=make_data(1600,seed+100); va_x,va_y=make_data(800,seed+200)
    model=SmallNet(); opt=torch.optim.Adam(model.parameters(),lr=0.012); masks,p=enumerate_trees(); rng=np.random.default_rng(seed+500)
    for _ in range(steps):
        idx=torch.randint(0,len(tr_x),(64,)); xb=tr_x[idx]; yb=tr_y[idx]
        m=torch.tensor(masks[rng.choice(len(masks),len(idx),p=p)] if kind=='forest' else (rng.random((len(idx),6))<.5),dtype=torch.float32)
        loss=((model(xb*m/.5)-yb)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    return ((model(va_x)-va_y)**2).mean().item()

def experiment():
    vals={k:[train_once(k,s) for s in range(3)] for k in ('bernoulli','forest')}
    return {k:{'runs':v,'mean':float(np.mean(v)),'std':float(np.std(v,ddof=1))} for k,v in vals.items()}

if __name__=='__main__':
    out={'math':math_check(),'experiment':experiment()}
    print(json.dumps(out,indent=2))