import os, json, math, random
import numpy as np
import torch
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

SEED = 2063

def seed(s=SEED):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)


def toy_check():
    # V_line = y^2/2 has a 1-D zero manifold in R^2: lambda=(2-1)/2=0.5.
    # V_point = (x^2+y^2)/2 is an isolated zero: lambda=2/2=1.
    n = 1001
    x = np.linspace(-4, 4, n); dx = x[1] - x[0]
    X, Y = np.meshgrid(x, x, indexing='ij')
    Vline = 0.5 * Y**2
    Vpoint = 0.5 * (X**2 + Y**2)
    betas = np.array([5., 8., 12., 20., 35., 60., 100., 180., 300.])
    Zline = np.array([(np.exp(-b*Vline).sum()*dx*dx) for b in betas])
    Zpoint = np.array([(np.exp(-b*Vpoint).sum()*dx*dx) for b in betas])
    # Omit beta=5 to reduce finite-box effects.
    sl = np.polyfit(np.log(betas[1:]), np.log(Zline[1:]), 1)[0]
    sp = np.polyfit(np.log(betas[1:]), np.log(Zpoint[1:]), 1)[0]
    ratio_slope = np.polyfit(np.log(betas[1:]), np.log((Zline/Zpoint)[1:]), 1)[0]
    widths = []
    for b in betas:
        w = np.exp(-b*Vline)
        widths.append(math.sqrt(float((w * Y**2).sum() / w.sum())))
    width_slope = np.polyfit(np.log(betas[1:]), np.log(widths[1:]), 1)[0]
    return {
        'lambda_scaling': {
            'line_predicted': 0.5, 'line_observed': float(-sl),
            'point_predicted': 1.0, 'point_observed': float(-sp),
            'ratio_exponent_predicted': 0.5, 'ratio_exponent_observed': float(ratio_slope)},
        'width_scaling': {'predicted_exponent': -0.5, 'observed_exponent': float(width_slope)},
        'passed': bool(abs(-sl-.5)<.06 and abs(-sp-1.)<.08 and abs(ratio_slope-.5)<.08 and abs(width_slope+.5)<.06)
    }

class MLP(torch.nn.Module):
    def __init__(self):
        super().__init__(); self.net=torch.nn.Sequential(torch.nn.Linear(64,96),torch.nn.ReLU(),torch.nn.Linear(96,10))
    def forward(self,x): return self.net(x)

def loss_acc(model, X, y):
    with torch.no_grad():
        z=model(X); return float(torch.nn.functional.cross_entropy(z,y)), float((z.argmax(1)==y).float().mean())

def perturb_loss(model, X, y, sigma=.02, reps=12):
    vals=[]; params=[p for p in model.parameters() if p.requires_grad]
    with torch.no_grad():
        original=[p.detach().clone() for p in params]
        for _ in range(reps):
            for p in params: p.add_(sigma*torch.randn_like(p))
            vals.append(loss_acc(model,X,y)[0])
            for p,o in zip(params,original): p.copy_(o)
    return float(np.mean(vals))

def train_compare(device):
    data=load_digits(); X=data.data.astype('float32')/16.; y=data.target.astype('int64')
    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
    Xt=torch.tensor(Xtr,device=device); yt=torch.tensor(ytr,device=device)
    Xe=torch.tensor(Xte,device=device); ye=torch.tensor(yte,device=device)
    results={}
    for method in ['adamw','zero_langevin']:
        seed(); m=MLP().to(device)
        opt=torch.optim.AdamW(m.parameters(),lr=2e-3,weight_decay=1e-4)
        # Same ordinary training budget.
        for step in range(350):
            opt.zero_grad(); L=torch.nn.functional.cross_entropy(m(Xt),yt); L.backward(); opt.step()
        pre_train,pre_test=loss_acc(m,Xt,yt),loss_acc(m,Xe,ye)
        if method=='zero_langevin':
            eps=0.02; eta=2e-5; bmin,bmax=0.5,8.
            for step in range(500):
                beta=bmin*(bmax/bmin)**min(step/250.,1.)
                m.zero_grad(); raw=torch.nn.functional.cross_entropy(m(Xt),yt)
                # max(raw-eps,0) has the intended zero-loss stratum gradient.
                objective=torch.relu(raw-eps); objective.backward()
                with torch.no_grad():
                    for p in m.parameters():
                        p.add_(-eta*beta*p.grad + math.sqrt(2*eta)*torch.randn_like(p))
            phase='500 Langevin steps'
        else: phase='none'
        tr=loss_acc(m,Xt,yt); te=loss_acc(m,Xe,ye)
        results[method]={'pre_phase_train_loss':pre_train[0],'pre_phase_test_loss':pre_test[0],
                         'train_loss':tr[0],'train_acc':tr[1],'test_loss':te[0],'test_acc':te[1],
                         'perturbed_train_loss':perturb_loss(m,Xt,yt), 'phase':phase}
    return results

def main():
    seed()
    try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'); toy=toy_check(); exp=train_compare(device)
    except Exception as e:
        device=torch.device('cpu'); seed(); toy=toy_check(); exp=train_compare(device)
        exp['fallback_error']=repr(e)
    result={'device':str(device),'toy':toy,'experiment':exp}
    os.makedirs('artifacts',exist_ok=True)
    with open('artifacts/results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()