import json, math, random
from pathlib import Path
import numpy as np
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

SEED = 464
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

# ---------- core fractional Mahalanobis radial head ----------
class FractionalRadialHead(nn.Module):
    def __init__(self, dim, exponents=(.25, .5, 1., 1.5, 2.), eps=1e-3):
        super().__init__()
        self.dim, self.exponents, self.eps = dim, tuple(exponents), eps
        self.mu = nn.Parameter(torch.zeros(2, dim))
        self.raw_diag = nn.Parameter(torch.zeros(2, dim))
        self.offdiag = nn.Parameter(torch.zeros(2, dim, dim))
        self.b = nn.Parameter(torch.zeros(len(exponents)))
        self.bias = nn.Parameter(torch.tensor(0.0))

    def factors(self):
        L = torch.tril(self.offdiag, diagonal=-1)
        diag = torch.nn.functional.softplus(self.raw_diag) + 1e-3
        return L + torch.diag_embed(diag)

    def covariances(self):
        L = self.factors()
        return L @ L.transpose(-1, -2) + self.eps * torch.eye(self.dim, device=L.device)

    def radii(self, z):
        # Cholesky(Sigma) followed by triangular solves avoids an explicit inverse.
        S = self.covariances()
        C = torch.linalg.cholesky(S)
        out = []
        for c in range(2):
            delta = z - self.mu[c]
            v = torch.linalg.solve_triangular(C[c], delta.T, upper=False).T
            out.append((v*v).sum(1))
        return torch.stack(out, 1)

    def forward(self, z):
        r = self.radii(z).clamp_min(1e-8)
        powers = torch.stack([r ** a for a in self.exponents], dim=-1)
        h = (powers * self.b).sum(-1)
        # b0 cancels. The remaining constant is the determinant/prior term.
        logdet = torch.linalg.slogdet(self.covariances())[1]
        return h[:, 1] - h[:, 0] + self.bias + .5 * (logdet[0] - logdet[1])


def init_radial(head, x, y):
    with torch.no_grad():
        for c in (0, 1):
            q = x[y == c]
            head.mu[c].copy_(q.mean(0))
            cov = torch.cov(q.T) + .05 * torch.eye(x.shape[1], device=x.device)
            L = torch.linalg.cholesky(cov)
            head.raw_diag[c].copy_(torch.log(torch.expm1(torch.diag(L).clamp_min(.002))))
            head.offdiag[c].copy_(torch.tril(L, diagonal=-1))
        # Start close to a Gaussian/QDA radial coefficient and let fractional terms adapt.
        head.b.zero_(); head.b[2] = -.5

class LinearHead(nn.Module):
    def __init__(self, d):
        super().__init__(); self.w = nn.Linear(d, 1)
    def forward(self, z): return self.w(z).squeeze(1)

class MLPHead(nn.Module):
    def __init__(self, d):
        super().__init__(); self.net = nn.Sequential(nn.Linear(d, 8), nn.Tanh(), nn.Linear(8, 1))
    def forward(self, z): return self.net(z).squeeze(1)

class QDAHead(nn.Module):
    def __init__(self, d):
        super().__init__(); self.d = d
    def fit(self, z, y):
        with torch.no_grad():
            self.mu0 = z[y < .5].mean(0); self.mu1 = z[y >= .5].mean(0)
            self.s0 = torch.cov(z[y < .5].T) + .05*torch.eye(self.d, device=z.device)
            self.s1 = torch.cov(z[y >= .5].T) + .05*torch.eye(self.d, device=z.device)
            self.i0 = torch.linalg.inv(self.s0); self.i1 = torch.linalg.inv(self.s1)
            self.const = .5*(torch.linalg.slogdet(self.s0)[1]-torch.linalg.slogdet(self.s1)[1])
    def forward(self, z):
        a=z-self.mu0; b=z-self.mu1
        r0=(a @ self.i0 * a).sum(1); r1=(b @ self.i1 * b).sum(1)
        return -.5*r1 + .5*r0 + self.const

