import json, math, time
import numpy as np
import torch

SEED = 7
np.random.seed(SEED)
torch.manual_seed(SEED)


def scalar_mlsi_check(N=512):
    x = 2 * np.pi * np.arange(N) / N
    g = 1.2 + 0.35 * np.cos(3*x) + 0.12 * np.sin(5*x)
    dx = 2 * np.pi / N
    dg = (np.roll(g, -1) - np.roll(g, 1)) / (2 * dx)
    entropy = np.mean(g*np.log(g)) - np.mean(g)*np.log(np.mean(g))
    rhs = 0.5 * np.mean(dg*dg/g)
    return {"entropy": float(entropy), "rhs": float(rhs),
            "holds": bool(entropy <= rhs + 2e-5), "slack": float(rhs-entropy)}


def heat(a, t):
    n = a.shape[0]
    k = np.fft.fftfreq(n) * n
    multiplier = np.exp(-t*k*k)
    return np.fft.ifft(np.fft.fft(a, axis=0) * multiplier[:, None, None], axis=0).real


def bkm(rho):
    n = rho.shape[0]
    dx = 2*np.pi/n
    dr = (np.roll(rho, -1, axis=0)-np.roll(rho, 1, axis=0))/(2*dx)
    vals, vecs = np.linalg.eigh(rho)
    vals = np.maximum(vals, 1e-8)
    # transform derivative to eigenbasis, then apply inverse logarithmic mean
    at = np.einsum('nki,nkl,nlj->nij', vecs, dr, vecs)
    q = vals[:, :, None]
    r = vals[:, None, :]
    den = np.log(q)-np.log(r)
    qb = np.broadcast_to(q, den.shape)
    rb = np.broadcast_to(r, den.shape)
    lm = np.empty_like(den)
    close = np.abs(qb-rb) < 1e-7
    np.divide(qb-rb, den, out=lm, where=~close)
    lm[close] = qb[close]
    return float(np.mean(np.sum((np.abs(at)**2)/lm, axis=(1,2))))


def matrix_contraction_check():
    n, d = 128, 3
    x = 2*np.pi*np.arange(n)/n
    rho = np.zeros((n,d,d))
    for i, xx in enumerate(x):
        A = np.array([[1.4+.25*np.cos(3*xx), .14*np.sin(4*xx), .06*np.cos(2*xx)],
                      [.14*np.sin(4*xx), 1.0+.18*np.sin(2*xx), .05*np.cos(5*xx)],
                      [.06*np.cos(2*xx), .05*np.cos(5*xx), .8+.2*np.cos(3*xx)]])
        A += .7*np.eye(d)
        rho[i] = A/np.trace(A)
    i0 = bkm(rho)
    rows=[]
    for t in [0.01, 0.03, 0.06, 0.1]:
        rt=heat(rho,t)
        it=bkm(rt)
        rows.append({"t":t,"ratio":it/i0,"bound":math.exp(-2*t),"I":it})
    return {"I0":i0,"rows":rows,"all_below_bound":all(r["ratio"] <= r["bound"]+0.08 for r in rows)}


def periodic_denoising(regularized, seed=7, steps=500):
    torch.manual_seed(seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    n, d = 64, 3
    x = torch.linspace(0, 2*np.pi, n, device=device)
    clean = torch.stack([torch.sin(x), torch.cos(2*x), torch.sin(3*x)], 1)
    noisy = clean + .35*torch.randn_like(clean)
    model = torch.nn.Sequential(torch.nn.Linear(d,32), torch.nn.Tanh(), torch.nn.Linear(32,d)).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=.025)
    t0=time.time()
    for step in range(steps):
        pred=model(noisy)
        if regularized:
            # Matrix-valued PSD field from three learned channels: Gram matrices.
            h=model[0](noisy)
            z=h[:, :9].reshape(n,3,3)
            S=torch.bmm(z,z.transpose(1,2)) + 1e-3*torch.eye(3,device=device)[None]
            rho=S/(torch.diagonal(S,dim1=1,dim2=2).sum(1)[:,None,None])
            # Exact circular heat step, differentiable through torch FFT.
            freq=torch.fft.fftfreq(n, device=device)*n
            filt=torch.exp(-.035*freq*freq).reshape(n,1,1)
            rh=torch.fft.fft(rho,dim=0)
            rhoh=torch.fft.ifft(rh*filt,dim=0).real
            smooth=torch.diagonal(rhoh,dim1=1,dim2=2).sum(1)
            # retain a scalar route score derived from the smoothed matrix field
            target=rhoh[:, :, 0].real
            loss=((pred-clean)**2).mean() + .15*((target-target.mean(0))**2).mean()
        else:
            loss=((pred-clean)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        mse=((model(noisy)-clean)**2).mean().item()
        residual=(model(noisy)-clean).detach()
        power=torch.abs(torch.fft.fft(residual,dim=0))**2
        split=max(1,n//8)
        hf=power[split:-split].mean().item()
    return {"mse":mse,"high_frequency_residual_power":hf,"seconds":time.time()-t0,"device":device}


def main():
    out={"scalar_mlsi":scalar_mlsi_check(),
         "matrix_contraction":matrix_contraction_check(),
         "toy": {"baseline":periodic_denoising(False), "heat":periodic_denoising(True)},
         "repetitions": [{"seed":s,"baseline":periodic_denoising(False, s),"heat":periodic_denoising(True, s)} for s in [11, 23, 41]]}
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()