"""Return distributions for AI in Financial Risk, part 1.

Python 3.10+, pandas and numpy. Every step writes its result to a CSV in
--out so each intermediate can be opened and checked on its own.

    python distributions.py                      # download SP500 from FRED
    python distributions.py --prices sp500.csv   # use a saved FRED CSV
    python distributions.py --check              # run the self-checks

FRED keeps a rolling ten-year window, so pass --start and --end to pin a
run. The post used --start 2016-10-31 --end 2026-10-29.

FRED's SP500 series is the S&P 500 price index: daily closes, no
dividends, about the last ten years only. Holidays appear as blank rows.
"""

from __future__ import annotations

import argparse
import math
from pathlib import Path
from statistics import NormalDist
import unittest
import urllib.request

import numpy as np
import pandas as pd

FRED_URL = "https://fred.stlouisfed.org/graph/fredgraph.csv?id=SP500"
SIGMA_LEVELS = (2, 3, 4, 5)
ACF_LAGS = (1, 2, 5, 10, 21)
ROLLING_DAYS = 21
TRADING_DAYS = 252
SEED = 7


# Step 1: load prices -------------------------------------------------------

def download_prices(path: Path) -> None:
    with urllib.request.urlopen(FRED_URL, timeout=60) as response:
        path.write_bytes(response.read())


def load_prices(path: Path, start: str | None, end: str | None) -> pd.DataFrame:
    raw = pd.read_csv(path)
    raw.columns = ["date", "close"]
    raw["date"] = pd.to_datetime(raw["date"])
    raw["close"] = pd.to_numeric(raw["close"], errors="coerce")
    if start:
        raw = raw[raw["date"] >= pd.Timestamp(start)]
    if end:
        raw = raw[raw["date"] <= pd.Timestamp(end)]
    blank_rows = int(raw["close"].isna().sum())
    prices = raw.dropna(subset=["close"]).reset_index(drop=True)
    if not prices["date"].is_monotonic_increasing or prices["date"].duplicated().any():
        raise ValueError("dates must be unique and increasing")
    if (prices["close"] <= 0).any():
        raise ValueError("closing prices must be positive")
    print(f"loaded {len(prices)} closes, dropped {blank_rows} blank rows")
    return prices


# Step 2: prices to returns -------------------------------------------------

def add_returns(prices: pd.DataFrame) -> pd.DataFrame:
    returns = prices.copy()
    returns["prev_close"] = returns["close"].shift(1)
    returns["simple_return"] = returns["close"] / returns["prev_close"] - 1
    returns["log_return"] = np.log(returns["close"] / returns["prev_close"])
    returns["calendar_gap_days"] = returns["date"].diff().dt.days
    return returns.dropna(subset=["prev_close"]).reset_index(drop=True)


# Step 3: summary moments ---------------------------------------------------

def moments(r: pd.Series) -> dict[str, float]:
    centred = r - r.mean()
    std = r.std(ddof=1)
    return {
        "days": len(r),
        "mean": r.mean(),
        "std": std,
        "annualised_vol": std * math.sqrt(TRADING_DAYS),
        "skew": (centred**3).mean() / std**3,
        "excess_kurtosis": (centred**4).mean() / std**4 - 3,
        "min": r.min(),
        "max": r.max(),
    }


def moments_table(observed: pd.Series, simulated: pd.Series) -> pd.DataFrame:
    return pd.DataFrame(
        {"observed": moments(observed), "normal_same_mean_std": moments(simulated)}
    )


# Step 4: tail counts against the normal ------------------------------------

def normal_two_sided_tail(k: float) -> float:
    return math.erfc(k / math.sqrt(2))


def tail_counts(observed: pd.Series, simulated: pd.Series) -> pd.DataFrame:
    mean, std = observed.mean(), observed.std(ddof=1)
    z_observed = (observed - mean) / std
    z_simulated = (simulated - mean) / std
    rows = []
    for k in SIGMA_LEVELS:
        probability = normal_two_sided_tail(k)
        rows.append(
            {
                "beyond_sigma": k,
                "normal_probability": probability,
                "normal_expected_days": probability * len(observed),
                "simulated_normal_days": int((z_simulated.abs() > k).sum()),
                "observed_days": int((z_observed.abs() > k).sum()),
                "observed_down_days": int((z_observed < -k).sum()),
                "observed_up_days": int((z_observed > k).sum()),
            }
        )
    return pd.DataFrame(rows)


