import json, math, time
from itertools import product
import numpy as np


def lattice(L):
    edges=[]
    for r in range(L):
        for c in range(L):
            u=r*L+c
            if c+1<L: edges.append((u,u+1))
            if r+1<L: edges.append((u,u+L))
    return edges


def energies(states, edges, J):
    # states: [S,N], values +/-1
    return -sum(J[k]*states[:,u]*states[:,v] for k,(u,v) in enumerate(edges))


def exact_distribution(L, beta, J):
    N=L*L
    states=np.array(list(product((-1,1), repeat=N)), dtype=np.int8)
    logw=-beta*energies(states, lattice(L), J)
    logw-=logw.max(); w=np.exp(logw); p=w/w.sum()
    return states,p


def conditional_oracle(states, p, i, prefix):
    mask=np.all(states[:,:i]==np.asarray(prefix)[None,:],axis=1)
    vals=states[mask,i]; ww=p[mask]
    zplus=ww[vals==1].sum(); zminus=ww[vals==-1].sum()
    return float(zplus/(zplus+zminus))


def all_prefix_data(states,p):
    X=[]; I=[]; Q=[]
    N=states.shape[1]
    # Prefix encoding: known spins, zero for unknown; position one-hot is appended.
    for s,ps in zip(states,p):
        # weighted selection is not needed: include every state and position,
        # with exact q at its realized prefix; duplicate prefixes are harmless.
        for i in range(N):
            pref=s[:i]
            q=conditional_oracle(states,p,i,pref)
            x=np.zeros(N,dtype=np.float32); x[:i]=pref
            X.append(x); I.append(i); Q.append(q)
    return np.asarray(X),np.asarray(I),np.asarray(Q)


def kw_partition(L, K):
    """Kac-Ward determinant for a square lattice embedded at integer coordinates."""
    edges=lattice(L); und=[]
    for u,v in edges: und += [(u,v),(v,u)]
    m=len(und); T=np.zeros((m,m),dtype=np.complex128)
    xy=lambda x:(x//L,x%L)
    for a,(u,v) in enumerate(und):
        ru,cu=xy(u); rv,cv=xy(v)
        vin=np.array([cv-cu,rv-ru],float) # incoming direction u -> v
        for b,(v2,w) in enumerate(und):
            if v2!=v or w==u: continue
            rw,cw=xy(w); vout=np.array([cw-cv,rw-rv],float)
            cross=vin[0]*vout[1]-vin[1]*vout[0]
            dot=vin.dot(vout)
            theta=0.5*math.atan2(cross,dot)
            # edge coupling belongs to the undirected edge
            kk=K[edges.index(tuple(sorted((u,v))))]
            T[a,b]=math.tanh(kk)*np.exp(1j*theta)
    sign,ld=np.linalg.slogdet(np.eye(m,dtype=complex)-T)
    # For planar ferromagnetic examples this branch is positive real.
    z=(2.0**(L*L))*np.prod(np.cosh(K))*np.sqrt(sign*np.exp(ld))
    return float(np.real_if_close(z).real)


def kw_check():
    rows=[]
    for L in (2,3):
        edges=lattice(L); K=np.array([0.17+0.04*((k%3)-1) for k in range(len(edges))])
        states=np.array(list(product((-1,1),repeat=L*L)),dtype=np.int8)
        z_enum=float(np.exp(np.logaddexp.reduce(sum(K[k]*states[:,u]*states[:,v] for k,(u,v) in enumerate(edges)))))
        z_kw=kw_partition(L,K)
        rows.append((L,z_enum,z_kw,abs(z_enum-z_kw)/z_enum))
    return rows


def train_model(X,I,Q, seed, soft, beta, steps=300, batch=128):
    import torch
    torch.manual_seed(seed); np.random.seed(seed)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        x=torch.tensor(X); ii=torch.nn.functional.one_hot(torch.tensor(I),X.shape[1]).float();
        inp=torch.cat([x,ii],1).to(dev); y=torch.tensor(Q,dtype=torch.float32).to(dev)
        rng=np.random.default_rng(seed+19); yy=torch.tensor((rng.random(len(Q))<Q).astype(np.float32)).to(dev)
        model=torch.nn.Sequential(torch.nn.Linear(inp.shape[1],48),torch.nn.Tanh(),torch.nn.Linear(48,1)).to(dev)
        opt=torch.optim.Adam(model.parameters(),lr=.025)
        rng=np.random.default_rng(seed+7)
        for _ in range(steps):
            ix=torch.tensor(rng.integers(0,len(Q),size=batch),device=dev)
            logit=model(inp[ix]).squeeze(1); target= y[ix] if soft else yy[ix]
            loss=torch.nn.functional.binary_cross_entropy_with_logits(logit,target)
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=torch.sigmoid(model(inp).squeeze(1)).cpu().numpy()
        eps=1e-7
        kl=Q*np.log((Q+eps)/(pred+eps))+(1-Q)*np.log((1-Q+eps)/(1-pred+eps))
        ce=-(Q*np.log(pred+eps)+(1-Q)*np.log(1-pred+eps))
        return float(np.mean(kl)),float(np.mean(ce)),dev
    except Exception as e:
        # retry CPU, as required for shared-GPU failures
        torch.manual_seed(seed); dev='cpu'
        inp=torch.tensor(np.concatenate([X,np.eye(X.shape[1])[I]],1),dtype=torch.float32)
        y=torch.tensor(Q,dtype=torch.float32); yy=torch.tensor((np.random.default_rng(seed+19).random(len(Q))<Q).astype(np.float32))
        model=torch.nn.Sequential(torch.nn.Linear(inp.shape[1],48),torch.nn.Tanh(),torch.nn.Linear(48,1)); opt=torch.optim.Adam(model.parameters(),lr=.025)
        rng=np.random.default_rng(seed+7)
        for _ in range(steps):
            ix=torch.tensor(rng.integers(0,len(Q),size=batch)); logit=model(inp[ix]).squeeze(1); target=y[ix] if soft else yy[ix]
            loss=torch.nn.functional.binary_cross_entropy_with_logits(logit,target); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): pred=torch.sigmoid(model(inp).squeeze(1)).numpy()
        eps=1e-7; kl=Q*np.log((Q+eps)/(pred+eps))+(1-Q)*np.log((1-Q+eps)/(1-pred+eps)); ce=-(Q*np.log(pred+eps)+(1-Q)*np.log(1-pred+eps))
        return float(np.mean(kl)),float(np.mean(ce)),dev


