"""Inference versus prediction for AI in Financial Risk, part 2.

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 prediction_vs_inference.py                      # download SP500 from FRED
    python prediction_vs_inference.py --prices sp500.csv   # use a saved FRED CSV
    python prediction_vs_inference.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-11-14 --end 2026-11-13.

The question asked of every row: given what is known at today's close,
will tomorrow's log return be below LOSS_THRESHOLD?
"""

from __future__ import annotations

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

import numpy as np
import pandas as pd

FRED_URL = "https://fred.stlouisfed.org/graph/fredgraph.csv?id=SP500"
LOSS_THRESHOLD = -0.01          # a "large loss" day: log return below -1%
FEATURES = ["return_1d", "abs_return_1d", "mean_abs_return_5d", "mean_abs_return_21d"]
TRAIN_SHARE, VALIDATION_SHARE = 0.60, 0.20
K_VALUES = (1, 5, 25, 100, 400)
BLOCK_DAYS = 21                 # block length for the bootstrap standard errors
BOOTSTRAP_RUNS = 500
SIMULATED_TRAIN_SIZES = (1_500, 15_000)
SIMULATED_RUNS = 100
SIMULATED_EVAL_POINTS = 500
COST_FALSE_ALARM = 1.0
COST_RATIOS = (5, 10, 20)       # a missed large loss costs this many false alarms
REPORT_COST_RATIO = 10
SEED = 7


# Step 1: load prices (same as part 1) -------------------------------------

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["log_return"] = np.log(returns["close"] / returns["prev_close"])
    return returns.dropna(subset=["prev_close"]).reset_index(drop=True)


# Step 3: features known today, label known tomorrow ------------------------

def add_features(returns: pd.DataFrame) -> pd.DataFrame:
    data = returns[["date", "log_return"]].copy()
    data["return_1d"] = data["log_return"]
    data["abs_return_1d"] = data["log_return"].abs()
    # Rolling windows end on today's row, so they use today and earlier only.
    data["mean_abs_return_5d"] = data["abs_return_1d"].rolling(5).mean()
    data["mean_abs_return_21d"] = data["abs_return_1d"].rolling(21).mean()
    return data


def add_label(data: pd.DataFrame) -> pd.DataFrame:
    labelled = data.copy()
    labelled["next_date"] = labelled["date"].shift(-1)
    labelled["next_log_return"] = labelled["log_return"].shift(-1)
    labelled["large_loss_next_day"] = (labelled["next_log_return"] < LOSS_THRESHOLD).astype(int)
    # The first 20 rows have no 21-day window; the last row has no tomorrow.
    complete = labelled.dropna(subset=FEATURES + ["next_log_return"])
    return complete.reset_index(drop=True)


# Step 4: split by time -----------------------------------------------------

def add_split(dataset: pd.DataFrame) -> pd.DataFrame:
    split = dataset.copy()
    train_end = int(len(split) * TRAIN_SHARE)
    validation_end = int(len(split) * (TRAIN_SHARE + VALIDATION_SHARE))
    split["split"] = "test"
    split.loc[: validation_end - 1, "split"] = "validation"
    split.loc[: train_end - 1, "split"] = "train"
    return split


