import json
import numpy as np
from scipy.optimize import nnls
from scipy.stats import wasserstein_distance

SEED = 1457
rng = np.random.default_rng(SEED)

def T(x):
    return 0.78*x + 0.28*np.sin(2.7*x)

class RBFCodec:
    def __init__(self, lo=-2.0, hi=2.0, m=25, sigma=.22, ngrid=1601):
        self.x = np.linspace(lo, hi, ngrid)
        self.dx = self.x[1] - self.x[0]
        self.centers = np.linspace(lo, hi, m)
        raw = np.exp(-.5*((self.x[:, None]-self.centers[None, :])/sigma)**2)
        # Normalize each basis to unit mass: q_i = integral psi_i = 1 exactly up to grid precision.
        self.A = raw / (raw.sum(axis=0)[None, :] * self.dx)
        self.q = np.trapz(self.A, self.x, axis=0)

    def fit(self, samples):
        edges = np.linspace(self.x[0], self.x[-1], len(self.x))
        h, _ = np.histogram(samples, bins=edges, density=False)
        dens = np.empty_like(self.x)
        dens[1:] = h / (len(samples) * self.dx)
        dens[0] = dens[1]
        c = nnls(self.A, dens)[0]
        return c / max(self.q @ c, 1e-12)

    def density(self, c):
        p = np.maximum(self.A @ c, 0.)
        return p / max(np.trapz(p, self.x), 1e-12)

    def sample(self, c, n):
        p = self.density(c)
        cdf = np.maximum.accumulate(np.cumsum(p) * self.dx)
        cdf /= cdf[-1]
        return np.interp(rng.random(n), np.r_[0., cdf[:-1]], self.x)

    def moments(self, c):
        p = self.density(c)
        mu = np.trapz(self.x*p, self.x)
        va = np.trapz((self.x-mu)**2*p, self.x)
        return mu, va

def mass_project(K, q):
    residual = q - q @ K
    return K + np.outer(q / (q @ q), residual)

def train_operator(codec, ntrain=260, npart=1800, ridge=1e-5):
    C, Y = [], []
    for _ in range(ntrain):
        w = rng.dirichlet(np.ones(len(codec.centers)) * .7)
        C.append(w)
        Y.append(codec.fit(T(codec.sample(w, npart))))
    C, Y = np.asarray(C).T, np.asarray(Y).T
    Kraw = Y @ C.T @ np.linalg.inv(C @ C.T + ridge*np.eye(C.shape[0]))
    K = mass_project(Kraw, codec.q)
    return K, Kraw

def learned_roll(codec, K, c0, horizons=(1, 5, 10, 25), ntruth=12000):
    truth = codec.sample(c0, ntruth)
    c = c0.copy(); out = []
    for k in range(1, max(horizons)+1):
        truth = T(truth); c = K @ c
        if k in horizons:
            out.append(wasserstein_distance(truth, codec.sample(c, ntruth)))
    return out

def gaussian_roll(c0, codec, horizons=(1, 5, 10, 25), ntruth=12000):
    truth = codec.sample(c0, ntruth)
    mu, var = codec.moments(c0); out = []
    for k in range(1, max(horizons)+1):
        truth = T(truth)
        s = np.sqrt(max(var, 1e-10)); z = np.array([mu-s, mu, mu+s]); wt = np.array([.25, .5, .25])
        mz = np.sum(wt*T(z)); var = np.sum(wt*(T(z)-mz)**2); mu = mz
        if k in horizons:
            out.append(wasserstein_distance(truth, rng.normal(mu, np.sqrt(max(var, 1e-10)), ntruth)))
    return out

def spectral_sweep():
    # Prediction: for M=gamma*A, transition occurs at gamma*rho(A)=1 and
    # asymptotic log norm slope is log(gamma*rho(A)).
    lam = .93; A = np.diag([lam] + [.45]*7); c = np.ones(8); rows=[]
    for gamma in [.70, .90, 1.00, 1.075, 1.10, 1.30]:
        M = gamma*A; z = c.copy(); norms = []
        for _ in range(100): norms.append(np.linalg.norm(z)); z = M @ z
        slope = np.polyfit(np.arange(50, 100), np.log(np.maximum(norms[50:], 1e-300)), 1)[0]
        rows.append({'gamma': gamma, 'predicted_rho': gamma*lam,
                     'observed_log_slope': float(slope), 'predicted_log_slope': float(np.log(gamma*lam)),
                     'observed_bounded': bool(norms[-1] <= norms[0])})
    return rows

def separation_sweep(codec, K):
    # Prediction: a multimodal PF representation should retain an advantage as
    # mode separation grows, while a single Gaussian loses shape information.
    rows = []
    for left, right in [(-.25, .25), (-.55, .55), (-.85, .85), (-1.15, 1.15)]:
        c0 = np.exp(-.5*((codec.centers-left)/.13)**2) + np.exp(-.5*((codec.centers-right)/.13)**2)
        c0 /= codec.q @ c0
        pf = learned_roll(codec, K, c0, horizons=(5,), ntruth=10000)[0]
        gauss = gaussian_roll(c0, codec, horizons=(5,), ntruth=10000)[0]
        rows.append({'separation': right-left, 'PF_W1': pf, 'Gaussian_W1': gauss, 'PF_minus_Gaussian': pf-gauss})
    return rows

def main():
    codec = RBFCodec(); K, Kraw = train_operator(codec)
    q = codec.q
    raw_mass_error = float(np.max(np.abs(q @ Kraw - q)))
    projected_mass_error = float(np.max(np.abs(q @ K - q)))
    c0 = np.zeros(len(q)); c0[5] = .5; c0[19] = .5
    horizons = (1, 5, 10, 25)
    pf = learned_roll(codec, K, c0, horizons)
    gauss = gaussian_roll(c0, codec, horizons)
    eig = np.linalg.eigvals(K)
    result = {
      'seed': SEED, 'basis': len(q), 'basis_integral_min_max': [float(q.min()), float(q.max())],
      'mass_error_raw': raw_mass_error, 'mass_error_after_projection': projected_mass_error,
      'spectral_radius_fitted_K': float(max(abs(eig))), 'spectral_prediction_sweep': spectral_sweep(),
      'separation_sweep': separation_sweep(codec, K),
      'w1_horizons': {'horizons': list(horizons), 'PF_RBF': pf, 'Gaussian_moment': gauss},
      'mean_w1_PF': float(np.mean(pf)), 'mean_w1_Gaussian': float(np.mean(gauss))}
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()