def main():
    kw=kw_check()
    # Quantitative prediction 1: KW equals enumeration to numerical precision.
    identity=[]
    for beta in (.2,.7,1.3):
        L=3; edges=lattice(L); J=np.array([1 if k%2 else -1 for k in range(len(edges))],float)
        s,p=exact_distribution(L,beta,J); X,I,Q=all_prefix_data(s,p)
        # Conditional CE - entropy equals KL, checked with a deliberately perturbed model.
        pred=np.clip(.15+.7*Q,.001,.999)
        ce=np.mean(-(Q*np.log(pred)+(1-Q)*np.log(1-pred)))
        ent=np.mean(-(Q*np.log(Q)+(1-Q)*np.log(1-Q)))
        kl=np.mean(Q*np.log(Q/pred)+(1-Q)*np.log((1-Q)/(1-pred)))
        # Prediction 2: sampled-label variance q(1-q), increasing near beta=0.
        var=float(np.mean(Q*(1-Q)))
        identity.append({'beta':beta,'identity_abs_err':abs((ce-ent)-kl),'mean_label_variance':var,'mean_q':float(Q.mean())})
    # Fixed tiny setup, same data and training budget.
    beta=.7; L=3; edges=lattice(L); J=np.array([1 if k%2 else -1 for k in range(len(edges))],float)
    s,p=exact_distribution(L,beta,J); X,I,Q=all_prefix_data(s,p)
    results=[]
    for b in (.2,.7,1.3):
        sb, pb=exact_distribution(L,b,np.ones(len(edges)))
        xb,ib,qb=all_prefix_data(sb,pb)
        out=[]
        for soft in (False,True):
            vals=[train_model(xb,ib,qb,seed=11+r,soft=soft,beta=b)[:2] for r in range(3)]
            out.append({'method':'soft_oracle' if soft else 'sampled_labels','kl_mean':float(np.mean([v[0] for v in vals])),'ce_mean':float(np.mean([v[1] for v in vals]))})
        results.append({'beta':b,'mean_q_variance':float(np.mean(qb*(1-qb))),'methods':out})
    # Prediction 3: exact enumeration grows 2^N while autoregressive pass is N outputs.
    timing=[]
    for L in (2,3,4):
        t=time.perf_counter(); exact_distribution(L,.5,np.ones(len(lattice(L)))); sec=time.perf_counter()-t
        timing.append({'L':L,'N':L*L,'enumeration_seconds':sec,'autoregressive_outputs':L*L})
    report={'kw_check':kw,'mechanism_checks':identity,'training_sweep':results,'scaling':timing}
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()