"""Monthly cross-sectional MAX backtest with realistic transaction costs.

Timing convention (no look-ahead)
---------------------------------
* The MAX signal for *formation* month ``p`` is built only from daily data
  inside month ``p``.
* Stocks are sorted into quintiles at the close of month ``p``.
* The portfolio is held through month ``p + 1`` and earns that month's return.

So a signal observed at time T only ever drives a position over (T, T+1].

Portfolio
---------
Equal-weighted quintiles.  Q1 = lowest MAX (long), Q5 = highest MAX (short).
The headline long-short return is ``Q1 - Q5`` (buy the boring stocks, sell the
lottery tickets).  Names whose next-month return is missing (mid-month delist)
are dropped and the leg is re-normalised, so the universe degrades gracefully.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd

from . import config
from .signals import max_signal, monthly_returns


@dataclass
class BacktestResult:
    quintile_returns: pd.DataFrame   # columns Q1..Qn, index = holding month-end
    long_short_gross: pd.Series
    long_short_net: pd.Series
    turnover: pd.Series              # combined long+short one-way turnover
    holdings: dict[pd.Timestamp, dict[str, list[str]]]
    params: config.BacktestParams


def _equal_weights(names: list[str]) -> pd.Series:
    if not names:
        return pd.Series(dtype="float64")
    return pd.Series(1.0 / len(names), index=names)


def _assign_quintiles(signal_row: pd.Series, n_q: int) -> pd.Series:
    """Map a cross-section of MAX values to quintile labels 1..n_q (1 = lowest)."""
    s = signal_row.dropna()
    if s.size < n_q:
        return pd.Series(dtype="float64")
    # ``rank`` then ``qcut`` on the ranks avoids failures from tied MAX values.
    ranks = s.rank(method="first")
    labels = pd.qcut(ranks, n_q, labels=range(1, n_q + 1))
    return pd.Series(labels.astype(int), index=s.index)


def run_backtest(
    prices: pd.DataFrame,
    params: config.BacktestParams | None = None,
) -> BacktestResult:
    """Run the monthly MAX quintile backtest on ``prices`` (adjusted close)."""
    params = params or config.BacktestParams()
    from .signals import daily_returns

    rets = daily_returns(prices)
    sig = max_signal(rets, n_high=params.max_n, min_days=params.min_days_in_month)
    mret = monthly_returns(rets)

    # Align everything on a monthly PeriodIndex and form next-month returns.
    sig_p = sig.copy()
    sig_p.index = sig_p.index.to_period("M")
    mret_p = mret.copy()
    mret_p.index = mret_p.index.to_period("M")
    fwd = mret_p.shift(-1).reindex(sig_p.index)  # value at p is the return of p+1

    n_q = params.n_quintiles
    cost_rate = params.cost_bps / 1e4

    q_cols = [f"Q{i}" for i in range(1, n_q + 1)]
    q_returns: dict[pd.Timestamp, dict[str, float]] = {}
    turnover_rec: dict[pd.Timestamp, float] = {}
    holdings: dict[pd.Timestamp, dict[str, list[str]]] = {}

    # Previous-period state for turnover: realised target weights per quintile.
    prev_target: dict[int, pd.Series] = {q: pd.Series(dtype="float64") for q in range(1, n_q + 1)}
    prev_ret: dict[int, pd.Series] = {q: pd.Series(dtype="float64") for q in range(1, n_q + 1)}

    for form_p in sig_p.index:
        quintiles = _assign_quintiles(sig_p.loc[form_p], n_q)
        if quintiles.empty:
            continue
        fwd_row = fwd.loc[form_p]
        if fwd_row.isna().all():
            continue  # no following month (end of sample)
        hold_ts = (form_p + 1).to_timestamp(how="end").normalize()

        q_ret_row: dict[str, float] = {}
        leg_turnover = 0.0
        holdings[hold_ts] = {}
        for q in range(1, n_q + 1):
            names = quintiles.index[quintiles == q].tolist()
            # Only hold names with a realised next-month return.
            r = fwd_row.reindex(names).dropna()
            held = r.index.tolist()
            holdings[hold_ts][f"Q{q}"] = held
            w = _equal_weights(held)
            q_ret_row[f"Q{q}"] = float((w * r.reindex(w.index)).sum()) if len(held) else np.nan

            # Turnover only matters for the traded legs (long Q1, short Qn).
            if q in (1, n_q):
                drift = prev_target[q] * (1.0 + prev_ret[q].reindex(prev_target[q].index).fillna(0.0))
                drift = drift / drift.sum() if drift.sum() != 0 else drift
                union = w.index.union(drift.index)
                to = (w.reindex(union, fill_value=0.0) - drift.reindex(union, fill_value=0.0)).abs().sum()
                leg_turnover += float(to)
            # Save state for next iteration's drift computation.
            prev_target[q] = w
            prev_ret[q] = r.reindex(w.index)

        q_returns[hold_ts] = q_ret_row
        turnover_rec[hold_ts] = leg_turnover

    qdf = pd.DataFrame.from_dict(q_returns, orient="index").reindex(columns=q_cols).sort_index()
    qdf.index.name = "month"
    turnover = pd.Series(turnover_rec).sort_index()
    turnover.index.name = "month"

    long_short_gross = (qdf["Q1"] - qdf[f"Q{n_q}"]).rename("long_short_gross")
    long_short_net = (long_short_gross - cost_rate * turnover).rename("long_short_net")

    return BacktestResult(
        quintile_returns=qdf,
        long_short_gross=long_short_gross,
        long_short_net=long_short_net,
        turnover=turnover,
        holdings=holdings,
        params=params,
    )


# ---------------------------------------------------------------------------
# Performance statistics
# ---------------------------------------------------------------------------
def performance_stats(returns: pd.Series, periods_per_year: int = 12) -> dict[str, float]:
    """Annualised performance summary for a monthly return series."""
    r = returns.dropna()
    if r.empty:
        return {k: float("nan") for k in
                ("ann_return", "ann_vol", "sharpe", "t_stat", "max_drawdown", "hit_rate", "n_months")}
    mean = r.mean()
    vol = r.std(ddof=1)
    ann_return = mean * periods_per_year
    ann_vol = vol * np.sqrt(periods_per_year)
    sharpe = ann_return / ann_vol if ann_vol > 0 else float("nan")
    t_stat = mean / (vol / np.sqrt(len(r))) if vol > 0 else float("nan")
    equity = (1.0 + r).cumprod()
    drawdown = equity / equity.cummax() - 1.0
    return {
        "ann_return": float(ann_return),
        "ann_vol": float(ann_vol),
        "sharpe": float(sharpe),
        "t_stat": float(t_stat),
        "max_drawdown": float(drawdown.min()),
        "hit_rate": float((r > 0).mean()),
        "n_months": int(len(r)),
    }


def equity_curve(returns: pd.Series) -> pd.Series:
    """Cumulative growth of $1 (geometric)."""
    r = returns.dropna()
    return (1.0 + r).cumprod()
