import json
import math
from pathlib import Path
import numpy as np
from scipy.special import ndtr, ndtri, logsumexp
from scipy.stats import norm
from sklearn.metrics import roc_auc_score

RNG = np.random.default_rng(1342)


def mixture_params(separation=4.0, prior_var=0.35**2):
    means = np.array([-separation / 2.0, separation / 2.0])
    weights = np.array([0.5, 0.5])
    return means, weights, prior_var


def posterior_components(y, noise_var, separation=4.0, prior_var=0.35**2):
    """p(x0|y) is a two-component Gaussian mixture under y=x0+N(0,noise_var)."""
    means, weights, _ = mixture_params(separation, prior_var)
    y = np.asarray(y)
    marginal_vars = prior_var + noise_var
    logw = np.log(weights)[None, :] + norm.logpdf(y[..., None], means, np.sqrt(marginal_vars))
    logw -= logsumexp(logw, axis=-1, keepdims=True)
    post_w = np.exp(logw)
    gain = prior_var / marginal_vars
    post_means = means + gain * (y[..., None] - means)
    post_var = prior_var * noise_var / marginal_vars
    return post_w, post_means, post_var


def posterior_cdf(x, y, noise_var, separation=4.0, prior_var=0.35**2):
    w, m, v = posterior_components(y, noise_var, separation, prior_var)
    return np.sum(w * ndtr((np.asarray(x)[..., None] - m) / np.sqrt(v)), axis=-1)


def posterior_logpdf(x, y, noise_var, separation=4.0, prior_var=0.35**2):
    w, m, v = posterior_components(y, noise_var, separation, prior_var)
    return logsumexp(np.log(w) + norm.logpdf(np.asarray(x)[..., None], m, np.sqrt(v)), axis=-1)


def transport_sample(y, noise_var, n, separation=4.0, prior_var=0.35**2, grid_n=5001):
    """Exact conditional-CDF transport T(y,z)=F^{-1}_{X|y}(Phi(z))."""
    y = np.asarray(y)
    z = RNG.normal(size=(n, len(y)))
    q = ndtr(z)
    lo = np.min(posterior_components(y, noise_var, separation, prior_var)[1]) - 8
    hi = np.max(posterior_components(y, noise_var, separation, prior_var)[1]) + 8
    # Vectorized inverse CDF by bisection; monotonicity is the key mechanism.
    l = np.full_like(q, lo, dtype=float); r = np.full_like(q, hi, dtype=float)
    yy = np.broadcast_to(y, q.shape[1:])
    for _ in range(45):
        mid = (l + r) / 2
        f = posterior_cdf(mid, yy[None, :], noise_var, separation, prior_var)
        l = np.where(f < q, mid, l); r = np.where(f < q, r, mid)
    return (l + r) / 2, z


def affine_params(y, noise_var, separation=4.0, prior_var=0.35**2):
    """Moment-matched Gaussian reverse kernel, standard affine reverse head."""
    w, m, v = posterior_components(y, noise_var, separation, prior_var)
    mean = np.sum(w * m, axis=-1)
    var = v + np.sum(w * (m - mean[..., None]) ** 2, axis=-1)
    return mean, var


def conditional_kl(y, noise_var, separation=4.0, prior_var=0.35**2):
    """KL(exact posterior || moment-matched Gaussian), quadrature by posterior samples."""
    x, _ = transport_sample(y, noise_var, 1800, separation, prior_var)
    mean, var = affine_params(y, noise_var, separation, prior_var)
    exact = posterior_logpdf(x, y[None, :], noise_var, separation, prior_var)
    gauss = norm.logpdf(x, mean[None, :], np.sqrt(var)[None, :])
    return np.mean(exact - gauss, axis=0)


def residual_auc(noise_var, separation=4.0, prior_var=0.35**2, n=5000):
    """Can a classifier predict y from reverse residual? AUC=0.5 means independence."""
    x, y = sample_joint(n, noise_var, separation, prior_var)
    # Exact transport residual: conditional probability integral transform.
    u = np.clip(posterior_cdf(x, y, noise_var, separation, prior_var), 1e-6, 1-1e-6)
    z_transport = ndtri(u)
    mu, var = affine_params(y, noise_var, separation, prior_var)
    z_affine = (x - mu) / np.sqrt(var)
    # Use |residual|, which captures state-dependent shape without fitting a classifier.
    # AUC is made orientation-invariant.
    def auc(z):
        a = roc_auc_score(y > np.median(y), np.abs(z))
        return max(a, 1-a)
    return auc(z_affine), auc(z_transport)


