import json
from pathlib import Path
import numpy as np
from scipy import stats, optimize

RNG = np.random.default_rng(1434)
EPS = 1e-10


def copula_logpdf(u1, u2, rho):
    """log h1(u1) + log h2(u2|u1) for a Gaussian copula."""
    z1, z2 = stats.norm.ppf(np.clip(u1, EPS, 1-EPS)), stats.norm.ppf(np.clip(u2, EPS, 1-EPS))
    rho = float(np.clip(rho, -0.999999, 0.999999))
    return (-0.5*np.log1p(-rho*rho)
            - (rho*rho*(z1*z1 + z2*z2) - 2*rho*z1*z2)/(2*(1-rho*rho)))


def conditional_cdf(u2, u1, rho):
    z1, z2 = stats.norm.ppf(np.clip(u1, EPS, 1-EPS)), stats.norm.ppf(np.clip(u2, EPS, 1-EPS))
    return stats.norm.cdf((z2-rho*z1)/np.sqrt(1-rho*rho))


def make_sample(n, rho, marginal_pair=("lognormal", "gamma"), seed=0):
    rng = np.random.default_rng(seed)
    z = rng.normal(size=(n, 2))
    z[:, 1] = rho*z[:, 0] + np.sqrt(1-rho*rho)*z[:, 1]
    u = stats.norm.cdf(z)
    if marginal_pair[0] == "lognormal":
        x1 = stats.lognorm.ppf(u[:, 0], s=0.8, scale=np.exp(0.3))
        f1 = stats.lognorm.pdf(x1, s=0.8, scale=np.exp(0.3))
    else:
        x1 = stats.gamma.ppf(u[:, 0], a=2.0, scale=1.2)
        f1 = stats.gamma.pdf(x1, a=2.0, scale=1.2)
    if marginal_pair[1] == "gamma":
        x2 = stats.gamma.ppf(u[:, 1], a=3.0, scale=0.7)
        f2 = stats.gamma.pdf(x2, a=3.0, scale=0.7)
    else:
        x2 = stats.lognorm.ppf(u[:, 1], s=0.5, scale=np.exp(0.2))
        f2 = stats.lognorm.pdf(x2, s=0.5, scale=np.exp(0.2))
    return x1, x2, u, f1, f2


def fit_rho(u):
    def objective(a):
        return -np.mean(copula_logpdf(u[:, 0], u[:, 1], np.tanh(a)))
    result = optimize.minimize_scalar(objective, bounds=(-4, 4), method="bounded")
    return float(np.tanh(result.x))


def pit_sweep():
    rows = []
    settings = [("lognormal", "gamma"), ("gamma", "lognormal"), ("lognormal", "lognormal"), ("gamma", "gamma")]
    for j, pair in enumerate(settings):
        _, _, u, _, _ = make_sample(30000, 0.0, pair, 100+j)
        for k in range(2):
            rows.append({"marginals": "%s/%s" % pair, "coordinate": k+1,
                         "mean": float(np.mean(u[:, k])), "variance": float(np.var(u[:, k])),
                         "ks": float(stats.kstest(u[:, k], "uniform").statistic)})
    return rows


def dependence_sweep():
    rows = []
    for i, rho in enumerate([0.0, 0.2, 0.5, 0.8]):
        _, _, u, _, _ = make_sample(40000, rho, ("lognormal", "gamma"), 200+i)
        observed = float(stats.spearmanr(u[:, 0], u[:, 1]).statistic)
        predicted = float(6/np.pi*np.arcsin(rho/2))
        kendall_observed = float(stats.kendalltau(u[:, 0], u[:, 1]).statistic)
        kendall_predicted = float(2/np.pi*np.arcsin(rho))
        rows.append({"rho": rho, "predicted_spearman": predicted, "observed_spearman": observed,
                     "abs_error": abs(observed-predicted),
                     "predicted_kendall": kendall_predicted,
                     "observed_kendall": kendall_observed,
                     "kendall_abs_error": abs(kendall_observed-kendall_predicted)})
    return rows


def likelihood_sweep():
    rows = []
    for i, rho in enumerate([0.0, 0.2, 0.5, 0.8]):
        _, _, u, f1, f2 = make_sample(50000, rho, ("lognormal", "gamma"), 300+i)
        marginal_nll = float(np.mean(-np.log(np.maximum(f1, EPS))-np.log(np.maximum(f2, EPS))))
        indep = marginal_nll
        joint = float(np.mean(-np.log(np.maximum(f1, EPS))-np.log(np.maximum(f2, EPS))
                              -copula_logpdf(u[:, 0], u[:, 1], rho)))
        observed_gain = indep-joint
        predicted_gain = float(-0.5*np.log(1-rho*rho))
        fitted = fit_rho(u)
        rows.append({"rho": rho, "independent_nll": indep, "copula_nll": joint,
                     "observed_nll_gain": observed_gain, "predicted_nll_gain": predicted_gain,
                     "fitted_rho": fitted})
    return rows


def derivative_check():
    u1, u2, rho = 0.37, 0.61, 0.65
    h = 1e-5
    numerical = (conditional_cdf(u2+h, u1, rho)-conditional_cdf(u2-h, u1, rho))/(2*h)
    z1, z2 = stats.norm.ppf(u1), stats.norm.ppf(u2)
    analytic = np.exp(copula_logpdf(np.array([u1]), np.array([u2]), rho))[0]
    return {"u1": u1, "u2": u2, "rho": rho, "finite_difference": float(numerical),
            "analytic_density": float(analytic), "relative_error": float(abs(numerical-analytic)/analytic)}


def main():
    pit = pit_sweep()
    dep = dependence_sweep()
    lik = likelihood_sweep()
    deriv = derivative_check()
    # A direct baseline-vs-idea mini comparison at substantial dependence.
    target = lik[-1]
    summary = {
        "predictions": {
            "PIT_uniformity": "mean -> 0.5, variance -> 1/12, and KS remains small for every continuous marginal",
            "rank_invariance": "For a Gaussian copula, Spearman = (6/pi) asin(rho/2) and Kendall tau = (2/pi) asin(rho), independent of marginal family",
            "dependence_gain": "independent NLL - copula NLL = -0.5 log(1-rho^2), exactly zero at rho=0"
        },
        "derivative_check": deriv,
        "pit_sweep": pit,
        "dependence_sweep": dep,
        "likelihood_sweep": lik,
        "mini_comparison": {"dependence": target["rho"], "baseline_independent_nll": target["independent_nll"],
                            "idea_copula_nll": target["copula_nll"], "idea_fitted_rho": target["fitted_rho"]}
    }
    Path("results.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))

if __name__ == "__main__":
    main()