import sys, json, time
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, train_model, evaluate, sweep_baseline, make_report
from poisson_symmetry_track import get_dataset as get_pde, META

SEEDS = tuple(range(8))
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 16
BATCH = 64
N = 8
D = N*N

def laplacian(n=N):
    A = torch.zeros(D, D)
    for i in range(n):
        for j in range(n):
            q=i*n+j; A[q,q]=4
            for ii,jj in [((i+1)%n,j),((i-1)%n,j),(i,(j+1)%n),(i,(j-1)%n)]:
                A[q,ii*n+jj] -= 1
    return A

L = laplacian()
K = L + 1e-5*torch.eye(D)
LAM = (4 - 2*torch.cos(2*torch.pi*torch.arange(N)/N)[:,None]
       - 2*torch.cos(2*torch.pi*torch.arange(N)/N)[None,:]).float()

class PoissonNet(nn.Module):
    def __init__(self, idea=False):
        super().__init__(); self.idea=idea
        self.trunk=nn.Sequential(nn.Linear(D,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh(),nn.Linear(64,D))
    def solve(self, b):
        if not self.idea:
            return torch.linalg.solve(K.to(device=b.device, dtype=b.dtype), b.unsqueeze(-1)).squeeze(-1)
        bh=torch.fft.fft2(b.reshape(-1,N,N), dim=(-2,-1))
        uh=bh/(LAM.to(device=b.device, dtype=b.dtype)+1e-5); uh[...,0,0]=0
        return torch.fft.ifft2(uh, dim=(-2,-1)).real.reshape(-1,D)
    def forward(self,x): return self.solve(self.trunk(x))

def ds(seed):
    a=get_pde(seed,400,200)
    return {k: torch.as_tensor(v, dtype=torch.float32) if k in ('xtr','ytr','xte','yte') else v for k,v in a.items()}

def factory(idea, cfg):
    def fn(seed):
        torch.manual_seed(seed); np.random.seed(seed)
        net=PoissonNet(idea)
        _, metric, _=train_model(net, ds(seed), epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_:None)
        return float(metric)
    return fn

def main():
    # All idea learning rates are included in the baseline sweep (search-space parity).
    base=sweep_baseline(lambda cfg: factory(False, cfg), [{'lr':lr} for lr in LRS], seeds=SEEDS)
    runs=[]
    for lr in LRS:
        vals=[factory(True, {'lr':lr})(s) for s in SEEDS]
        runs.append({'lr':lr,'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)})
    chosen=min(runs,key=lambda z:z['mean'])
    idea=dict(chosen); idea['sweep']=runs; idea['cfg']={'lr':chosen['lr'],'nearby_tested':LRS}
    # Core trained-model signature: compare observed residual and equivariance of its output.
    torch.manual_seed(0); net=PoissonNet(True)
    trained, metric, _=train_model(net, ds(0), epochs=EPOCHS, lr=chosen['lr'], batch=BATCH, log=lambda *_:None)
    with torch.no_grad():
        dev=next(trained.parameters()).device; x=ds(0)['xte'][:32].to(dev); pred=trained(x)
        Kd=K.to(device=dev, dtype=pred.dtype)
        residual=((pred@Kd.T)-trained.trunk(x)).norm(dim=1)/(trained.trunk(x).norm(dim=1)+1e-8)
        # Translation action is a mesh automorphism; test output equivariance on the trained net.
        xx=x.reshape(-1,N,N); shifted=torch.roll(xx,1,dims=1).reshape(-1,D)
        equiv=(trained(shifted)-torch.roll(pred.reshape(-1,N,N),1,dims=1).reshape(-1,D)).norm(dim=1)/(pred.norm(dim=1)+1e-8)
        comm=float(torch.linalg.norm(L@torch.roll(torch.eye(D),1,dims=0)-torch.roll(torch.eye(D),1,dims=0)@L))
        off=float((torch.fft.fft2(L.reshape(N,N,N,N).diagonal(dim1=1,dim2=3),dim=(-2,-1)) if False else torch.tensor(0.)).item())
    sig={'prediction':'symmetry-basis implicit solve preserves the discretized solution and commutes with translations',
         'trained_model_observed_max_relative_solve_residual':float(residual.max()),
         'trained_model_observed_max_translation_equivariance_error':float(equiv.max()),
         'operator_commutator_frobenius':comm,
         'predicted_residual_and_equivariance':0.0,
         'confirmed':bool(float(residual.max())<2e-4 and float(equiv.max())<2e-4 and comm<1e-6),
         'note':'signature is measured from the trained idea model on held-out benchmark inputs'}
    report=make_report('symmetric_periodic_poisson','matched_mlp_implicit_poisson',base,idea,sig)
    report['custom_track']={'name':META['name'],'file':'poisson_symmetry_track.py','domain':META['domain']}
    report['stage2_config']={'epochs':EPOCHS,'batch':BATCH,'seeds':list(SEEDS),'lr_union':LRS,'track_reason':'PDE inverse problem with periodic Poisson operator; no built-in track had this structure'}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()