"""Drawdown-scenario screen: which funds held up / gained when equities
had their worst episodes in recent years?

Scenarios are DETECTED from the index (IVV, S&P 500) rather than
hard-coded: contiguous peak-to-trough episodes where the index fell
>= 8% (MIN_DD); one scenario per calendar year = that year's worst
peak->trough (so the 2023 rate shock and the 2022 bear market stay
separate, and shallow years drop out).

For every screened fund we measure its total return over each window
(from its own Adj Close, using its first print after the peak and its
last print on/before the trough) and rank by how many scenarios it was
positive in.

Output: fundlab/drawdown_results.json + console table.
Usage:  python -m fundlab.drawdown
"""

from __future__ import annotations

import json
from pathlib import Path

import pandas as pd

from fundlab.decompose import DATA

HERE = Path(__file__).parent
RESULTS = HERE / "drawdown_results.json"
INDEX = "ivv"
MIN_DD = 0.08
SINCE = "2022-01-01"
LOCAL_ANCHOR_MONTHS = 12


def index_series(sym: str = INDEX) -> pd.Series:
    f = DATA / f"{sym}-history.csv"
    s = (pd.read_csv(f, parse_dates=["Date"], index_col="Date")
         ["Adj Close"].dropna())
    return s[~s.index.duplicated(keep="last")].sort_index()


def detect_episodes(p: pd.Series, min_dd: float = MIN_DD,
                    since: str = SINCE) -> list[dict]:
    """One scenario per calendar year: the year's worst peak->trough.

    Per-year bands keep distinct crashes separate (the 2023 rate shock
    is its own window, not "the 2022 bear market part 2") and drop
    shallow years (2024's Aug-5 dip was only ~5%).
    """
    p = p[p.index >= since]
    out = []
    for year, g in p.groupby(p.index.year):
        rm = g.cummax()
        dd = g / rm - 1
        if dd.min() > -min_dd:
            continue
        trough = dd.idxmin()
        peak = rm.loc[:trough].idxmax()
        out.append({"peak": peak, "trough": trough,
                    "min_dd": float(dd.min()), "year": int(year)})
    out.sort(key=lambda e: e["peak"])
    return out


_NAMED = {2022: "2022 bear mkt", 2023: "2023 rate shock",
           2024: "2024 vol spike", 2025: "2025 tariff crash",
           2026: "2026 Q1 drawdown"}


def _label(e: dict) -> str:
    return _NAMED.get(e["year"], f"{e['year']} drawdown")


def fund_windows(sym: str, eps: list[dict]) -> dict[str, float]:
    """Fund total return over each episode, from one CSV load."""
    f = DATA / f"{sym}-history.csv"
    if not f.exists():
        return {}
    try:
        s = (pd.read_csv(f, parse_dates=["Date"], index_col="Date")
             ["Adj Close"].dropna())
        s = s[~s.index.duplicated(keep="last")].sort_index()
    except Exception:
        return {}
    if len(s) < 30:
        return {}
    vals = s.to_numpy()
    idx = s.index
    out: dict[str, float] = {}
    for e in eps:
        peak = pd.Timestamp(e["peak"])
        trough = pd.Timestamp(e["trough"])
        i0 = idx.searchsorted(peak, side="right")   # first print after peak
        i1 = idx.searchsorted(trough, side="right") - 1  # last print <= trough
        if i0 >= len(vals) or i1 < 0 or i1 <= i0:
            continue
        a, b = vals[i0], vals[i1]
        if a <= 0 or not (pd.notna(a) and pd.notna(b)):
            continue
        out[e["label"]] = float(b / a - 1)
    return out


def run() -> dict:
    idx = index_series()
    eps = detect_episodes(idx)
    for e in eps:
        e["label"] = _label(e)
        e["peak"] = str(e["peak"].date())
        e["trough"] = str(e["trough"].date())
    fr = json.loads((HERE / "factor_results.json").read_text())

    out: dict[str, dict] = {}
    for sym, meta in fr.items():
        rets = fund_windows(sym, eps)
        if not rets:
            continue
        n_pos = sum(1 for r in rets.values() if r > 0)
        out[sym] = {
            "name": meta.get("name", ""),
            "verdict": meta.get("verdict", ""),
            "t5": meta.get("alpha_t_5y"),
            "corr_port": meta.get("corr_portfolio"),
            "rets": rets,
            "n_avail": len(rets),
            "n_pos": n_pos,
            "min_ret": min(rets.values()),
            "max_ret": max(rets.values()),
            # positive in every scenario it had data for
            "all_pos": n_pos == len(rets),
        }
    res = {"index": INDEX, "since": SINCE, "min_dd": MIN_DD,
           "episodes": eps, "funds": out}
    RESULTS.write_text(json.dumps(res, indent=1))
    return res


def _print(res: dict) -> None:
    print(f"Index: {res['index']}  episodes (min drawdown "
          f"{res['min_dd']*100:.0f}% since {res['since']}):")
    for e in res["episodes"]:
        print(f"  {e['label']:<20} {e['peak']} -> {e['trough']}  "
              f"({e['min_dd']*100:.1f}%)")
    funds = {s: v for s, v in res["funds"].items()
             if v["verdict"].startswith("CANDIDATE")}
    funds = dict(sorted(funds.items(),
                        key=lambda kv: (-kv[1]["n_pos"],
                                        -kv[1]["n_avail"],
                                        kv[1]["min_ret"])))
    ep = res["episodes"]
    hdr = "  ".join(f"{e['label'][:9]:>10}" for e in ep)

    def row(s: str, v: dict) -> str:
        cols = "  ".join(
            (f"{v['rets'][e['label']]*100:+8.1f}%"
             if e["label"] in v["rets"] else f"{'n/a':>10}")
            for e in ep)
        cp = v.get("corr_port")
        cps = f"{cp:+.2f}" if isinstance(cp, (int, float)) else "  -"
        t5 = (f"{v['t5']:+.1f}" if isinstance(v.get("t5"),
                                              (int, float)) else "  -")
        return (f"  {s.upper()[:6]:<7}{v['name'][:34]:<35} {cols}  "
                f"corr{cps:>5}  t5 {t5}")

    for label_ in ("5/5 (all)", "4/5", "3/5"):
        want = {"5/5 (all)": lambda v: v["n_pos"] == 5 and v["n_avail"] >= 4,
                "4/5": lambda v: v["n_pos"] == 4,
                "3/5": lambda v: v["n_pos"] == 3}[label_]
        grp = [(s, v) for s, v in funds.items() if want(v)]
        if not grp:
            continue
        print(f"\n== {label_} positive  ({len(grp)})  "
              f"{'fund':<6}{'name':<35} {hdr}  corr  5y-t")
        for s, v in grp[:25]:
            print(row(s, v))


if __name__ == "__main__":
    _print(run())