scanner/performance.py

"""Auto-evaluate the scoring algorithm and self-tune the factor weights.

After each scan we already persist a snapshot containing every ticker's
factor values, score, and ``last_close`` (see :mod:`scanner.history`).
Combined with the OHLCV cache, those snapshots are enough to compute the
realised forward returns of past predictions and the information
coefficient (IC) - the rank correlation between a score and what actually
happened over the next ``HORIZON_DAYS`` trading days.

:func:`auto_improve` evaluates the *current* weights, then runs a small
random search over alternative weight vectors and picks the one with the
highest IC.  If the candidate beats the baseline by at least
``IC_IMPROVEMENT_THRESHOLD`` the new weights are persisted to
``LEARNED_WEIGHTS_PATH`` and a row is appended to
``PERFORMANCE_LOG_PATH``.  The next scan can pick those weights up via
:func:`load_learned_weights`.

The search is deliberately small (a few hundred candidates) so it runs in
a few seconds in a background thread after each scan.
"""

from __future__ import annotations

import json
import os
from datetime import datetime
from typing import Optional

import numpy as np
import pandas as pd

from .data_fetcher import _cache_load
from .history import list_snapshots
from . import paths
from .scorer import FACTOR_KEYS, robust_zscore

# How many trading days forward to evaluate each snapshot over
HORIZON_DAYS = 5
# Need at least this many snapshots with valid forward returns
MIN_SNAPSHOTS_FOR_TUNING = 5
# Per-snapshot, need at least this many tickers with both factors and returns
MIN_VALID_TICKERS_PER_SNAPSHOT = 30
# Search width
N_RANDOM_CANDIDATES = 200
# Only accept a new weight vector if it beats the baseline IC by this margin
IC_IMPROVEMENT_THRESHOLD = 0.005


# ---------------------------------------------------------------------------
# Per-snapshot evaluation table
# ---------------------------------------------------------------------------

def _entry_exit_close(frame: pd.DataFrame, snap_date: datetime,
                      horizon: int) -> Optional[tuple[float, float]]:
    """Look up entry/exit close prices for a ticker around ``snap_date``.

    ``entry`` = close on first trading day on or after ``snap_date``.
    ``exit``  = close ``horizon`` trading days later.

    Returns ``None`` if the cache does not extend far enough.
    """
    if frame is None or frame.empty or "Close" not in frame.columns:
        return None
    if "Date" in frame.columns:
        dates = pd.to_datetime(frame["Date"]).values
    else:
        dates = pd.to_datetime(frame.index).values
    closes = frame["Close"].astype(float).values
    target = np.datetime64(snap_date.date())
    where = np.where(dates >= target)[0]
    if len(where) == 0:
        return None
    i0 = int(where[0])
    i1 = i0 + int(horizon)
    if i1 >= len(closes):
        return None
    p0, p1 = float(closes[i0]), float(closes[i1])
    if not (np.isfinite(p0) and np.isfinite(p1)) or p0 <= 0:
        return None
    return p0, p1


def _build_eval_table(snap_ts: datetime, snap_df: pd.DataFrame,
                      cache: dict[str, pd.DataFrame], horizon: int) -> pd.DataFrame:
    """For a single snapshot, build a frame of forward returns and per-snapshot
    z-scored factor values.  Empty frame if the snapshot is too thin.
    """
    if snap_df is None or snap_df.empty:
        return pd.DataFrame()
    needed = ["ticker"] + FACTOR_KEYS
    if not all(c in snap_df.columns for c in needed):
        return pd.DataFrame()
    work = snap_df[needed].dropna(subset=FACTOR_KEYS).copy()
    work["ticker"] = work["ticker"].astype(str)

    rets: dict[str, float] = {}
    for t in work["ticker"]:
        pair = _entry_exit_close(cache.get(t), snap_ts, horizon)
        if pair is None:
            continue
        rets[t] = (pair[1] / pair[0]) - 1.0
    if not rets:
        return pd.DataFrame()

    work["fwd_ret"] = work["ticker"].map(rets)
    work = work.dropna(subset=["fwd_ret"])
    if len(work) < MIN_VALID_TICKERS_PER_SNAPSHOT:
        return pd.DataFrame()
    # Per-snapshot z-score
    for k in FACTOR_KEYS:
        work[f"z_{k}"] = robust_zscore(work[k])
    return work[["ticker", "fwd_ret"] + [f"z_{k}" for k in FACTOR_KEYS]]


def collect_eval_tables(horizon: int = HORIZON_DAYS,
                        cache: Optional[dict[str, pd.DataFrame]] = None
                        ) -> list[pd.DataFrame]:
    """Return one evaluation table per snapshot that has enough forward data."""
    cache = cache if cache is not None else _cache_load()
    if not cache:
        return []
    tables: list[pd.DataFrame] = []
    for ts, path in list_snapshots():
        try:
            snap_df = pd.read_parquet(path)
        except Exception:
            continue
        tbl = _build_eval_table(ts, snap_df, cache, horizon)
        if not tbl.empty:
            tables.append(tbl)
    return tables


