"""Unit tests for portfolio construction, metrics and data loading robustness.

Regressions covered:
  * _rebalance_dates must return the first TRADING day of each period.
    (The old code returned the Grouper period-END labels — Sundays /
    month-ends — which never exist in a trading-day index, so every
    non-daily rebalance scheme silently never fired.)
  * metrics.beta_alpha must not divide by zero on a flat benchmark.
  * _read_csv_clean must coerce downloader nil tokens ('%!f(<nil>)')
    to NaN, keeping price columns numeric (a single such token used to
    turn the whole column into string and crash every pct_change).

Run:  .venv/bin/python tests/test_portfolio.py
"""
from __future__ import annotations

import sys
import tempfile
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import data  # noqa: E402
import metrics  # noqa: E402
from portfolio import _rebalance_dates, portfolio_returns  # noqa: E402

PASS, FAIL = 0, 0


def check(name: str, ok: bool) -> None:
    global PASS, FAIL
    PASS += bool(ok)
    FAIL += not ok
    print(("  ok    " if ok else "  FAIL  ") + name)


def main() -> int:
    idx = pd.bdate_range("2023-01-02", "2023-12-29")   # full business year
    i0, i1 = idx[0], idx[-1]

    # --- rebalance dates land on real trading days ----------------------
    for freq in ("1W", "1ME", "QE", "YE"):
        dates = _rebalance_dates(idx, freq)
        check(f"{freq}: all rebalance dates are trading days",
              bool(dates) and dates <= set(idx))
        check(f"{freq}: first day included", i0 in dates)

    # bdate_range has every weekday, so the first trading day of each
    # week/month is exactly each Monday / the min of each month's slice
    weekly = sorted(_rebalance_dates(idx, "1W"))
    check("1W: rebalances exactly on the Mondays (not the Sunday labels)",
          weekly == [d for d in idx if d.weekday() == 0])
    monthly = sorted(_rebalance_dates(idx, "1ME"))
    check("1ME: first trading day of each month",
          monthly == [d for m in range(1, 13)
                      for d in [idx[idx.month == m].min()]])
    check("YE: one rebalance, on the first day",
          _rebalance_dates(idx, "YE") == {i0})
    check("1D: every day", _rebalance_dates(idx, "1D") == set(idx))

    # --- the fix actually makes non-daily rebalancing happen -----------
    p = pd.DataFrame({"a": np.linspace(10, 12, len(idx)),
                      "b": np.linspace(10, 8, len(idx))}, index=idx)
    r_none = portfolio_returns(p, {"a": 0.5, "b": 0.5}, rebalance=None)
    r_wk = portfolio_returns(p, {"a": 0.5, "b": 0.5}, rebalance="1W")
    check("1W: rebalances fire (turnover > 0 on rebalance days)",
          r_wk.turnover.loc[weekly[1:]].gt(0).all()
          and float(r_wk.turnover.sum()) > float(r_none.turnover.sum()) + 1e-9)
    check("1W: rebalance_dates reported on trading days",
          set(r_wk.rebalance_dates) == set(weekly) and
          set(r_wk.rebalance_dates) <= set(idx))

    # --- beta_alpha: flat benchmark must not divide by zero ------------
    r = pd.Series(np.random.default_rng(0).normal(0, 0.01, 100), index=idx[:100])
    flat = pd.Series(0.0, index=idx[:100])
    beta, alpha = metrics.beta_alpha(r, flat, 0.0)
    check("beta_alpha: flat benchmark -> beta 0, no crash",
          beta == 0.0 and np.isfinite(alpha))

    # --- nil tokens in CSVs stay numeric --------------------------------
    with tempfile.TemporaryDirectory() as td:
        f = Path(td) / "x-history.csv"
        f.write_text("Date,Open,High,Low,Close,Adj Close,Volume\n"
                     "2024-01-02,1.0,1.1,0.9,1.0,1.0,100\n"
                     "2024-01-03,1.0,1.1,0.9,%!f(<nil>),%!f(<nil>),100\n"
                     "2024-01-04,1.0,1.1,0.9,1.2,1.2,100\n")
        df = data._read_csv_clean(f)
        close = df["Close"]
        check("nil token: Close column stays numeric",
              pd.api.types.is_numeric_dtype(close))
        check("nil token: the bad cell becomes NaN, others intact",
              close.isna().sum() == 1 and close.iloc[0] == 1.0 and close.iloc[2] == 1.2)
        # the real pipeline drops NaN rows before pct_change
        rets = df["Close"].dropna().pct_change().fillna(0.0)
        check("nil token: dropna + pct_change works (the old crash)",
              abs(float(rets.iloc[-1]) - (1.2 / 1.0 - 1)) < 1e-9)

    print(f"\n{PASS} passed, {FAIL} failed")
    return 1 if FAIL else 0


if __name__ == "__main__":
    sys.exit(main())