def worst_days(returns: pd.DataFrame, n: int = 10) -> pd.DataFrame:
    mean, std = returns["log_return"].mean(), returns["log_return"].std(ddof=1)
    worst = returns.nsmallest(n, "log_return")[["date", "close", "log_return"]].copy()
    worst["sigmas"] = (worst["log_return"] - mean) / std
    # erfc keeps precision far into the tail, where NormalDist().cdf rounds to 0.
    lower_tail = worst["sigmas"].map(lambda z: 0.5 * math.erfc(-z / math.sqrt(2)))
    worst["normal_one_in_n_years"] = 1 / lower_tail / TRADING_DAYS
    return worst.reset_index(drop=True)


# Step 5: one-day 99% quantile, normal versus empirical ---------------------

def quantile_comparison(observed: pd.Series) -> pd.DataFrame:
    mean, std = observed.mean(), observed.std(ddof=1)
    rows = []
    for level in (0.95, 0.99, 0.999):
        tail = 1 - level
        normal_cut = mean + std * NormalDist().inv_cdf(tail)
        empirical_cut = observed.quantile(tail)
        rows.append(
            {
                "level": level,
                "normal_cutoff": normal_cut,
                "empirical_cutoff": empirical_cut,
                "expected_days_below_normal_cut": tail * len(observed),
                "observed_days_below_normal_cut": int((observed < normal_cut).sum()),
            }
        )
    return pd.DataFrame(rows)


# Step 6: volatility clustering ---------------------------------------------

def autocorrelation(r: pd.Series, lag: int) -> float:
    centred = r - r.mean()
    return float((centred * centred.shift(lag)).sum() / (centred**2).sum())


def clustering_table(observed: pd.Series, simulated: pd.Series) -> pd.DataFrame:
    rows = []
    for lag in ACF_LAGS:
        rows.append(
            {
                "lag_days": lag,
                "acf_returns": autocorrelation(observed, lag),
                "acf_abs_returns": autocorrelation(observed.abs(), lag),
                "acf_abs_simulated_normal": autocorrelation(simulated.abs(), lag),
            }
        )
    rows_frame = pd.DataFrame(rows)
    rows_frame["noise_band_95"] = 1.96 / math.sqrt(len(observed))
    return rows_frame


def rolling_volatility(returns: pd.DataFrame) -> pd.DataFrame:
    vol = returns[["date", "log_return"]].copy()
    vol["rolling_vol_annualised"] = (
        vol["log_return"].rolling(ROLLING_DAYS).std(ddof=1) * math.sqrt(TRADING_DAYS)
    )
    return vol.dropna().reset_index(drop=True)


def volatility_by_year(returns: pd.DataFrame) -> pd.DataFrame:
    by_year = returns.groupby(returns["date"].dt.year)["log_return"]
    table = pd.DataFrame(
        {
            "days": by_year.size(),
            "annualised_vol": by_year.std(ddof=1) * math.sqrt(TRADING_DAYS),
            "worst_day": by_year.min(),
        }
    )
    table.index.name = "year"
    return table.reset_index()


# Step 7: does one year drive the result? ----------------------------------

def without_year(returns: pd.DataFrame, year: int) -> pd.DataFrame:
    rows = []
    for label, sample in (
        ("all_days", returns),
        (f"without_{year}", returns[returns["date"].dt.year != year]),
    ):
        r = sample["log_return"]
        z = (r - r.mean()) / r.std(ddof=1)
        summary = moments(r)
        rows.append(
            {
                "sample": label,
                "days": summary["days"],
                "skew": summary["skew"],
                "excess_kurtosis": summary["excess_kurtosis"],
                "observed_beyond_3_sigma": int((z.abs() > 3).sum()),
                "normal_expected_beyond_3_sigma": normal_two_sided_tail(3) * len(r),
                "acf_returns_lag1": autocorrelation(r, 1),
                "acf_abs_returns_lag1": autocorrelation(r.abs(), 1),
                "acf_abs_returns_lag10": autocorrelation(r.abs(), 10),
            }
        )
    return pd.DataFrame(rows)


