import json, math, random
from pathlib import Path
import numpy as np

SEED = 913
random.seed(SEED)
np.random.seed(SEED)

# Memory-retaining profile h(x, eta)=eta^alpha F(x/eta^beta).
def profile(x, eta, alpha, beta):
    u = x / (eta ** beta)
    return eta ** alpha * np.exp(-0.5*u*u) * (1.0 + 0.35*np.cos(3.0*u))

def effective_width(x, y):
    w = np.maximum(y, 0.0)**2
    return float(np.sqrt(np.sum(w*x*x)/(np.sum(w)+1e-30)))

def log_slope(a, b):
    return float(np.polyfit(np.log(a), np.log(b), 1)[0])

def math_sweep():
    alpha, beta = 0.70, 0.62
    etas = np.geomspace(0.35, 3.0, 13)
    x = np.linspace(-10, 10, 20001)
    peaks = np.array([np.max(profile(x,e,alpha,beta)) for e in etas])
    widths = np.array([effective_width(x, profile(x,e,alpha,beta)) for e in etas])
    alpha_hat = log_slope(etas, peaks)
    beta_hat = log_slope(etas, widths)

    # Prediction 1: amplitude slope is alpha.
    # Prediction 2: width slope is beta.
    # Prediction 3: with x'=x/L and eta'=eta/L^(1/beta),
    # h(x',eta') = L^(-alpha/beta) h(x,eta), so normalized residual is zero.
    eta = 1.1
    exact_residuals, wrong_eta_residuals = [], []
    for L in [1.25, 1.5, 2.0, 3.0]:
        ep = eta / (L ** (1.0/beta))
        lhs = profile(x/L, ep, alpha, beta)
        rhs = (L ** (-alpha/beta))*profile(x, eta, alpha, beta)
        exact_residuals.append(np.linalg.norm(lhs-rhs)/(np.linalg.norm(rhs)+1e-12))
        ep_wrong = eta/L
        lhs_wrong = profile(x/L, ep_wrong, alpha, beta)
        wrong_eta_residuals.append(np.linalg.norm(lhs_wrong-rhs)/(np.linalg.norm(rhs)+1e-12))

    return {
        "true_alpha": alpha, "fitted_alpha": alpha_hat,
        "true_beta": beta, "fitted_beta": beta_hat,
        "amplitude_abs_error": abs(alpha_hat-alpha),
        "width_abs_error": abs(beta_hat-beta),
        "rg_L": [1.25,1.5,2.0,3.0],
        "rg_exact_residual": exact_residuals,
        "rg_wrong_eta_residual": wrong_eta_residuals,
        "predictions": [
            "log peak amplitude slope equals alpha",
            "log profile width slope equals beta",
            "matched RG rescaling gives zero residual while eta/L does not"
        ]
    }

def mini_experiment():
    # A deliberately tiny signal classification task: class is encoded by a
    # narrow/wide Gaussian and amplitude. Inputs are randomly shifted profiles.
    try:
        import torch
        import torch.nn as nn
        from torch.utils.data import DataLoader, TensorDataset
        torch.manual_seed(SEED)
        device = "cpu"
        # CPU is the reliable fallback when shared CUDA convolution engines fail.
    except Exception as e:
        return {"error": "torch unavailable: "+str(e)}

    ntrain, nval, length = 1600, 500, 64
    def make(n, heldout=False):
        X, Y, E = [], [], []
        for _ in range(n):
            y = np.random.randint(0,2)
            # Validation uses a shifted scale range, testing extrapolation.
            eta = np.random.uniform(0.55,0.95) if (not heldout and y==0) else None
            if eta is None:
                eta = np.random.uniform(1.05,1.55) if heldout else np.random.uniform(1.05,1.45)
            amp = 1.0 if y==0 else 1.55
            xx = np.linspace(-3,3,length)
            shift = np.random.uniform(-0.25,0.25)
            sig = eta*(0.70 if y==0 else 0.95)
            z = amp*np.exp(-0.5*((xx-shift)/sig)**2)
            z += np.random.normal(0,0.045,length)
            X.append(z.astype(np.float32)); Y.append(y); E.append(eta)
        return torch.tensor(np.array(X)), torch.tensor(Y), torch.tensor(np.array(E),dtype=torch.float32)
    Xtr,Ytr,Etr = make(ntrain,False); Xv,Yv,Ev = make(nval,True)
    train = DataLoader(TensorDataset(Xtr,Ytr,Etr), batch_size=64, shuffle=True)

    class Baseline(nn.Module):
        def __init__(self):
            super().__init__(); self.net=nn.Sequential(nn.Conv1d(1,12,5,stride=2,padding=2),nn.ReLU(),nn.Conv1d(12,20,5,stride=2,padding=2),nn.ReLU(),nn.AdaptiveAvgPool1d(1)); self.fc=nn.Linear(20,2)
        def forward(self,x,e): return self.fc(self.net(x[:,None,:]).squeeze(-1))
    class Retained(nn.Module):
        def __init__(self):
            super().__init__(); self.conv=nn.Sequential(nn.Conv1d(1,12,5,stride=2,padding=2),nn.ReLU(),nn.Conv1d(12,20,5,stride=2,padding=2),nn.ReLU()); self.film=nn.Linear(1,40); self.fc=nn.Linear(20,2)
        def forward(self,x,e):
            h=self.conv(x[:,None,:]); ab=self.film(torch.log(e[:,None]+1e-5)); a,b=ab[:,:20,None],ab[:,20:,None]; h=h*(1+a)+b; return self.fc(h.mean(-1))
    def train_eval(model):
        model.to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3); lossfn=nn.CrossEntropyLoss()
        for _ in range(10):
            model.train()
            for x,y,e in train:
                x,y,e=x.to(device),y.to(device),e.to(device); opt.zero_grad(); loss=lossfn(model(x,e),y); loss.backward(); opt.step()
        model.eval(); correct=0; total=0; losses=[]
        with torch.no_grad():
            for x,y,e in DataLoader(TensorDataset(Xv,Yv,Ev),batch_size=128):
                out=model(x.to(device),e.to(device)); losses.append(float(lossfn(out,y.to(device)))); correct += int((out.argmax(1).cpu()==y).sum()); total += len(y)
        return {"accuracy":correct/total,"loss":float(np.mean(losses))}
    try:
        base=train_eval(Baseline()); idea=train_eval(Retained())
    except Exception as exc:
        # Retry from scratch on CPU for shared/unsupported CUDA backends.
        device = "cpu"
        base=train_eval(Baseline()); idea=train_eval(Retained())
    return {"device":device,"baseline":base,"idea":idea,"train_size":ntrain,"validation_size":nval}

def main():
    result={"math":math_sweep(),"mini_experiment":mini_experiment()}
    Path("results.json").write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__ == "__main__": main()