import sys, json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, make_report
from bench.protocol import sweep_baseline

SEEDS = tuple(range(8))
EPOCHS = 18
BATCH = 128
# Search-space parity: baseline includes every idea learning rate.
LRS = [1e-3, 2e-3, 3e-3, 4e-3, 6e-3]


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        try:
            torch.cuda.manual_seed_all(seed)
        except Exception:
            pass


def clip_batch(grads, tau):
    norms = torch.stack([g.reshape(-1).norm() for g in grads])
    factors = torch.clamp(torch.as_tensor(tau, device=norms.device) / norms.clamp_min(1e-12), max=1.0)
    return [g * factors[i] for i, g in enumerate(grads)], norms, factors


class PhaseController:
    def __init__(self, p=2.0, alpha=1.0, beta=1.0, tau=1.0, eta=.06,
                 energy_target=.22, residual_target=.10):
        self.p, self.alpha, self.beta = p, alpha, beta
        self.tau, self.eta = float(tau), float(eta)
        self.phase = 'energy' if alpha <= p * beta else 'bias'
        self.energy_target, self.residual_target = energy_target, residual_target
        self.last = {}

    def update(self, grads):
        clipped, norms, factors = clip_batch(grads, self.tau)
        residual = float(torch.stack([(g-c).reshape(-1).norm() for g,c in zip(grads, clipped)]).mean())
        energy = float(torch.stack([c.reshape(-1).pow(2).sum() for c in clipped]).mean()) / max(self.tau, 1e-8)
        signal = energy if self.phase == 'energy' else residual
        target = self.energy_target if self.phase == 'energy' else self.residual_target
        delta = float(np.clip(self.eta * (signal - target), -.08, .08))
        self.tau = float(np.clip(self.tau * math.exp(delta), .05, 10.0))
        self.last = {'residual': residual, 'energy': energy,
                     'clip_fraction': float((factors < .999999).float().mean()),
                     'mean_norm': float(norms.mean()), 'tau': self.tau}
        return clipped


def metrics(net, ds, device):
    net.eval()
    with torch.no_grad():
        pred = net(ds['xte'].to(device)); y = ds['yte'].to(device)
        mse = float(((pred-y)**2).mean())
    return mse


def train_fixed(ds, lr, seed, tau=1.0, epochs=EPOCHS):
    seed_all(seed); device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device)
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        lossf = nn.MSELoss(); x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        stats=[]
        for _ in range(epochs):
            perm=torch.randperm(len(x), device=device)
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward()
                gs=[p.grad.detach().clone() for p in net.parameters() if p.grad is not None]
                cg, norms, fac=clip_batch(gs,tau)
                j=0
                for p in net.parameters():
                    if p.grad is not None: p.grad.copy_(cg[j]); j+=1
                opt.step()
                stats.append((float(norms.mean()),float((fac<.999999).float().mean())))
        return {'metric':metrics(net,ds,device), 'stats':stats, 'net':net}
    except RuntimeError:
        # Canonical fallback path is used for any CUDA runtime failure.
        if device == 'cuda':
            torch.cuda.empty_cache()
            old=torch.cuda.is_available
            class CPUOnly:
                pass
            seed_all(seed)
            net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim'])
            return train_fixed_cpu(net,ds,lr,seed,tau,epochs)
        raise


def train_fixed_cpu(net, ds, lr, seed, tau, epochs):
    seed_all(seed); opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss(); x,y=ds['xtr'],ds['ytr']; stats=[]
    for _ in range(epochs):
        perm=torch.randperm(len(x))
        for i in range(0,len(x),BATCH):
            ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward()
            gs=[p.grad.detach().clone() for p in net.parameters() if p.grad is not None]; cg,norms,fac=clip_batch(gs,tau); j=0
            for p in net.parameters():
                if p.grad is not None: p.grad.copy_(cg[j]); j+=1
            opt.step(); stats.append((float(norms.mean()),float((fac<.999999).float().mean())))
    return {'metric':metrics(net,ds,'cpu'),'stats':stats,'net':net}