def split_summary(dataset: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for name in ("train", "validation", "test"):
        part = dataset[dataset["split"] == name]
        rows.append(
            {
                "split": name,
                "first_date": part["date"].min().date(),
                "last_date": part["date"].max().date(),
                "days": len(part),
                "large_loss_days": int(part["large_loss_next_day"].sum()),
                "base_rate": part["large_loss_next_day"].mean(),
            }
        )
    return pd.DataFrame(rows)


def standardise(train_x: np.ndarray, other_x: np.ndarray) -> np.ndarray:
    # Mean and standard deviation come from the training rows only.
    return (other_x - train_x.mean(axis=0)) / train_x.std(axis=0, ddof=1)


# Model 1: logistic regression, fitted by Newton's method -------------------

def sigmoid(z: np.ndarray) -> np.ndarray:
    return 1 / (1 + np.exp(-z))


def add_intercept(x: np.ndarray) -> np.ndarray:
    return np.column_stack([np.ones(len(x)), x])


def fit_logistic(x: np.ndarray, y: np.ndarray, max_steps: int = 50) -> tuple[np.ndarray, np.ndarray]:
    """Return (coefficients, covariance). x must already hold the intercept."""
    beta = np.zeros(x.shape[1])
    for _ in range(max_steps):
        p = sigmoid(x @ beta)
        gradient = x.T @ (y - p)
        hessian = x.T @ (x * (p * (1 - p))[:, None])
        step = np.linalg.solve(hessian, gradient)
        beta = beta + step
        if np.abs(step).max() < 1e-10:
            break
    p = sigmoid(x @ beta)
    hessian = x.T @ (x * (p * (1 - p))[:, None])
    return beta, np.linalg.inv(hessian)


def predict_logistic(beta: np.ndarray, x: np.ndarray) -> np.ndarray:
    return sigmoid(x @ beta)


# Model 2: k nearest neighbours ---------------------------------------------

def predict_knn(train_x: np.ndarray, train_y: np.ndarray, x: np.ndarray, k_values) -> dict[int, np.ndarray]:
    """For each k, the share of large-loss days among the k nearest training rows."""
    squared_distance = (
        (x**2).sum(axis=1)[:, None] - 2 * x @ train_x.T + (train_x**2).sum(axis=1)[None, :]
    )
    largest_k = min(max(k_values), len(train_x))
    nearest = np.argpartition(squared_distance, largest_k - 1, axis=1)[:, :largest_k]
    nearest_distance = np.take_along_axis(squared_distance, nearest, axis=1)
    order = np.argsort(nearest_distance, axis=1)
    nearest_labels = train_y[np.take_along_axis(nearest, order, axis=1)]
    running_share = np.cumsum(nearest_labels, axis=1) / np.arange(1, largest_k + 1)
    return {k: running_share[:, min(k, largest_k) - 1] for k in k_values}


# Step 5: the inference view ------------------------------------------------

def two_sided_p_value(z: float) -> float:
    return math.erfc(abs(z) / math.sqrt(2))


def block_bootstrap_standard_errors(x: np.ndarray, y: np.ndarray, seed: int = SEED) -> np.ndarray:
    """Refit on resamples made of whole 21-day blocks, so clusters stay together."""
    rng = np.random.default_rng(seed)
    blocks_needed = math.ceil(len(y) / BLOCK_DAYS)
    last_start = len(y) - BLOCK_DAYS
    fitted = []
    for _ in range(BOOTSTRAP_RUNS):
        starts = rng.integers(0, last_start + 1, blocks_needed)
        rows = np.concatenate([np.arange(s, s + BLOCK_DAYS) for s in starts])[: len(y)]
        beta, _ = fit_logistic(x[rows], y[rows])
        fitted.append(beta)
    return np.array(fitted).std(axis=0, ddof=1)


def inference_table(train_x: np.ndarray, train_y: np.ndarray) -> pd.DataFrame:
    x = add_intercept(standardise(train_x, train_x))
    beta, covariance = fit_logistic(x, train_y)
    standard_error = np.sqrt(np.diag(covariance))
    block_standard_error = block_bootstrap_standard_errors(x, train_y)
    table = pd.DataFrame(
        {
            "term": ["intercept"] + FEATURES,
            "coefficient": beta,
            "odds_ratio_per_1_sd": np.exp(beta),
            "standard_error": standard_error,
            "z": beta / standard_error,
            "block_standard_error": block_standard_error,
            "block_z": beta / block_standard_error,
        }
    )
    table["p_value"] = table["z"].map(two_sided_p_value)
    table["block_p_value"] = table["block_z"].map(two_sided_p_value)
    return table


# Step 6: the prediction view -----------------------------------------------

def brier(y: np.ndarray, p: np.ndarray) -> float:
    return float(((p - y) ** 2).mean())


def all_predictions(dataset: pd.DataFrame) -> pd.DataFrame:
    """One probability per row and model. Every model is fitted on train only."""
    train = dataset[dataset["split"] == "train"]
    train_x_raw = train[FEATURES].to_numpy()
    train_y = train["large_loss_next_day"].to_numpy()
    train_x = standardise(train_x_raw, train_x_raw)
    all_x = standardise(train_x_raw, dataset[FEATURES].to_numpy())

    predictions = dataset[["date", "next_date", "split", "large_loss_next_day"]].copy()
    predictions["base_rate"] = train_y.mean()
    beta, _ = fit_logistic(add_intercept(train_x), train_y)
    predictions["logistic"] = predict_logistic(beta, add_intercept(all_x))
    for k, p in predict_knn(train_x, train_y, all_x, K_VALUES).items():
        predictions[f"knn_{k}"] = p
    return predictions


def model_names() -> list[str]:
    return ["base_rate", "logistic"] + [f"knn_{k}" for k in K_VALUES]


def prediction_scores(predictions: pd.DataFrame, splits=("train", "validation")) -> pd.DataFrame:
    rows = []
    for split in splits:
        part = predictions[predictions["split"] == split]
        y = part["large_loss_next_day"].to_numpy()
        base = brier(y, part["base_rate"].to_numpy())
        for model in model_names():
            score = brier(y, part[model].to_numpy())
            rows.append(
                {"split": split, "model": model, "brier": score, "skill_vs_base_rate": 1 - score / base}
            )
    return pd.DataFrame(rows)


def choose_model(scores: pd.DataFrame) -> str:
    validation = scores[scores["split"] == "validation"]
    return str(validation.loc[validation["brier"].idxmin(), "model"])


# Step 7: bias and variance, on a simulated market with a known truth -------

def true_probability(features: pd.DataFrame, calm_cut: float, stress_cut: float) -> np.ndarray:
    """The simulated truth: three volatility regimes, and bad days matter in stress."""
    vol = features["mean_abs_return_21d"].to_numpy()
    fell_today = features["return_1d"].to_numpy() < LOSS_THRESHOLD
    p = np.full(len(features), 0.10)
    p[vol < calm_cut] = 0.03
    p[vol >= stress_cut] = 0.30
    p[(vol >= stress_cut) & fell_today] = 0.50
    return p


def bias_variance_table(dataset: pd.DataFrame, seed: int = SEED) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    calm_cut = dataset["mean_abs_return_21d"].quantile(0.50)
    stress_cut = dataset["mean_abs_return_21d"].quantile(0.90)
    feature_rows = dataset[FEATURES].reset_index(drop=True)
    truth = true_probability(feature_rows, calm_cut, stress_cut)
    raw_x = feature_rows.to_numpy()

    eval_rows = rng.integers(0, len(raw_x), SIMULATED_EVAL_POINTS)
    eval_truth = truth[eval_rows]
    noise = float((eval_truth * (1 - eval_truth)).mean())

    rows = []
    for train_size in SIMULATED_TRAIN_SIZES:
        runs: dict[str, list[np.ndarray]] = {name: [] for name in model_names()[1:]}
        for _ in range(SIMULATED_RUNS):
            train_rows = rng.integers(0, len(raw_x), train_size)
            train_y = (rng.random(train_size) < truth[train_rows]).astype(float)
            train_x = standardise(raw_x[train_rows], raw_x[train_rows])
            eval_x = standardise(raw_x[train_rows], raw_x[eval_rows])
            beta, _ = fit_logistic(add_intercept(train_x), train_y)
            runs["logistic"].append(predict_logistic(beta, add_intercept(eval_x)))
            for k, p in predict_knn(train_x, train_y, eval_x, K_VALUES).items():
                runs[f"knn_{k}"].append(p)
        for model, run_predictions in runs.items():
            stacked = np.array(run_predictions)            # runs x evaluation points
            bias_squared = float(((stacked.mean(axis=0) - eval_truth) ** 2).mean())
            variance = float(stacked.var(axis=0).mean())
            rows.append(
                {
                    "train_rows": train_size,
                    "model": model,
                    "noise": noise,
                    "bias_squared": bias_squared,
                    "variance": variance,
                    "expected_brier": noise + bias_squared + variance,
                    "excess_over_noise": bias_squared + variance,
                }
            )
    return pd.DataFrame(rows)


# Step 8: the confusion matrix, with costs ----------------------------------

def confusion(y: np.ndarray, alarm: np.ndarray) -> dict[str, int]:
    return {
        "true_positive": int(((alarm == 1) & (y == 1)).sum()),
        "false_positive": int(((alarm == 1) & (y == 0)).sum()),
        "false_negative": int(((alarm == 0) & (y == 1)).sum()),
        "true_negative": int(((alarm == 0) & (y == 0)).sum()),
    }


def total_cost(counts: dict[str, int], cost_ratio: float) -> float:
    return counts["false_positive"] * COST_FALSE_ALARM + counts["false_negative"] * COST_FALSE_ALARM * cost_ratio


def break_even_threshold(cost_ratio: float) -> float:
    # Alarm when p * cost_of_miss > (1 - p) * cost_of_false_alarm.
    return 1 / (1 + cost_ratio)


def threshold_table(predictions: pd.DataFrame, model: str) -> pd.DataFrame:
    """Score every candidate threshold on validation, for each cost ratio."""
    validation = predictions[predictions["split"] == "validation"]
    y = validation["large_loss_next_day"].to_numpy()
    p = validation[model].to_numpy()
    candidates = np.round(np.arange(0.01, 0.61, 0.01), 2)
    rows = []
    for cost_ratio in COST_RATIOS:
        for threshold in candidates:
            counts = confusion(y, (p >= threshold).astype(int))
            rows.append(
                {"cost_ratio": cost_ratio, "threshold": threshold, **counts,
                 "cost": total_cost(counts, cost_ratio)}
            )
    return pd.DataFrame(rows)


def chosen_thresholds(thresholds: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for cost_ratio, group in thresholds.groupby("cost_ratio"):
        best = group.loc[group["cost"].idxmin()]
        rows.append(
            {
                "cost_ratio": cost_ratio,
                "break_even_threshold": break_even_threshold(cost_ratio),
                "validation_best_threshold": best["threshold"],
                "validation_cost_at_best": best["cost"],
            }
        )
    return pd.DataFrame(rows)


def decision_rules(model: str, cost_ratio: float, validation_threshold: float) -> list[tuple[str, str, float]]:
    """(rule name, probability column, threshold). The thresholds are fixed before test."""
    return [
        ("never_alarm", "base_rate", 2.0),
        ("always_alarm", "base_rate", 0.0),
        (f"{model}_at_0.50", model, 0.50),
        (f"{model}_at_break_even", model, break_even_threshold(cost_ratio)),
        (f"{model}_at_validation_best", model, validation_threshold),
    ]


def confusion_table(predictions: pd.DataFrame, split: str, rules, cost_ratio: float) -> pd.DataFrame:
    part = predictions[predictions["split"] == split]
    y = part["large_loss_next_day"].to_numpy()
    rows = []
    for name, column, threshold in rules:
        alarm = (part[column].to_numpy() >= threshold).astype(int)
        counts = confusion(y, alarm)
        caught = counts["true_positive"] + counts["false_negative"]
        alarms = counts["true_positive"] + counts["false_positive"]
        rows.append(
            {
                "rule": name,
                "threshold": threshold,
                **counts,
                "accuracy": (counts["true_positive"] + counts["true_negative"]) / len(y),
                "recall": counts["true_positive"] / caught if caught else float("nan"),
                "precision": counts["true_positive"] / alarms if alarms else float("nan"),
                "cost": total_cost(counts, cost_ratio),
            }
        )
    return pd.DataFrame(rows)


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

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


def run(prices_path: Path | None, out: Path, start: str | None, end: str | None) -> 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)

    dataset = add_split(add_label(add_features(returns)))
    dataset.to_csv(out / "03_dataset.csv", index=False)
    pd.options.display.float_format = "{:.6g}".format
    save(split_summary(dataset), out, "04_split_summary.csv")

    train = dataset[dataset["split"] == "train"]
    save(
        inference_table(train[FEATURES].to_numpy(), train["large_loss_next_day"].to_numpy()),
        out, "05_inference_coefficients.csv",
    )

    predictions = all_predictions(dataset)
    predictions.to_csv(out / "06_predictions.csv", index=False)
    scores = prediction_scores(predictions)
    save(scores, out, "06b_prediction_scores.csv")
    model = choose_model(scores)
    print(f"\nmodel chosen on validation: {model}")

    save(bias_variance_table(dataset), out, "07_bias_variance.csv")

    thresholds = threshold_table(predictions, model)
    thresholds.to_csv(out / "08_thresholds_validation.csv", index=False)
    chosen = chosen_thresholds(thresholds)
    save(chosen, out, "08b_chosen_thresholds.csv")

    # Everything above was decided without the test split. Now it is scored once.
    validation_threshold = float(
        chosen.loc[chosen["cost_ratio"] == REPORT_COST_RATIO, "validation_best_threshold"].iloc[0]
    )
    rules = decision_rules(model, REPORT_COST_RATIO, validation_threshold)
    save(prediction_scores(predictions, splits=("test",)), out, "09_test_scores.csv")
    save(confusion_table(predictions, "test", rules, REPORT_COST_RATIO), out, "09b_test_confusion.csv")


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

class Checks(unittest.TestCase):
    def test_logistic_recovers_known_coefficients(self):
        rng = np.random.default_rng(1)
        x = add_intercept(rng.normal(size=(50_000, 2)))
        true_beta = np.array([-2.0, 0.8, -0.5])
        y = (rng.random(50_000) < sigmoid(x @ true_beta)).astype(float)
        beta, covariance = fit_logistic(x, y)
        self.assertTrue(np.allclose(beta, true_beta, atol=0.05))
        self.assertTrue((np.sqrt(np.diag(covariance)) < 0.03).all())

    def test_knn_with_every_row_is_the_base_rate(self):
        rng = np.random.default_rng(2)
        train_x, train_y = rng.normal(size=(40, 3)), (rng.random(40) < 0.3).astype(float)
        p = predict_knn(train_x, train_y, rng.normal(size=(5, 3)), (1, 40))
        self.assertTrue(np.allclose(p[40], train_y.mean()))
        self.assertTrue(set(np.unique(p[1])) <= {0.0, 1.0})

    def test_knn_finds_the_nearest_row(self):
        train_x = np.array([[0.0], [1.0], [5.0]])
        train_y = np.array([0.0, 1.0, 0.0])
        p = predict_knn(train_x, train_y, np.array([[0.9], [4.0]]), (1, 2))
        self.assertEqual(list(p[1]), [1.0, 0.0])
        self.assertEqual(list(p[2]), [0.5, 0.5])

    def test_features_do_not_see_tomorrow(self):
        dates = pd.bdate_range("2024-01-01", periods=40)
        close = 100 * np.exp(np.cumsum(np.random.default_rng(3).normal(0, 0.01, 40)))
        base = add_label(add_features(add_returns(pd.DataFrame({"date": dates, "close": close}))))
        changed = close.copy()
        changed[-1] = changed[-1] * 0.5           # crash on the very last day
        after = add_label(add_features(add_returns(pd.DataFrame({"date": dates, "close": changed}))))
        self.assertTrue(np.allclose(base[FEATURES], after[FEATURES]))
        self.assertEqual(after["large_loss_next_day"].iloc[-1], 1)

    def test_confusion_counts_and_cost(self):
        y = np.array([1, 1, 0, 0, 0])
        counts = confusion(y, np.array([1, 0, 1, 0, 0]))
        self.assertEqual(sum(counts.values()), 5)
        self.assertEqual(total_cost(counts, 10), 1 * 1 + 1 * 10)
        self.assertAlmostEqual(break_even_threshold(10), 1 / 11)


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("--out", type=Path, default=Path("part02_output"))
    parser.add_argument("--check", action="store_true", help="run self-checks")
    args = parser.parse_args()
    if args.check:
        unittest.main(argv=["prediction_vs_inference"], verbosity=2)
    else:
        run(args.prices, args.out, args.start, args.end)


if __name__ == "__main__":
    main()