import json, math, random
from pathlib import Path
import numpy as np
from scipy.integrate import quad
import torch
from torch import nn

SEED = 282
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

# Positive Mellin mixture h(a)=sum softplus(w)(1+softplus(s))^-a.
class MellinGate(nn.Module):
    def __init__(self, L=6, hidden=8):
        super().__init__()
        self.L = L
        self.wraw = nn.Parameter(torch.zeros(L))
        self.sraw = nn.Parameter(torch.linspace(-2.5, 1.5, L))
        self.base = nn.Sequential(nn.Linear(1, hidden), nn.Tanh(), nn.Linear(hidden, 1))
    def forward(self, z, a):
        b = torch.nn.functional.softplus(self.base(z))
        w = torch.nn.functional.softplus(self.wraw)
        s = torch.nn.functional.softplus(self.sraw)
        h = torch.exp(torch.log(w)[None, :] - a[:, None] * torch.log1p(s)[None, :]).sum(1, keepdim=True)
        return b * h

class FreeGate(nn.Module):
    def __init__(self, hidden=12):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, 1), nn.Softplus())
    def forward(self, z, a):
        return self.net(torch.cat([z, a[:, None]], 1))

def gamma_identity_check():
    # Adaptive quadrature verifies the gamma-normalized identity for an atomic measure.
    mu = np.array([0.7, 1.2, 0.4]); rates = np.array([0.0, 0.8, 3.0])
    aa = np.array([0.6, 1.0, 2.5, 5.0])
    lhs = []
    for a in aa:
        f = lambda y: y**(a-1) * np.exp(-y) * np.sum(mu*np.exp(-rates*y))
        lhs.append(quad(f, 0.0, np.inf, epsabs=1e-12, epsrel=1e-12)[0] / math.gamma(a))
    lhs = np.array(lhs)
    rhs = np.sum(mu[None,:] * (1+rates[None,:])**(-aa[:,None]), axis=1)
    return float(np.max(np.abs(lhs-rhs))), lhs.tolist(), rhs.tolist()

def structure_check():
    w = np.array([0.4, 1.1, 0.8, 0.3]); s = np.array([0., .2, 1.5, 6.])
    def h(a): return np.sum(w * (1+s)**(-a))
    grid = np.linspace(.3, 8., 100)
    # finite differences: (-1)^n h^(n) >= 0, and log convexity h h''-(h')^2 >=0.
    deriv_min = []
    for n in range(5):
        exact = np.sum(w[None,:] * (-np.log1p(s)[None,:])**n * (1+s)[None,:]**(-grid[:,None]), axis=1)
        deriv_min.append(float(np.min(((-1)**n)*exact)))
    a0, d = .7, .45
    H = np.array([[h(a0+(i+j)*d) for j in range(4)] for i in range(4)])
    eig = np.linalg.eigvalsh(H)
    hp = np.sum(w[None,:] * (-np.log1p(s)[None,:]) * (1+s)[None,:]**(-grid[:,None]), axis=1)
    hpp = np.sum(w[None,:] * np.log1p(s)[None,:]**2 * (1+s)[None,:]**(-grid[:,None]), axis=1)
    logconv_min = float(np.min(np.array([h(x) for x in grid])*hpp-hp**2))
    return deriv_min, float(eig.min()), logconv_min

def fit_model(model, ztr, atr, ytr, zte, ate, yte, steps=1800):
    opt = torch.optim.Adam(model.parameters(), lr=.025)
    for _ in range(steps):
        pred = model(ztr, atr)
        loss = ((pred-ytr)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        train = float(((model(ztr,atr)-ytr)**2).mean())
        test = float(((model(zte,ate)-yte)**2).mean())
    return train, test, sum(p.numel() for p in model.parameters())

def toy_experiment():
    # Feature is constant; target is a positive Mellin response, with observation noise.
    rng = np.random.default_rng(SEED)
    z = torch.zeros(180,1)
    a = torch.linspace(0.5, 6.0, 180)
    truew = torch.tensor([.25,.7,.5]); trues = torch.tensor([.05,.8,4.0])
    target = sum(truew[i]*(1+trues[i])**(-a) for i in range(3))
    noisy = target + torch.tensor(rng.normal(0,.012, len(a)), dtype=torch.float32)
    trainmask = (a <= 4.5)
    ztr,atr,ytr = z[trainmask],a[trainmask],noisy[trainmask,None]
    zte,ate,yte = z[~trainmask],a[~trainmask],target[~trainmask,None]
    # Equal-ish parameter count: free MLP has 49, Mellin has 6+6+ (1*8+8+8+1)=33.
    # Use a scalar offset feature to make both models genuinely order-dependent.
    torch.manual_seed(SEED)
    free = fit_model(FreeGate(hidden=9), ztr,atr,ytr,zte,ate,yte)
    torch.manual_seed(SEED)
    mellin = fit_model(MellinGate(L=6,hidden=8), ztr,atr,ytr,zte,ate,yte)
    return {"free_mlp": free, "positive_mellin": mellin, "train_orders": int(trainmask.sum()), "extrapolation_orders": int((~trainmask).sum())}

def main():
    ident = gamma_identity_check(); struct = structure_check(); exp = toy_experiment()
    out = {"gamma_identity_max_abs_error": ident[0], "identity_lhs": ident[1], "identity_rhs": ident[2], "derivative_nonnegative_min_by_order": struct[0], "hankel_min_eigenvalue": struct[1], "log_convexity_min": struct[2], "toy": exp}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()