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, evaluate, sweep_baseline, make_report
from bench.protocol import DEFAULT_SEEDS

V, D = .7, .015

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

def derivs(net, z):
    z = z.detach().requires_grad_(True)
    u = net(z).reshape(-1)
    g = torch.autograd.grad(u.sum(), z, create_graph=True)[0]
    ux, ut = g[:,0], g[:,1]
    gx = torch.autograd.grad(ux.sum(), z, create_graph=True)[0]
    return u, ux, ut, gx[:,0]

def points(seed, n):
    g = torch.Generator().manual_seed(seed)
    return torch.rand(n, 2, generator=g)

def weak_loss(net, z, n_patch=8, q=8, lam=.002, tau=.08):
    # Each patch uses q quadrature points and a compact tent weight. The
    # weighted local regression is the weak-form analogue of fitting A_j c=b_j.
    _, ux, ut, uxx = derivs(net, z)
    zz = z.reshape(n_patch, q, 2)
    # Local tents centered at the patch's first quadrature point, with a fixed
    # compact radius; normalization removes patch-volume scale.
    centers = zz[:, 0:1, :]
    dist = torch.abs(zz - centers)
    ph = ((1 - dist[:, :, 0] / .35).clamp_min(0) *
          (1 - dist[:, :, 1] / .35).clamp_min(0))
    A = torch.stack([ux.reshape(n_patch, q), uxx.reshape(n_patch, q)], -1)
    bb = ut.reshape(n_patch, q)
    crows=[]
    eye = torch.eye(2, device=z.device)
    for j in range(n_patch):
        w = ph[j]
        aj = A[j].detach() * w[:, None]
        bj = bb[j].detach() * w
        c = torch.linalg.solve(aj.T @ aj + .01*eye, aj.T @ bj)
        c = torch.sign(c) * torch.relu(torch.abs(c)-lam)
        crows.append(c)
    c = torch.stack(crows)
    support = c.abs() > tau
    modal = support.float().mean(0) >= .5
    cbar = torch.where(modal, c.mean(0), torch.zeros_like(c.mean(0)))
    # Retain gradients through weak A,b while stopping them through decisions.
    residual = (A * ph[:, :, None]).sum(1) @ cbar.detach() - (bb * ph).sum(1)
    residual = residual / (ph.sum(1) + 1e-6)
    consistency = ((c-cbar.detach())**2).mean()
    return residual.pow(2).mean() + .15*consistency, float(support.float().mean()), int(modal.sum())

def train(seed, cfg, idea):
    seed_all(seed)
    ds = get_dataset('pde_patch_consensus', seed, n_train=400, n_test=200)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net = make_model('mlp_tiny', tuple(ds['xtr'].shape[1:]), 1).to(device)
    except RuntimeError:
        device = 'cpu'; net = make_model('mlp_tiny', tuple(ds['xtr'].shape[1:]), 1)
    xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device)
    opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
    last = (0.0, 0)
    for ep in range(cfg['epochs']):
        net.train(); perm = torch.randperm(len(xtr), device=device)
        for ii in range(0, len(xtr), 128):
            ix = perm[ii:ii+128]
            loss = ((net(xtr[ix]).reshape(-1) - ytr[ix])**2).mean()
            z = points(seed*10000 + ep, 64).to(device)
            if idea:
                wl, agree, nterms = weak_loss(net, z, lam=cfg['lam'], tau=cfg['tau'])
                loss = loss + cfg['weight'] * wl
                last = (agree, nterms)
            else:
                _, ux, ut, uxx = derivs(net, z)
                loss = loss + cfg['weight'] * ((ut + V*ux - D*uxx)**2).mean()
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
    net.eval()
    with torch.no_grad():
        metric = float(((net(ds['xte'].to(device)).reshape(-1) - ds['yte'].to(device))**2).mean())
    z = points(seed+777, 128).to(device)
    _, ux, ut, uxx = derivs(net, z)
    point_rms = float(torch.sqrt(((ut + V*ux - D*uxx)**2).mean()).detach().cpu())
    obs_rms = math.sqrt(metric)
    return metric, {'pointwise_pde_rms': point_rms,
                    'prediction_observation_rms': obs_rms,
                    'local_support_rate': last[0], 'modal_terms': last[1]}

def make_train(cfg, idea):
    def fn(seed): return train(seed,cfg,idea)[0]
    return fn

def main():
    # Equal union: every lr and residual weight used by either method is swept on baseline.
    grid=[{'lr':lr,'weight':w,'epochs':20,'lam':lam,'tau':tau} for lr,w,lam,tau in
          [(.001,.15,.002,.08),(.003,.30,.002,.08),(.01,.15,.002,.08)]]
    base=sweep_baseline(lambda c: make_train(c,False),grid,seeds=(0,1,2,3))
    idea_candidates=[]
    for c in grid:
        r=evaluate(make_train(c,True),seeds=DEFAULT_SEEDS)
        idea_candidates.append({'cfg':c,'result':r})
    best=min(idea_candidates,key=lambda x:x['result']['mean'])
    sig=[]
    for s in DEFAULT_SEEDS:
        _,q=train(s,best['cfg'],True); sig.append(q)
    bsig=[]
    for s in DEFAULT_SEEDS:
        _,q=train(s,base['best_cfg'],False); bsig.append(q)
    means=lambda rows,k: float(np.mean([r[k] for r in rows]))
    extra={'prediction':'weak residual integration plus support consensus should reduce noisy residual variability and increase regional support agreement',
      'observed_trained_models':{
       'baseline_pointwise_pde_rms_mean':means(bsig,'pointwise_pde_rms'),
       'idea_pointwise_pde_rms_mean':means(sig,'pointwise_pde_rms'),
       'baseline_observation_rms_mean':means(bsig,'prediction_observation_rms'),
       'idea_observation_rms_mean':means(sig,'prediction_observation_rms'),
       'idea_local_support_rate_mean':means(sig,'local_support_rate'),
       'idea_modal_terms_mean':means(sig,'modal_terms')},
      'confirmed': bool(means(sig,'pointwise_pde_rms') <= means(bsig,'pointwise_pde_rms')*1.05 and means(sig,'local_support_rate')>0)}
    report=make_report('pde_patch_consensus','mlp_tiny',base,best['result'],extra)
    report['custom_track']={'name':'pde_patch_consensus','file':'pde_consensus_track.py','domain':'pde'}
    report['idea_sweep']=[{'cfg':x['cfg'],'mean':x['result']['mean']} for x in idea_candidates]
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()