# ---------------------------------------------------------------------------
# Metrics & search
# ---------------------------------------------------------------------------

def _metrics_for_weights(tables: list[pd.DataFrame],
                         weights: dict) -> tuple[float, float, int]:
    """Aggregate (mean Spearman IC, hit rate, n_periods) for a weight dict."""
    ics, hits = [], []
    for tbl in tables:
        score = sum(tbl[f"z_{k}"] * float(weights.get(k, 0.0)) for k in FACTOR_KEYS)
        if not np.isfinite(score).any() or score.abs().sum() == 0:
            continue
        ic = score.corr(tbl["fwd_ret"], method="spearman")
        if ic is not None and np.isfinite(ic):
            ics.append(float(ic))
        hit = float((np.sign(score) == np.sign(tbl["fwd_ret"])).mean())
        if np.isfinite(hit):
            hits.append(hit)
    if not ics:
        return float("nan"), float("nan"), 0
    return float(np.mean(ics)), (float(np.mean(hits)) if hits else float("nan")), len(ics)


def evaluate(weights: dict, horizon: int = HORIZON_DAYS) -> dict:
    """Public: metrics for *weights* on all available history."""
    tables = collect_eval_tables(horizon)
    ic, hit, n = _metrics_for_weights(tables, weights)
    return {"mean_ic": ic, "hit_rate": hit, "n_periods": n, "horizon_days": horizon}


def _normalize(vec: np.ndarray) -> np.ndarray:
    s = float(vec.sum())
    if s <= 1e-9:
        return np.ones_like(vec) / len(vec)
    return vec / s


