"""Signal construction: daily returns, the MAX(N) lottery proxy, and IVOL.

All signals for a given formation month use only data *within that month*, and
positions formed from them are held in the *following* month (see
``src.backtest``).  This strict separation is what rules out look-ahead bias.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

from . import config


def daily_returns(prices: pd.DataFrame) -> pd.DataFrame:
    """Simple daily returns from adjusted close prices."""
    return prices.sort_index().pct_change()


def _month_key(index: pd.DatetimeIndex) -> pd.PeriodIndex:
    return index.to_period("M")


def max_signal(
    rets: pd.DataFrame,
    n_high: int = config.DEFAULT_MAX_N,
    min_days: int = config.MIN_DAYS_IN_MONTH,
) -> pd.DataFrame:
    """Monthly MAX(N): mean of the ``n_high`` largest daily returns each month.

    Parameters
    ----------
    rets
        Daily returns, dates on the index, tickers on the columns.
    n_high
        Number of highest daily returns to average (1 -> classic MAX).
    min_days
        Minimum valid daily observations in a month for a value to be produced;
        otherwise the cell is ``NaN`` (handles halts / partial listings).

    Returns
    -------
    DataFrame indexed by month-end timestamp, columns = tickers.
    """
    if n_high < 1:
        raise ValueError("n_high must be >= 1")

    def _top_mean(block: pd.Series) -> float:
        vals = block.dropna()
        if len(vals) < min_days:
            return np.nan
        k = min(n_high, len(vals))
        return vals.nlargest(k).mean()

    grouped = rets.groupby(_month_key(rets.index))
    out = grouped.agg(lambda col: _top_mean(col))
    out.index = out.index.to_timestamp(how="end").normalize()
    out.index.name = "month"
    return out


def monthly_returns(rets: pd.DataFrame) -> pd.DataFrame:
    """Compound daily returns to total return within each calendar month."""
    grouped = rets.groupby(_month_key(rets.index))
    out = grouped.apply(lambda block: (1.0 + block).prod(min_count=1) - 1.0)
    out.index = out.index.to_timestamp(how="end").normalize()
    out.index.name = "month"
    return out


def month_valid_counts(rets: pd.DataFrame) -> pd.DataFrame:
    """Count of non-NaN daily returns per ticker per month."""
    out = rets.groupby(_month_key(rets.index)).count()
    out.index = out.index.to_timestamp(how="end").normalize()
    out.index.name = "month"
    return out


def idiosyncratic_vol(
    rets: pd.DataFrame,
    factors: pd.DataFrame,
    min_days: int = config.MIN_DAYS_IN_MONTH,
) -> pd.DataFrame:
    """Monthly idiosyncratic volatility from a daily FF3 regression.

    For each ticker and each calendar month we regress the stock's daily excess
    return on the contemporaneous Mkt-RF, SMB and HML factors and take the
    standard deviation of the residuals (Ang, Hodrick, Xing, Zhang style).  The
    result is on the same monthly grid as :func:`max_signal`, enabling a direct
    cross-sectional correlation between MAX and IVOL.
    """
    fac = factors.reindex(rets.index).dropna(how="any", subset=["MKT_RF", "SMB", "HML", "RF"])
    common = rets.index.intersection(fac.index)
    rets = rets.loc[common]
    fac = fac.loc[common]

    # Design matrix [1, Mkt-RF, SMB, HML] shared across tickers each day.
    X_all = np.column_stack([
        np.ones(len(fac)),
        fac["MKT_RF"].to_numpy(),
        fac["SMB"].to_numpy(),
        fac["HML"].to_numpy(),
    ])
    rf = fac["RF"].to_numpy()
    month_periods = _month_key(common)

    results: dict[pd.Timestamp, dict[str, float]] = {}
    for period, pos in pd.Series(np.arange(len(common)), index=month_periods).groupby(level=0):
        rows = pos.to_numpy()
        if len(rows) < min_days:
            continue
        X = X_all[rows]
        excess = rets.iloc[rows].sub(rf[rows], axis=0)
        month_ts = period.to_timestamp(how="end").normalize()
        ivol_row: dict[str, float] = {}
        # Solve once; loop tickers (small universe keeps this cheap and clear).
        for tkr in rets.columns:
            y = excess[tkr].to_numpy()
            mask = np.isfinite(y)
            if mask.sum() < min_days:
                ivol_row[tkr] = np.nan
                continue
            Xm, ym = X[mask], y[mask]
            beta, *_ = np.linalg.lstsq(Xm, ym, rcond=None)
            resid = ym - Xm @ beta
            ivol_row[tkr] = float(np.std(resid, ddof=1))
        results[month_ts] = ivol_row

    out = pd.DataFrame.from_dict(results, orient="index").sort_index()
    out.index.name = "month"
    return out.reindex(columns=rets.columns)
