"""Compare candidate return distributions on 1871-2024 annual US stock returns.

Fits each family to the simulator's own annual S&P composite log total returns
by maximum likelihood and reports log-likelihood, AIC and BIC, so the choice of
Tukey g-and-h for the simulator (simulation_engine/tail_shapes.py) can be
checked against the obvious alternatives. NumPy only; the published copy runs
on the published CSV.

Usage:
    python -m research.tail_family_comparison --out frontend/public/blog-assets/tail-family-comparison.json
    python tail_family_comparison.py --returns-dir . --out out.json   # published copy, next to the CSVs
"""

from __future__ import annotations

import argparse
import json
from math import erf, lgamma, log, pi, sqrt
from pathlib import Path
from typing import Any, Callable

import numpy as np

DATA_DIR = Path(__file__).resolve().parent.parent / "simulation_engine/historical/data"
# Stocks over the full record; bonds from 1926, when measured returns start
# (the earlier series is a yield-duration approximation; tail_calibration.py).
SERIES = {"stocks": ("equity.csv", 1871), "bonds": ("bonds.csv", 1926)}
STRETCH_START = -1.0


def load(path: Path, first: int = 1871, last: int = 2024) -> np.ndarray:
    rows = (line for line in path.read_text().splitlines() if line[:1].isdigit())
    series = {int(y): float(r) for y, r in (row.split(",")[:2] for row in rows)}
    return np.asarray(np.log1p([series[y] for y in range(first, last + 1)]), dtype=float)


_erf = np.vectorize(erf)


def normal_logpdf(z: np.ndarray) -> np.ndarray:
    return np.asarray(-0.5 * z * z - 0.5 * log(2 * pi), dtype=float)


def normal_cdf(z: np.ndarray) -> np.ndarray:
    return np.asarray(0.5 * (1.0 + _erf(z / sqrt(2.0))), dtype=float)


# Each family: (parameter count, starting point, negative log-likelihood).
def family_normal(x: np.ndarray) -> float:
    return float(-np.sum(normal_logpdf((x - x.mean()) / x.std()) - np.log(x.std())))


def nll_student_t(p: np.ndarray, x: np.ndarray) -> float:
    mu, log_s, log_nu = p
    s, nu = np.exp(log_s), np.exp(log_nu) + 2.0
    z = (x - mu) / s
    const = lgamma((nu + 1) / 2) - lgamma(nu / 2) - 0.5 * log(nu * pi)
    return float(-np.sum(const - np.log(s) - (nu + 1) / 2 * np.log1p(z * z / nu)))


def nll_skew_normal(p: np.ndarray, x: np.ndarray) -> float:
    xi, log_w, alpha = p
    w = np.exp(log_w)
    z = (x - xi) / w
    density = 2.0 / w * np.exp(normal_logpdf(z)) * normal_cdf(alpha * z)
    return float(-np.sum(np.log(np.maximum(density, 1e-300))))


def nll_stretch(p: np.ndarray, x: np.ndarray) -> float:
    """The simulator's former transform: draws below -1 stretched by k."""
    a, log_s, log_k = p
    s, k = np.exp(log_s), 1.0 + np.exp(log_k)
    y = (x - a) / s
    z = np.where(y < STRETCH_START, STRETCH_START + (y - STRETCH_START) / k, y)
    jac = np.where(y < STRETCH_START, k, 1.0)
    return float(-np.sum(normal_logpdf(z) - np.log(jac) - np.log(s)))


def _gh(z: np.ndarray, g: float, h: float) -> np.ndarray:
    core = z if abs(g) < 1e-9 else np.expm1(g * z) / g
    return np.asarray(core * np.exp(h * z * z / 2.0), dtype=float)


def nll_g_and_h(p: np.ndarray, x: np.ndarray) -> float:
    a, log_b, g, h = p
    if h < 0 or h > 0.24:
        return 1e12
    b = np.exp(log_b)
    target = (x - a) / b
    lo, hi = np.full_like(target, -12.0), np.full_like(target, 12.0)
    for _ in range(80):
        mid = (lo + hi) / 2.0
        above = _gh(mid, g, h) > target
        hi, lo = np.where(above, mid, hi), np.where(above, lo, mid)
    z = (lo + hi) / 2.0
    tail = np.exp(h * z * z / 2.0)
    d = tail * (1.0 + h * z * z) if abs(g) < 1e-9 else tail * (np.exp(g * z) + h * z * np.expm1(g * z) / g)
    return float(-np.sum(normal_logpdf(z) - np.log(b * d)))


