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

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

META = {
    'name': 'burgers_periodic_split',
    'domain': 'pde',
    'description': 'Periodic underresolved Burgers one-step regression on a 32-point grid.'
}


def _D(v):
    n = v.shape[-1]
    # x=2*pi*j/n, so Fourier integer wavenumbers are fftfreq(n)*n.
    k = torch.fft.rfftfreq(n, d=1.0/n).to(v.device)
    return torch.fft.irfft(1j * k * torch.fft.rfft(v, dim=-1), n=n, dim=-1)


def split_quadratic(u, alpha):
    beta = 1.0 - float(alpha)
    return float(alpha) * _D(0.5 * u * u) + beta * u * _D(u)


class SplitQuadraticNet(nn.Module):
    def __init__(self, alpha):
        super().__init__()
        self.alpha = float(alpha)
        self.inp = nn.Conv1d(1, 16, 1)
        self.mix = nn.Conv1d(16, 16, 1)
        self.out = nn.Conv1d(16, 1, 1)
        self.gamma = nn.Parameter(torch.tensor(0.1))

    def forward(self, x):
        z = torch.tanh(self.inp(x))
        h = torch.tanh(self.mix(z))
        # The only intervention is the split coefficient in this quadratic block.
        return self.out(h + self.gamma * split_quadratic(h, self.alpha))


def get_dataset(seed, n_train=400, n_test=160):
    rng = np.random.default_rng(int(seed)); n = 32
    x = np.arange(n) * 2 * np.pi / n

    def make(m):
        xs, ys = [], []
        for _ in range(int(m)):
            u = np.zeros(n)
            # Deliberately weakly underresolved high modes expose product aliasing.
            for mode in rng.choice(np.arange(5, 16), size=4, replace=False):
                u += rng.normal(0, 0.16 / np.sqrt(mode)) * np.sin(mode*x + rng.uniform(0, 2*np.pi))
                u += rng.normal(0, 0.16 / np.sqrt(mode)) * np.cos(mode*x + rng.uniform(0, 2*np.pi))
            dt = 0.025
            uh = np.fft.rfft(u); k = np.fft.rfftfreq(n, d=1.0/n)
            du = np.fft.irfft(1j*k*uh, n=n)
            # Target is a stable, filtered Burgers Euler step generated by padded product.
            mgrid = 48; up = np.zeros(mgrid, complex); q=n//2
            up[:q] = uh[:q] * (mgrid/n); up[-q:] = uh[-q:] * (mgrid/n)
            vp = np.fft.irfft(up, n=mgrid)
            kp=np.fft.rfftfreq(mgrid,d=1.0/mgrid)
            flux=np.fft.irfft(1j*kp*np.fft.rfft(0.5*vp*vp),n=mgrid)
            fluxh=np.fft.rfft(flux); trunc=np.zeros(n,complex)
            trunc[:q]=fluxh[:q]*(n/mgrid); trunc[-q:]=fluxh[-q:]*(n/mgrid)
            target=u-dt*(np.fft.irfft(trunc,n=n))
            xs.append(u.astype('float32')); ys.append(target.astype('float32'))
        return np.asarray(xs)[:,None,:], np.asarray(ys)[:,None,:]

    xtr,ytr=make(n_train); xte,yte=make(n_test)
    return {'xtr':xtr,'ytr':ytr,'xte':xte,'yte':yte,'task':'regression','metric':'mse',
            'input_shape':(1,n),'out_dim':n}


def run():
    torch.set_num_threads(4)
    # All idea learning rates are also evaluated for the baseline.
    lrs = [1e-3, 3e-3, 1e-2]
    grid = [{'lr': lr, 'alpha': 1.0} for lr in lrs]
    idea_grid = [{'lr': lr, 'alpha': 1/3} for lr in lrs]

    def factory(cfg):
        def train(seed):
            torch.manual_seed(seed); np.random.seed(seed)
            d=get_dataset(seed,400,160)
            d={k:(torch.from_numpy(v) if isinstance(v,np.ndarray) else v) for k,v in d.items()}
            net,metric,_=train_model(SplitQuadraticNet(cfg['alpha']),d,epochs=18,lr=cfg['lr'],batch=128,log=lambda *_:None)
            return float(metric) if metric is not None else float('inf')
        return train

    base=sweep_baseline(factory,grid)
    # Equal-sized idea sweep; its selected setting is evaluated on all eight paired seeds.
    idea_sweep=sweep_baseline(factory,idea_grid)
    idea=idea_sweep['full']
    rep=make_report('burgers_periodic_split','custom_split_quadratic_net',base,idea,{
        'predicted': {'aliasing_root_alpha_for_tau4_i1_j1': 1/3,
                      'prediction': 'alpha=1/3 reduces the dominant unresolved quadratic-product aliasing proxy'},
        'observed': {'baseline_alpha': 1.0, 'idea_alpha': 1/3,
                     'baseline_test_mse': base['full']['mean'],
                     'idea_test_mse': idea['mean'],
                     'idea_sweep': idea_sweep['sweep']},
        'confirmed': False,
        'measurement_note': 'The trained-model aliasing proxy was not retained by the generic harness metric path; confirmation is false rather than inferred from the formula.'
    })
    rep['custom_track']={'name':META['name'],'file':'split_bench.py','domain':'pde'}
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__ == '__main__': run()