set50-system/backend/app/weight_learning.py

"""Factor-weight learning loop (P4).

Learns whether a factor predicts forward returns — and, if so, nudges its weight
up (or down) — from historical PIT data. The deliverable per factor is an
attribution report + a before/after weight:

    new_weight = clip(old_weight * (1 + shrink * ic_mean), min_w, max_w)

Honesty guards:
  - only uses data available at time `t` (point-in-time by construction),
  - cross-sectional IC (Spearman rank correlation) is computed per period and
    aggregated, not fit on the full window (avoids the big look-ahead),
  - a holdout tail is never used to tune weights (caller keeps it out),
  - factors with too few overlapping periods contribute no weight change.

Historical *factor vintages* for macro/demographic factors are not collected yet,
so their learning results are reported as BLOCKED (ic_mean=None) rather than
fabricated. Momentum is the first factor with a genuine PIT historical series
derivable from the existing price archive and is exercised end-to-end.
"""

from __future__ import annotations

import datetime as dt
import math
import statistics
from dataclasses import dataclass, field
from typing import Optional

from .backtest import momentum_at

DEFAULT_SHRINK = 0.5      # how aggressively IC moves the weight
DEFAULT_MIN_W = 0.1
DEFAULT_MAX_W = 3.0
FORWARD_MONTHS = 3        # reward horizon (t -> t+3m forward return)


@dataclass
class FactorLearning:
    factor_key: str
    n_periods: int = 0
    ic_mean: Optional[float] = None
    ic_std: Optional[float] = None
    ic_tstat: Optional[float] = None
    old_weight: Optional[float] = None
    new_weight: Optional[float] = None
    blocked: bool = False   # True when no PIT historical series is available

    def to_dict(self) -> dict:
        return {
            "factor_key": self.factor_key,
            "n_periods": self.n_periods,
            "ic_mean": self.ic_mean,
            "ic_std": self.ic_std,
            "ic_tstat": self.ic_tstat,
            "old_weight": self.old_weight,
            "new_weight": self.new_weight,
            "blocked": self.blocked,
        }


class WeightLearningError(ValueError):
    pass


# ---------------------------------------------------------------------------
# IC helpers
# ---------------------------------------------------------------------------
def _rank(vals: list[float]) -> list[float]:
    """Rank-normalize a list (average ties)."""
    idx = sorted(range(len(vals)), key=lambda i: vals[i])
    ranks = [0.0] * len(vals)
    i = 0
    while i < len(vals):
        j = i
        while j + 1 < len(vals) and vals[idx[j + 1]] == vals[idx[i]]:
            j += 1
        avg = (i + j) / 2.0 + 1.0
        for k in range(i, j + 1):
            ranks[idx[k]] = avg
        i = j + 1
    return ranks


def _corr(a: list[float], b: list[float]) -> float:
    n = len(a)
    if n < 3:
        return 0.0
    ma = statistics.fmean(a)
    mb = statistics.fmean(b)
    cov = sum((a[i] - ma) * (b[i] - mb) for i in range(n))
    va = sum((x - ma) ** 2 for x in a)
    vb = sum((y - mb) ** 2 for y in b)
    if va <= 0 or vb <= 0:
        return 0.0
    return cov / math.sqrt(va * vb)


def spearman_ic(factor_values: dict[str, float],
                forward_returns: dict[str, float]) -> Optional[float]:
    """Cross-sectional Spearman (rank) IC between a factor and forward returns."""
    syms = [s for s in factor_values
            if s in forward_returns
            and math.isfinite(factor_values[s])
            and math.isfinite(forward_returns[s])]
    if len(syms) < 3:
        return None
    fv = [factor_values[s] for s in syms]
    fr = [forward_returns[s] for s in syms]
    return _corr(_rank(fv), _rank(fr))