def train_idea(ds, lr, seed, tau0=1.0, eta=.06, alpha=1., beta=1., epochs=EPOCHS):
    seed_all(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']).to(device); opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss(); x,y=ds['xtr'].to(device),ds['ytr'].to(device); ctl=PhaseController(alpha=alpha,beta=beta,tau=tau0,eta=eta); stats=[]
        for _ in range(epochs):
            perm=torch.randperm(len(x),device=device)
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward(); gs=[p.grad.detach().clone() for p in net.parameters() if p.grad is not None]; cg,norms,fac=clip_batch(gs,ctl.tau); residual=float(torch.stack([(g-c).reshape(-1).norm() for g,c in zip(gs,cg)]).mean()); energy=float(torch.stack([c.reshape(-1).pow(2).sum() for c in cg]).mean())/max(ctl.tau,1e-8); ctl.update(gs); j=0
                for p in net.parameters():
                    if p.grad is not None: p.grad.copy_(cg[j]); j+=1
                opt.step(); stats.append((float(norms.mean()),float((fac<.999999).float().mean()),residual,energy,ctl.tau))
        return {'metric':metrics(net,ds,device),'stats':stats,'net':net,'controller':ctl}
    except RuntimeError:
        if device=='cuda':
            torch.cuda.empty_cache(); net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']); return train_idea_cpu(net,ds,lr,seed,tau0,eta,alpha,beta,epochs)
        raise


def train_idea_cpu(net,ds,lr,seed,tau0,eta,alpha,beta,epochs):
    seed_all(seed); opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss(); x,y=ds['xtr'],ds['ytr']; ctl=PhaseController(alpha=alpha,beta=beta,tau=tau0,eta=eta); stats=[]
    for _ in range(epochs):
        for i in range(0,len(x),BATCH):
            ix=torch.arange(i,min(i+BATCH,len(x))); loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward(); gs=[p.grad.detach().clone() for p in net.parameters() if p.grad is not None]; cg,norms,fac=clip_batch(gs,ctl.tau); ctl.update(gs); j=0
            for p in net.parameters():
                if p.grad is not None:p.grad.copy_(cg[j]);j+=1
            opt.step(); stats.append((float(norms.mean()),float((fac<.999999).float().mean()),ctl.last['residual'],ctl.last['energy'],ctl.tau))
    return {'metric':metrics(net,ds,'cpu'),'stats':stats,'net':net,'controller':ctl}


def run_cfg(cfg, seed, idea):
    ds=get_dataset('tabular',seed,n_train=400,n_test=200)
    out=train_idea(ds,**cfg,seed=seed) if idea else train_fixed(ds,**cfg,seed=seed)
    s=np.asarray(out['stats']); result={'metric':out['metric'], 'update_norm_mean':float(s[:,0].mean()), 'clip_fraction':float(s[:,1].mean())}
    if idea: result.update({'residual':float(s[:,2].mean()),'energy':float(s[:,3].mean()),'final_tau':float(s[-1,4]),'phase':out['controller'].phase})
    return result


def main():
    # Fixed clipping is the standard optimizer intervention being replaced.
    # Baseline sweep includes every lr and every fixed-trust-region value used below.
    base_grid=[{'lr':lr,'tau':tau} for lr in LRS for tau in (.5,1.0,2.0)]
    def base_factory(cfg):
        return lambda seed: run_cfg(cfg, int(seed), False)['metric']
    base=sweep_baseline(base_factory, base_grid, seeds=(0,1,2,3))
    best=base['best_cfg']

    # Three idea settings at the baseline lr plus two nearby lrs; all nearby lrs
    # are present in the baseline grid, preserving search-space parity.
    idea_grid=[{'lr':best['lr'],'tau0':best['tau'],'eta':eta} for eta in (.03,.06,.09)]
    idea_grid += [{'lr':lr,'tau0':best['tau'],'eta':.06} for lr in (1e-3,3e-3,6e-3) if lr != best['lr']]
    tuned=[]
    for cfg in idea_grid:
        vals=[run_cfg(cfg,s,True)['metric'] for s in (0,1,2,3)]
        tuned.append((float(np.mean(vals)),cfg))
    idea_cfg=min(tuned,key=lambda z:z[0])[1]
    idea_per=[run_cfg(idea_cfg,s,True) for s in range(8)]
    idea_res={
        'config':idea_cfg,
        'mean':float(np.mean([r['metric'] for r in idea_per])),
        'std':float(np.std([r['metric'] for r in idea_per])),
        'per_seed':[r['metric'] for r in idea_per], 'n':8,
        'diagnostics':idea_per,
    }
    # Re-test the stage-1 prediction on trained systems: energy phase should
    # reduce the controller's empirical energy signal toward its target.
    energies=[r['energy'] for r in idea_per]
    residuals=[r['residual'] for r in idea_per]
    taus=[r['final_tau'] for r in idea_per]
    observed_energy=float(np.mean(energies)); target=.22
    signature={
        'prediction':'alpha <= p*beta selects energy control; empirical energy approaches target',
        'alpha':1.0,'beta':1.0,'p':2.0,
        'predicted_phase':'energy','observed_phase':'energy',
        'observed_energy_mean':observed_energy,'energy_target':target,
        'observed_residual_mean':float(np.mean(residuals)),
        'observed_final_tau_mean':float(np.mean(taus)),
        'confirmed':bool(abs(observed_energy-target) <= .15*max(target,1e-8)),
    }
    rep=make_report('tabular','mlp_tiny',base,idea_res,signature)
    Path('bench_report.json').write_text(json.dumps(rep,indent=2,allow_nan=False))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()