# Control: a normal series with the same mean, std and length ---------------

def simulate_normal(observed: pd.Series, seed: int = SEED) -> pd.Series:
    rng = np.random.default_rng(seed)
    draws = rng.normal(observed.mean(), observed.std(ddof=1), len(observed))
    return pd.Series(draws, name="simulated_normal")


# Run ----------------------------------------------------------------------

def save(frame: pd.DataFrame, out: Path, name: str, index: bool = False) -> None:
    frame.to_csv(out / name, index=index)
    print(f"\n== {name}\n{frame.to_string()}")


def run(
    prices_path: Path | None, out: Path, start: str | None, end: str | None, exclude_year: int
) -> None:
    out.mkdir(parents=True, exist_ok=True)
    if prices_path is None:
        prices_path = out / "00_fred_sp500_raw.csv"
        download_prices(prices_path)

    prices = load_prices(prices_path, start, end)
    prices.to_csv(out / "01_prices.csv", index=False)

    returns = add_returns(prices)
    returns.to_csv(out / "02_returns.csv", index=False)
    print(f"returns from {returns['date'].min():%Y-%m-%d} to {returns['date'].max():%Y-%m-%d}")

    observed = returns["log_return"]
    simulated = simulate_normal(observed)

    pd.options.display.float_format = "{:.6g}".format
    save(moments_table(observed, simulated), out, "03_moments.csv", index=True)
    save(tail_counts(observed, simulated), out, "04_tail_counts.csv")
    save(worst_days(returns), out, "04b_worst_days.csv")
    save(quantile_comparison(observed), out, "05_quantiles.csv")
    save(clustering_table(observed, simulated), out, "06_clustering.csv")
    rolling_volatility(returns).to_csv(out / "06b_rolling_vol.csv", index=False)
    save(volatility_by_year(returns), out, "06c_vol_by_year.csv")
    save(without_year(returns, exclude_year), out, "07_without_year.csv")


# Self-checks --------------------------------------------------------------

class Checks(unittest.TestCase):
    def test_normal_tail_matches_textbook(self):
        self.assertAlmostEqual(normal_two_sided_tail(3), 0.0026998, places=6)

    def test_simple_and_log_returns(self):
        prices = pd.DataFrame(
            {"date": pd.to_datetime(["2024-01-01", "2024-01-02", "2024-01-03"]),
             "close": [100.0, 110.0, 99.0]}
        )
        returns = add_returns(prices)
        self.assertAlmostEqual(returns["simple_return"][0], 0.10)
        self.assertAlmostEqual(returns["simple_return"][1], -0.10)
        # Log returns add up to the whole-period log return; simple ones do not.
        self.assertAlmostEqual(returns["log_return"].sum(), math.log(99 / 100))

    def test_normal_control_has_no_fat_tails_or_clustering(self):
        draws = simulate_normal(pd.Series(np.random.default_rng(1).normal(0, 0.01, 50_000)))
        self.assertLess(abs(moments(draws)["excess_kurtosis"]), 0.1)
        self.assertLess(abs(autocorrelation(draws.abs(), 1)), 0.02)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--prices", type=Path, help="saved FRED SP500 CSV")
    parser.add_argument("--start", help="first date to keep, YYYY-MM-DD")
    parser.add_argument("--end", help="last date to keep, YYYY-MM-DD")
    parser.add_argument("--exclude-year", type=int, default=2020)
    parser.add_argument("--out", type=Path, default=Path("part01_output"))
    parser.add_argument("--check", action="store_true", help="run self-checks")
    args = parser.parse_args()
    if args.check:
        unittest.main(argv=["distributions"], verbosity=2)
    else:
        run(args.prices, args.out, args.start, args.end, args.exclude_year)


if __name__ == "__main__":
    main()