# ---------- data and training ----------
def make_data(kind, n=5000):
    y = np.random.randint(0, 2, n)
    means = np.where(y[:, None] == 0, [-1.15, 0.0], [1.15, 0.0])
    # Slightly distinct anisotropic class geometry makes the radius distinction visible.
    scales = np.where(y[:, None] == 0, [1.0, .65], [1.0, .65])
    if kind == 'student':
        df = 2.5
        raw = np.random.randn(n, 2) / np.sqrt(np.random.chisquare(df, n)[:, None] / df)
    else:
        raw = np.random.randn(n, 2)
    x = means + raw * scales
    # fixed stratified-ish split by random permutation
    p = np.random.RandomState(SEED + (1 if kind == 'student' else 2)).permutation(n)
    cut = int(.7*n)
    return (torch.tensor(x[p[:cut]], dtype=torch.float32, device=DEVICE),
            torch.tensor(y[p[:cut]], dtype=torch.float32, device=DEVICE),
            torch.tensor(x[p[cut:]], dtype=torch.float32, device=DEVICE),
            torch.tensor(y[p[cut:]], dtype=torch.float32, device=DEVICE))

def fit(model, x, y, radial=False, epochs=180):
    model.to(DEVICE); model.train()
    if radial: init_radial(model, x, y.long())
    opt = torch.optim.Adam(model.parameters(), lr=.025 if radial else .01, weight_decay=1e-4)
    for _ in range(epochs):
        opt.zero_grad(); logits = model(x)
        loss = nn.functional.binary_cross_entropy_with_logits(logits, y)
        if radial:
            L = model.factors()
            # modest covariance and coefficient regularization for stable learned geometry
            loss = loss + 1e-5*(L*L).sum() + 1e-5*(model.b*model.b).sum()
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
    model.eval()
    with torch.no_grad():
        logits = model(x_test_global)
        yy = y_test_global
        nll = nn.functional.binary_cross_entropy_with_logits(logits, yy).item()
        acc = ((logits > 0) == (yy > .5)).float().mean().item()
    return nll, acc

# Exact Student-t radial log generator and basis approximation sanity check.
def math_check():
    t = np.geomspace(1e-5, 1e3, 500)
    df = 2.5
    exact = -0.5*(df+2)*np.log1p(t/df)
    exps = np.array([.25,.5,1.,1.5,2.])
    A = np.column_stack([np.ones_like(t)] + [t**a for a in exps])
    # fit on the numerically relevant log-radius range; report in-range error
    grid = np.geomspace(1e-4, 100, 400)
    Ag = np.column_stack([np.ones_like(grid)] + [grid**a for a in exps])
    eg = -0.5*(df+2)*np.log1p(grid/df)
    coef, *_ = np.linalg.lstsq(Ag, eg, rcond=None)
    fit = A @ coef
    ii = (t >= 1e-4) & (t <= 100)
    # Cholesky solve and explicit inverse agree for an SPD test matrix.
    S = np.array([[1.4,.25],[.25,.8]])
    L = np.linalg.cholesky(S); delta = np.array([.7,-1.1])
    r_solve = np.sum(np.linalg.solve(L, delta)**2)
    r_inv = delta @ np.linalg.inv(S) @ delta
    return {'basis_max_abs_error_t_1e-4_to_100': float(np.max(np.abs(fit[ii]-exact[ii]))),
            'basis_rmse_t_1e-4_to_100': float(np.sqrt(np.mean((fit[ii]-exact[ii])**2))),
            'cholesky_vs_inverse_abs_error': float(abs(r_solve-r_inv)),
            'all_basis_radii_finite': bool(np.isfinite(A).all())}

if __name__ == '__main__':
    result = {'device': DEVICE, 'math_check': math_check(), 'datasets': {}}
    for kind in ('gaussian', 'student'):
        xtr, ytr, xte, yte = make_data(kind)
        x_test_global, y_test_global = xte, yte
        vals = {}
        qda = QDAHead(2).to(DEVICE); qda.fit(xtr, ytr)
        with torch.no_grad():
            qlog = qda(xte); qloss = nn.functional.binary_cross_entropy_with_logits(qlog, yte).item()
            vals['qda_frozen'] = {'test_logloss': qloss, 'test_accuracy': float(((qlog > 0)==(yte>.5)).float().mean())}
        for name, cls, radial in [('linear', LinearHead, False), ('mlp', MLPHead, False), ('fractional_radial', lambda d: FractionalRadialHead(d), True)]:
            torch.manual_seed(SEED + len(name) + (0 if kind=='gaussian' else 10))
            model = cls(2)
            vals[name] = dict(zip(('test_logloss','test_accuracy'), fit(model, xtr, ytr, radial)))
            if radial:
                with torch.no_grad():
                    vals[name]['max_radius_test'] = float(model.radii(xte).max().cpu())
        result['datasets'][kind] = vals
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))