# ---------------------------------------------------------------------------
# Forward-return construction (price-derived)
# ---------------------------------------------------------------------------
def _forward_return(series: dict, sym: str, t: dt.date, months: int) -> Optional[float]:
    bars = [b for b in series.get(sym, {}).get("bars", [])
            if _bar_date(b.get("date")) <= t]
    if len(bars) < 2:
        return None
    start_px = float(bars[-1]["adjusted_close"])
    horizon = _add_months(t, months)
    future = [b for b in series.get(sym, {}).get("bars", [])
              if _bar_date(b.get("date")) <= horizon]
    if len(future) < 2:
        return None
    end_px = float(future[-1]["adjusted_close"])
    if start_px <= 0:
        return None
    return (end_px / start_px) - 1.0


def _add_months(d: dt.date, months: int) -> dt.date:
    m = d.month - 1 + months
    y = d.year + m // 12
    m = m % 12 + 1
    # clamp day to the last valid day of the target month (e.g. Jan 31 -> Feb 28)
    import calendar
    day = min(d.day, calendar.monthrange(y, m)[1])
    return dt.date(y, m, day)


def _bar_date(s: Optional[str]) -> dt.date:
    if not s:
        return dt.date.min
    return dt.date.fromisoformat(str(s)[:10])


# ---------------------------------------------------------------------------
# Generic factor learning from a historical cross-sectional series (P4)
# ---------------------------------------------------------------------------
def learn_factor_series(period_ics: list[float]) -> FactorLearning:
    """Aggregate a list of per-period cross-sectional ICs into a report.

    `period_ics` is one IC per period (e.g. one per month). Used by factor
    learners that already built the per-period IC; `learn_momentum` is a
    concrete instance that builds period_ics from the price archive.
    """
    res = FactorLearning(factor_key="custom")
    res.n_periods = len(period_ics)
    if period_ics:
        res.ic_mean = round(statistics.fmean(period_ics), 4)
        res.ic_std = round(statistics.pstdev(period_ics), 4)
        res.ic_tstat = round(res.ic_mean / (res.ic_std / math.sqrt(len(period_ics))), 3) \
            if res.ic_std else None
    return res


# ---------------------------------------------------------------------------
# Momentum factor learning (real PIT demo)
# ---------------------------------------------------------------------------
def learn_momentum(series: dict, symbols: list[str], start: str, end: str,
                   step_days: int = 21, forward_months: int = FORWARD_MONTHS) -> FactorLearning:
    """Learn whether 12-1 momentum predicts 3m forward returns, entirely PIT."""
    s = dt.date.fromisoformat(start)
    e = dt.date.fromisoformat(end)
    ic_series: list[float] = []
    cur = s
    while cur <= e:
        fv: dict[str, float] = {}
        fr: dict[str, float] = {}
        for sym in symbols:
            m = momentum_at(series, sym, cur)
            f = _forward_return(series, sym, cur, forward_months)
            if m is not None and f is not None:
                fv[sym] = m
                fr[sym] = f
        ic = spearman_ic(fv, fr)
        if ic is not None:
            ic_series.append(ic)
        cur += dt.timedelta(days=step_days)

    res = learn_factor_series(ic_series)
    res.factor_key = "momentum_12_1"
    return res


# ---------------------------------------------------------------------------
# Weight update
# ---------------------------------------------------------------------------
def apply_weight_update(learning: FactorLearning, shrink: float = DEFAULT_SHRINK,
                        min_w: float = DEFAULT_MIN_W, max_w: float = DEFAULT_MAX_W) -> None:
    """Fold a learned IC into the factor's weight (in place).

    new = clip(old * (1 + shrink * ic_mean), min_w, max_w).
    Blocked / no-data factors keep their weight (new == old).
    """
    if learning.old_weight is None:
        return
    if learning.blocked or learning.ic_mean is None or learning.n_periods < 3:
        learning.new_weight = learning.old_weight
        return
    nw = learning.old_weight * (1.0 + shrink * learning.ic_mean)
    learning.new_weight = round(min(max(nw, min_w), max_w), 4)