def optimize_weights(
    base_weights: dict,
    horizon: int = HORIZON_DAYS,
    n_random: int = N_RANDOM_CANDIDATES,
    seed: int = 13,
    tables: Optional[list[pd.DataFrame]] = None,
) -> tuple[dict, dict]:
    """Random search over Dirichlet samples + a few structured candidates.

    Returns ``(weights, metrics)``.  ``metrics["tuned"]`` is True only if
    the search found weights that beat the baseline by
    :data:`IC_IMPROVEMENT_THRESHOLD`.
    """
    if tables is None:
        tables = collect_eval_tables(horizon)

    base_ic, base_hit, base_n = _metrics_for_weights(tables, base_weights)
    base_metrics = {
        "mean_ic": base_ic, "hit_rate": base_hit, "n_periods": base_n,
        "horizon_days": horizon, "tuned": False, "baseline_ic": base_ic,
        "ic_gain": 0.0,
    }
    if base_n < MIN_SNAPSHOTS_FOR_TUNING:
        return dict(base_weights), base_metrics

    rng = np.random.default_rng(seed)
    base_vec = _normalize(np.array([base_weights.get(k, 0.0) for k in FACTOR_KEYS],
                                   dtype=float))
    # Structured candidates: current, uniform, one-hot-heavy
    candidates: list[np.ndarray] = [base_vec, np.ones(len(FACTOR_KEYS)) / len(FACTOR_KEYS)]
    for i in range(len(FACTOR_KEYS)):
        v = np.full(len(FACTOR_KEYS), 0.05)
        v[i] = 0.8
        candidates.append(_normalize(v))
    # Broad Dirichlet (explore)
    candidates.extend(rng.dirichlet(np.ones(len(FACTOR_KEYS)) * 1.0,
                                    size=max(1, n_random // 2)))
    # Narrow Dirichlet around current (exploit)
    alpha = base_vec * 10.0 + 0.5
    candidates.extend(rng.dirichlet(alpha, size=max(1, n_random // 2)))

    best_vec = base_vec
    best_ic = base_ic if np.isfinite(base_ic) else -np.inf
    best_hit = base_hit
    best_n = base_n
    for vec in candidates:
        vec = _normalize(np.asarray(vec, dtype=float))
        cand = {k: float(vec[i]) for i, k in enumerate(FACTOR_KEYS)}
        ic, hit, n = _metrics_for_weights(tables, cand)
        if not np.isfinite(ic) or n < MIN_SNAPSHOTS_FOR_TUNING:
            continue
        if ic > best_ic:
            best_ic, best_hit, best_n, best_vec = ic, hit, n, vec

    improved = (np.isfinite(base_ic)
                and (best_ic - base_ic) >= IC_IMPROVEMENT_THRESHOLD)
    if not improved:
        return dict(base_weights), base_metrics

    learned = {k: float(best_vec[i]) for i, k in enumerate(FACTOR_KEYS)}
    metrics = {
        "mean_ic": float(best_ic),
        "hit_rate": float(best_hit) if np.isfinite(best_hit) else float("nan"),
        "n_periods": int(best_n),
        "horizon_days": horizon,
        "tuned": True,
        "baseline_ic": float(base_ic),
        "ic_gain": float(best_ic - base_ic),
    }
    return learned, metrics


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------

def load_learned_weights() -> Optional[dict]:
    """Read the most recently saved learned weights, or None if absent."""
    if not os.path.exists(paths.LEARNED_WEIGHTS_PATH):
        return None
    try:
        with open(paths.LEARNED_WEIGHTS_PATH, "r", encoding="utf-8") as fh:
            data = json.load(fh)
        if isinstance(data, dict) and all(k in data for k in FACTOR_KEYS):
            return {k: float(data[k]) for k in FACTOR_KEYS}
    except Exception:
        return None
    return None


def load_learned_meta() -> dict:
    """Return the saved metrics blob alongside learned weights, or empty dict."""
    if not os.path.exists(paths.LEARNED_WEIGHTS_PATH):
        return {}
    try:
        with open(paths.LEARNED_WEIGHTS_PATH, "r", encoding="utf-8") as fh:
            data = json.load(fh)
        out = {}
        if isinstance(data, dict):
            out["weights"] = {k: float(data[k]) for k in FACTOR_KEYS if k in data}
            out["metrics"] = data.get("_metrics", {})
            out["saved_at"] = data.get("_saved_at")
        return out
    except Exception:
        return {}


def save_learned_weights(weights: dict, metrics: Optional[dict] = None) -> bool:
    try:
        os.makedirs(os.path.dirname(paths.LEARNED_WEIGHTS_PATH) or ".", exist_ok=True)
        payload: dict = {k: float(weights.get(k, 0.0)) for k in FACTOR_KEYS}
        payload["_metrics"] = metrics or {}
        payload["_saved_at"] = datetime.utcnow().isoformat()
        with open(paths.LEARNED_WEIGHTS_PATH, "w", encoding="utf-8") as fh:
            json.dump(payload, fh, indent=2)
        return True
    except Exception:
        return False


def append_performance_log(metrics: dict, weights: dict) -> None:
    row = {**metrics,
           "saved_at": datetime.utcnow().isoformat(),
           **{f"w_{k}": float(weights.get(k, 0.0)) for k in FACTOR_KEYS}}
    df_new = pd.DataFrame([row])
    try:
        os.makedirs(os.path.dirname(paths.PERFORMANCE_LOG_PATH) or ".", exist_ok=True)
        if os.path.exists(paths.PERFORMANCE_LOG_PATH):
            try:
                old = pd.read_parquet(paths.PERFORMANCE_LOG_PATH)
                df_new = pd.concat([old, df_new], ignore_index=True).tail(500)
            except Exception:
                pass
        df_new.to_parquet(paths.PERFORMANCE_LOG_PATH, index=False)
    except Exception:
        pass


def load_performance_log() -> pd.DataFrame:
    if not os.path.exists(paths.PERFORMANCE_LOG_PATH):
        return pd.DataFrame()
    try:
        return pd.read_parquet(paths.PERFORMANCE_LOG_PATH)
    except Exception:
        return pd.DataFrame()


# ---------------------------------------------------------------------------
# Entry point used by app
# ---------------------------------------------------------------------------

def auto_improve(base_weights: dict, horizon: int = HORIZON_DAYS) -> dict:
    """End-to-end: evaluate baseline, search, persist if improved, log row.

    Returns a dict with ``weights`` (the persisted set: either learned or
    baseline) and ``metrics`` (the metrics dict from the search).  Safe to
    call from a background thread; never raises.
    """
    try:
        tables = collect_eval_tables(horizon)
        if len(tables) < MIN_SNAPSHOTS_FOR_TUNING:
            metrics = {"mean_ic": float("nan"), "hit_rate": float("nan"),
                       "n_periods": len(tables), "horizon_days": horizon,
                       "tuned": False,
                       "note": f"Need at least {MIN_SNAPSHOTS_FOR_TUNING} snapshots with forward data."}
            append_performance_log(metrics, base_weights)
            return {"weights": dict(base_weights), "metrics": metrics}

        learned, metrics = optimize_weights(base_weights, horizon=horizon,
                                            tables=tables)
        if metrics.get("tuned"):
            save_learned_weights(learned, metrics)
        append_performance_log(metrics,
                               learned if metrics.get("tuned") else base_weights)
        return {"weights": learned, "metrics": metrics}
    except Exception as exc:  # never let background work crash the app
        return {"weights": dict(base_weights),
                "metrics": {"error": str(exc), "tuned": False,
                            "horizon_days": horizon, "n_periods": 0,
                            "mean_ic": float("nan"), "hit_rate": float("nan")}}