def minimize(f: Callable[[np.ndarray], float], starts: list[list[float]]) -> tuple[np.ndarray, float]:
    best: tuple[np.ndarray, float] | None = None
    for start in starts:
        simplex = [np.array(start, dtype=float)]
        simplex += [simplex[0] + np.eye(len(start))[i] * 0.1 for i in range(len(start))]
        values = [f(v) for v in simplex]
        for _ in range(4000):
            order = np.argsort(values)
            simplex, values = [simplex[i] for i in order], [values[i] for i in order]
            if abs(values[-1] - values[0]) < 1e-11:
                break
            centroid = np.mean(simplex[:-1], axis=0)
            reflected = centroid + (centroid - simplex[-1])
            fr = f(reflected)
            if fr < values[0]:
                expanded = centroid + 2 * (centroid - simplex[-1])
                fe = f(expanded)
                simplex[-1], values[-1] = (expanded, fe) if fe < fr else (reflected, fr)
            elif fr < values[-2]:
                simplex[-1], values[-1] = reflected, fr
            else:
                contracted = centroid + 0.5 * (simplex[-1] - centroid)
                fc = f(contracted)
                if fc < values[-1]:
                    simplex[-1], values[-1] = contracted, fc
                else:
                    simplex = [simplex[0] + 0.5 * (v - simplex[0]) for v in simplex]
                    values = [f(v) for v in simplex]
        i = int(np.argmin(values))
        if best is None or values[i] < best[1]:
            best = (simplex[i], values[i])
    assert best is not None
    return best


def compare(x: np.ndarray) -> dict[str, Any]:
    n = len(x)
    m, s = float(x.mean()), float(x.std())
    fits = {
        "normal (log-normal returns)": (2, family_normal(x)),
        "Student t": (3, minimize(lambda p: nll_student_t(p, x), [[m, log(s), log(5.0)], [m, log(s), log(20.0)]])[1]),
        "skew-normal (Azzalini)": (3, minimize(lambda p: nll_skew_normal(p, x), [[m + s, log(s), -1.0], [m, log(s), -3.0]])[1]),
        "left-tail stretch (former model)": (3, minimize(lambda p: nll_stretch(p, x), [[m, log(s), log(0.5)], [m, log(s * 0.9), log(1.0)]])[1]),
        "Tukey g-and-h (current model)": (4, minimize(lambda p: nll_g_and_h(p, x), [[float(np.median(x)), log(s), g0, 0.05] for g0 in (-0.4, -0.1, 0.1)])[1]),
    }
    rows: list[dict[str, Any]] = []
    for name, (k, nll) in fits.items():
        ll = -nll
        rows.append({"family": name, "parameters": k, "log_likelihood": round(ll, 2),
                     "aic": round(2 * k - 2 * ll, 2), "bic": round(k * log(n) - 2 * ll, 2)})
    rows.sort(key=lambda r: r["aic"])
    z = (x - x.mean()) / x.std()
    return {"n": n, "skewness": round(float(np.mean(z**3)), 3), "excess_kurtosis": round(float(np.mean(z**4) - 3), 3),
            "families": rows}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--returns-dir", type=Path, default=DATA_DIR,
                        help="directory holding equity.csv and bonds.csv (year,return)")
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    result: dict[str, Any] = {}
    for name, (filename, first) in SERIES.items():
        result[name] = {"years": [first, 2024], **compare(load(args.returns_dir / filename, first))}
        print(name)
        for row in result[name]["families"]:
            print(f"  {row['family']:36s} k={row['parameters']} logL={row['log_likelihood']:7.2f} AIC={row['aic']:8.2f} BIC={row['bic']:8.2f}")
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, indent=1) + "\n")


if __name__ == "__main__":
    main()