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
    ic_train: Optional[float] = None   # walk-forward train-window mean IC
    ic_holdout: Optional[float] = None  # holdout-window mean IC (out-of-sample)
    validated: bool = False  # True only when the gate fully passes
    gate_notes: list = field(default_factory=list)  # why not validated

    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,
            "ic_train": self.ic_train,
            "ic_holdout": self.ic_holdout,
            "old_weight": self.old_weight,
            "new_weight": self.new_weight,
            "blocked": self.blocked,
            "validated": self.validated,
            "gate_notes": self.gate_notes,
        }


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


# Default gate thresholds (honest, conservative — from the user's P4 guardrail:
# no auto-apply until minimum sample + holdout/walk-forward + baseline pass).
MIN_SAMPLE_PERIODS = 12      # at least a year of monthly periods for any learn
MIN_HOLDOUT_PERIODS = 6      # holdout window must be meaningful, not trivial
MIN_TRAIN_PERIODS = 6        # train window must be non-trivial
# A factor "beats baseline (random)" when both train and holdout IC are
# positive with a t-stat magnitude beyond this (else it is indistinguishable
# from noise and must not move weights).
MIN_IC_TSTAT = 1.0
BASELINE_IC = 0.0            # naive benchmark: no predictive power (IC = 0)


def apply_validation_gate(learning: FactorLearning,
                          train_ics: list[float],
                          holdout_ics: list[float]) -> FactorLearning:
    """Apply the strict holdout/walk-forward + baseline gate.

    Sets ``ic_train`` / ``ic_holdout`` and decides ``validated``. A factor is
    only ``validated=True`` when ALL of:

      - total periods >= MIN_SAMPLE_PERIODS
      - train and holdout windows each >= their minimums
      - train IC and holdout IC agree in sign AND are positive (beat baseline)
      - the pooled IC is significant enough (|t| > MIN_IC_TSTAT)

    This is intentionally strict: unvalidated factors keep their weight (the
    caller must never auto-apply).
    """
    total = len(train_ics) + len(holdout_ics)
    notes: list[str] = []
    if total < MIN_SAMPLE_PERIODS:
        notes.append(f"sample too small ({total} < {MIN_SAMPLE_PERIODS})")

    def _mean(xs):
        return round(statistics.fmean(xs), 4) if xs else None
    learning.ic_train = _mean(train_ics)
    learning.ic_holdout = _mean(holdout_ics)

    validated = True
    if len(train_ics) < MIN_TRAIN_PERIODS:
        notes.append(f"train too small ({len(train_ics)} < {MIN_TRAIN_PERIODS})")
        validated = False
    if len(holdout_ics) < MIN_HOLDOUT_PERIODS:
        notes.append(f"holdout too small ({len(holdout_ics)} < {MIN_HOLDOUT_PERIODS})")
        validated = False
    # train & holdout must both be positive (beat baseline) and agree in sign
    if validated:
        if not train_ics or not holdout_ics:
            notes.append("missing train or holdout IC")
            validated = False
        elif learning.ic_train <= BASELINE_IC or learning.ic_holdout <= BASELINE_IC:
            notes.append(f"IC not above baseline ({learning.ic_train}/{learning.ic_holdout})")
            validated = False
        elif (learning.ic_train < 0) != (learning.ic_holdout < 0):
            notes.append("train and holdout IC disagree in sign")
            validated = False
    # significance: pooled t-stat must clear the bar
    pooled = (train_ics + holdout_ics) if validated else []
    if validated and pooled:
        m = statistics.fmean(pooled)
        sd = statistics.pstdev(pooled)
        t = (m / (sd / math.sqrt(len(pooled)))) if sd else None
        if t is None or abs(t) <= MIN_IC_TSTAT:
            notes.append(f"IC not significant (|t|={t} <= {MIN_IC_TSTAT})")
            validated = False

    learning.validated = validated
    learning.gate_notes = notes
    return learning


# ---------------------------------------------------------------------------
# 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).

    **Never auto-applies an unvalidated factor.** A weight only moves when the
    factor is ``validated`` (passed minimum sample + holdout/walk-forward +
    baseline via ``apply_validation_gate``). Blocked, no-data, insufficient, or
    non-significant factors keep their weight (new == old) and stay
    ``validated=False``.
    """
    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
    if not learning.validated:
        # gate not passed -> do NOT move the weight (honest, no auto-apply)
        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)


def learn_momentum_gated(series: dict, symbols: list[str], start: str, end: str,
                         step_days: int = 21, forward_months: int = FORWARD_MONTHS,
                         holdout_fraction: float = 0.3) -> FactorLearning:
    """Learn 12-1 momentum with a strict holdout validation gate.

    Builds per-period PIT ICs across [start, end], then splits chronologically
    into train (first 1-holdout_fraction) and holdout (last fraction) windows
    and applies ``apply_validation_gate``. Returns a FactorLearning whose
    ``validated`` is True only if the factor clears the gate; unvalidated runs
    keep their weight (the caller must not auto-apply).
    """
    s = dt.date.fromisoformat(start)
    e = dt.date.fromisoformat(end)
    all_ics: 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:
            all_ics.append(ic)
        cur += dt.timedelta(days=step_days)

    res = learn_factor_series(all_ics)
    res.factor_key = "momentum_12_1"

    # chronologically split train / holdout (walk-forward style)
    holdout_n = max(0, int(round(len(all_ics) * holdout_fraction)))
    if holdout_n > 0:
        split = len(all_ics) - holdout_n
        train_ics = all_ics[:split]
        holdout_ics = all_ics[split:]
        res = apply_validation_gate(res, train_ics, holdout_ics)
    else:
        res.gate_notes.append("no holdout window (too few periods)")
        res.validated = False
    return res