def sample_joint(n, noise_var, separation=4.0, prior_var=0.35**2):
    means, weights, _ = mixture_params(separation, prior_var)
    c = RNG.choice(2, size=n, p=weights)
    x = RNG.normal(means[c], np.sqrt(prior_var))
    y = x + RNG.normal(0, np.sqrt(noise_var), n)
    return x, y


def math_check():
    # Bayes identity: integrate the joint density relation at random points.
    x, y = sample_joint(3000, 0.5)
    prior = 0.5 * norm.pdf(x, -2, .35) + 0.5 * norm.pdf(x, 2, .35)
    likelihood = norm.pdf(y, x, np.sqrt(.5))
    marginal = .5 * norm.pdf(y, -2, np.sqrt(.35**2+.5)) + .5 * norm.pdf(y, 2, np.sqrt(.35**2+.5))
    post = np.exp(posterior_logpdf(x, y, .5))
    bayes_relerr = np.median(np.abs(post - likelihood*prior/marginal) / (post + 1e-12))
    # PIT uniformity and transport conditional moment agreement.
    u = posterior_cdf(x, y, .5)
    ks_like = float(np.max(np.abs(np.sort(u) - (np.arange(len(u)) + .5) / len(u))))
    return {"bayes_median_relative_error": float(bayes_relerr), "pit_sup_deviation": ks_like}


def run():
    out = {"math_check": math_check(), "separation_sweep": [], "noise_sweep": [], "independence": []}
    # Prediction 1: mixture separation past posterior SD produces increasing affine KL; transport stays ~0.
    for d in [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]:
        ys = np.linspace(-3.5, 3.5, 15)
        kl = conditional_kl(ys, .35**2, d)
        out["separation_sweep"].append({"separation": d, "affine_kl_mean": float(np.mean(kl)), "transport_kl_expected": 0.0})
    # Prediction 2: as forward noise increases, posterior becomes less multimodal and affine gap shrinks.
    for nv in [.03, .10, .25, .5, 1.0, 2.0, 4.0]:
        ys = np.linspace(-4, 4, 17)
        kl = conditional_kl(ys, nv, 4.0)
        out["noise_sweep"].append({"noise_var": nv, "affine_kl_mean": float(np.mean(kl)), "transport_kl_expected": 0.0})
    # Prediction 3: exact PIT transport gives state-independent residual, affine residual does not.
    for nv in [.03, .1, .5, 1.0, 2.0]:
        a, t = residual_auc(nv, 4.0)
        out["independence"].append({"noise_var": nv, "affine_auc": float(a), "transport_auc": float(t), "ideal_auc": .5})

    # Quantitative transition prediction: equal-weight posterior components become
    # visibly separated when their mean gap exceeds two posterior standard deviations.
    # gap/sd = d*sqrt(nv)/(sqrt(prior_var)*sqrt(prior_var+nv)).
    prior = .35**2
    d = 4.0
    # Solve d^2*nv/(prior*(prior+nv)) = 4.
    predicted_nv = (4 * prior**2) / (d**2 - 4 * prior)
    out["bimodality_transition"] = {"predicted_noise_var_gap_eq_2sd": predicted_nv, "sweep": []}
    for nv in [0.01, 0.03, 0.05, predicted_nv, 0.10, 0.25, 1.0]:
        w, m, v = posterior_components(np.array([0.0]), nv, d, prior)
        ratio = float((m[0, 1] - m[0, 0]) / np.sqrt(v))
        out["bimodality_transition"]["sweep"].append({"noise_var": float(nv), "posterior_gap_over_sd": ratio})

    # Marginal sampling check: transport samples should match the mixture; affine
    # samples have the same first two moments but lose mixture shape.
    xtrue, ytest = sample_joint(6000, .5, 4.0, prior)
    xt, _ = transport_sample(ytest[:600], .5, 10, 4.0, prior)
    am, av = affine_params(ytest[:600], .5, 4.0, prior)
    xa = am[None, :] + np.sqrt(av)[None, :] * RNG.normal(size=xt.shape)
    out["sample_moments"] = {
        "true_mean": float(np.mean(xtrue)), "transport_mean": float(np.mean(xt)), "affine_mean": float(np.mean(xa)),
        "true_second_moment": float(np.mean(xtrue**2)), "transport_second_moment": float(np.mean(xt**2)), "affine_second_moment": float(np.mean(xa**2)),
        "true_excess_kurtosis": float(np.mean((xtrue-np.mean(xtrue))**4)/np.var(xtrue)**2-3),
        "transport_excess_kurtosis": float(np.mean((xt-np.mean(xt))**4)/np.var(xt)**2-3),
        "affine_excess_kurtosis": float(np.mean((xa-np.mean(xa))**4)/np.var(xa)**2-3)}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    run()