"""CEF screen: stage 1 of the CEF pass.

Reads cef_universe.json (SEC active-CEF report + ticker map) and local
price files (goget-downloaded).  Per ticker:
  t5 / t12        total return (Adj Close)
  vol5            annualized daily-return vol, 5y
  maxdd5          deepest peak->trough drawdown, 5y
  scen_*          return over each of the 5 market-crash episodes
                  (drawdown.py episodes, fund_windows)
  payout_12m      distribution proxy = 12m adj return - 12m price
                  return (Close = market price, Adj Close = total
                  return incl. reinvested distributions)
  p5y_share       5y version of the same

CEF-specific context for interpreting the output:
- price = MARKET price; premium/discount dynamics are on top of NAV
  (stage 2 adds the NPORT NAV-per-share series)
- distributions are NOT the fund's discretion to skip - a CEF usually
  keeps paying even in drawdowns (that's what the leverage is for),
  so the payout proxy is more stable than for open-end funds
- BDCs in the set pay ordinary income + occasional ROC (stage 2
  flags them from the EDGAR form history)

Output: fundlab/cef_screen.json
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd

from fundlab import drawdown

HERE = Path(__file__).parent
UNIVERSE = HERE / "cef_universe.json"
EPS = json.loads((HERE / "drawdown_results.json").read_text())["episodes"]
OUT = HERE / "cef_screen.json"
SCEN_KEYS = ["2022 bear mkt", "2023 rate shock", "2024 vol spike",
             "2025 tariff crash", "2026 Q1 drawdown"]


def _load(sym: str) -> pd.DataFrame | None:
    f = Path.home() / "prog/fin/stocks" / f"{sym.lower()}-history.csv"
    if not f.exists():
        return None
    try:
        d = pd.read_csv(f, parse_dates=["Date"])
    except Exception:
        return None
    d = d.set_index("Date").sort_index()
    d = d[~d.index.duplicated(keep="last")]
    return d


def _perf(d: pd.DataFrame) -> dict:
    adj = d["Adj Close"].dropna()
    px = d["Close"].dropna()
    out: dict = {}
    for tag, n in (("5", 5 * 252), ("12", 252)):
        a = adj.tail(n)
        p = px.tail(n)
        out[f"t{tag}"] = float(a.iloc[-1] / a.iloc[0] - 1) if len(a) > 20 else None
        out[f"p{tag}"] = float(p.iloc[-1] / p.iloc[0] - 1) if len(p) > 20 else None
    a5 = adj.tail(5 * 252)
    if len(a5) > 20:
        r = a5.pct_change().dropna()
        out["vol5"] = float(r.std() * np.sqrt(252))
        peak = np.maximum.accumulate(a5.to_numpy())
        out["maxdd5"] = float(((a5.to_numpy() / peak) - 1).min())
    else:
        out["vol5"] = out["maxdd5"] = None
    if out.get("t12") is not None and out.get("p12") is not None:
        out["payout_12m"] = out["t12"] - out["p12"]
    if out.get("t5") is not None and out.get("p5") is not None:
        out["payout_5y_share"] = (out["t5"] - out["p5"]) / out["t5"] \
            if abs(out["t5"]) > 1e-4 else None
    return out


def screen_one(sym: str) -> dict:
    s = sym.lower()          # price files are lowercase
    d = _load(s)
    if d is None or len(d) < 30:
        return {"error": "no data"}
    out = _perf(d)
    out.update(drawdown.fund_windows(s, EPS))
    return out


def run() -> dict:
    uni = json.loads(UNIVERSE.read_text())
    res: dict = {}
    for t in sorted(uni):
        info = dict(uni[t])
        info["sym"] = t.lower()
        s = screen_one(t)
        info.update(s)
        info["n_pos_scen"] = sum(
            1 for k in SCEN_KEYS if info.get(k) is not None and info[k] > 0)
        res[t] = info
    n_ok = sum(1 for v in res.values() if "t5" in v)
    print(f"{n_ok}/{len(res)} CEFs screened", flush=True)
    OUT.write_text(json.dumps(res, indent=1, default=str))
    print(f"wrote {OUT}")
    return res


def _fmt(x, pct: bool = True, dec: int = 2) -> str:
    if x is None:
        return ""
    if pct:
        return f"{x:+.{dec}%}"
    return f"{x:.{dec}f}"


def _print(res: dict, top: int = 30, sort: str = "t5") -> None:
    rows = [v for v in res.values() if v.get("t5") is not None]
    print(f"{'fund':6} {'name':42} {'5yTR':>8} {'vol':>6} {'maxDD':>8}"
          f" {'2022':>8} {'2023':>8} {'2025t':>8} {'2026':>8}"
          f" {'dist12':>8} {'n+':>3}")
    for v in sorted(rows, key=lambda x: -x.get(sort, -9))[:top]:
        print(f"{v['sym']:6} {v.get('name', '')[:42]:42}"
              f" {_fmt(v.get('t5')):>8} {_fmt(v.get('vol5'), False):>6}"
              f" {_fmt(v.get('maxdd5')):>8}"
              f" {_fmt(v.get('2022 bear mkt')):>8}"
              f" {_fmt(v.get('2023 rate shock')):>8}"
              f" {_fmt(v.get('2025 tariff crash')):>8}"
              f" {_fmt(v.get('2026 Q1 drawdown')):>8}"
              f" {_fmt(v.get('payout_12m')):>8}"
              f" {v.get('n_pos_scen', 0):>3}")


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