"""overnight_recovery_top10, ported from quant-trading `strategies/robust_momentum.py`
(`RobustMomentumStrategy("overnight_recovery", top_n=10, rebal_freq=10, buffer=2.0)`).

Score: mean cross-sectional percentile rank of overnight-return momentum (12 months,
skipping the last one) and distance from the 126-day low. Every `rebalance_every`
sessions the book holds the top_n names equal-weighted, keeping a held name while its
rank is within top_n * (1 + buffer); in between nothing trades and weights drift.
"""

import datetime as dt

import numpy as np
import pandas as pd

# Rows the source zeroes before trading; overnight momentum needs 180 + 21 of them.
WARMUP = 252


def overnight_momentum(opens: pd.DataFrame, closes: pd.DataFrame, window: int = 252, skip: int = 21) -> pd.DataFrame:
    """Sum of log overnight returns (open / previous close), Lou, Polk & Skouras (2019)."""
    overnight = np.log(opens / closes.ffill().shift(1))
    return overnight.rolling(window, min_periods=180).sum().shift(skip)


def recovery(closes: pd.DataFrame, window: int = 126) -> pd.DataFrame:
    c = closes.ffill()
    return c / c.rolling(window, min_periods=window).min() - 1


def signal_overnight_recovery(
    opens: pd.DataFrame, closes: pd.DataFrame, member: pd.DataFrame | None = None
) -> pd.DataFrame:
    """Row t uses bars through t's close. `member` (dates x symbols) limits the ranks to
    each day's index members; factors still use each name's full price history. Live,
    every column is a current member. A name without a close on a row is not eligible."""
    factors = [overnight_momentum(opens, closes), recovery(closes)]
    if member is not None:
        factors = [f.where(member) for f in factors]
    score = sum(f.rank(axis=1, pct=True) for f in factors) / len(factors)
    return score.where(closes.notna())


def select_with_buffer(score: pd.Series, held: set[str], top_n: int, buffer: float) -> dict[str, float]:
    """Equal weights on the top_n names of `score`, keeping held names while ranked
    within top_n * (1 + buffer). Ties keep column order, as in the source."""
    valid = score.dropna()
    if len(valid) < top_n:
        raise ValueError(f"only {len(valid)} names have a score, need {top_n}")
    order = valid.index[np.argsort(-valid.to_numpy(), kind="stable")]
    rank = {symbol: i for i, symbol in enumerate(order)}
    chosen = sorted((s for s in held if rank.get(s, np.inf) < top_n * (1 + buffer)), key=rank.__getitem__)[:top_n]
    chosen += [s for s in order if s not in chosen][: top_n - len(chosen)]
    return {symbol: 1.0 / top_n for symbol in chosen}


def sessions_until_rebalance(sessions: pd.DatetimeIndex, anchor: dt.date, every: int) -> int | None:
    """0 when the last of `sessions` is a rebalance session: the first session on or
    after `anchor` and every `every` sessions after it. None before the anchor."""
    n = int((sessions.date >= anchor).sum()) - 1
    return None if n < 0 else -n % every