tests/test_performance.py

"""Tests for the auto-tuning module."""

from __future__ import annotations

import os
from datetime import datetime, timedelta

import numpy as np
import pandas as pd
import pytest

from scanner.data_fetcher import _cache_save
from scanner.history import save_snapshot
from scanner.performance import (
    HORIZON_DAYS, MIN_SNAPSHOTS_FOR_TUNING, _metrics_for_weights,
    append_performance_log, auto_improve, collect_eval_tables, evaluate,
    load_learned_meta, load_learned_weights, load_performance_log,
    optimize_weights, save_learned_weights,
)
from scanner.scorer import DEFAULT_WEIGHTS, FACTOR_KEYS


# ---------------------------------------------------------------------------
# Helpers to build synthetic snapshots + matching OHLCV cache
# ---------------------------------------------------------------------------

def _build_cache_for_targets(snapshot_specs: list[tuple[datetime, dict[str, float]]],
                             horizon: int) -> dict[str, pd.DataFrame]:
    """Build an OHLCV cache where each (snap_date, ticker) realises the
    specified forward return after ``horizon`` business days.
    """
    earliest = min(ts for ts, _ in snapshot_specs) - timedelta(days=10)
    latest = max(ts for ts, _ in snapshot_specs) + timedelta(days=horizon * 3 + 10)
    dates = pd.bdate_range(earliest.date(), latest.date()).tolist()
    all_tickers: set[str] = set()
    for _, m in snapshot_specs:
        all_tickers.update(m.keys())

    cache: dict[str, pd.DataFrame] = {}
    for t in all_tickers:
        closes = np.full(len(dates), 100.0)
        for ts, target_map in snapshot_specs:
            if t not in target_map:
                continue
            target_ret = float(target_map[t])
            # find first index where date >= ts
            try:
                entry_idx = next(i for i, d in enumerate(dates)
                                 if d.date() >= ts.date())
            except StopIteration:
                continue
            exit_idx = entry_idx + horizon
            if exit_idx >= len(closes):
                continue
            closes[entry_idx] = 100.0
            closes[exit_idx] = 100.0 * (1.0 + target_ret)
        df = pd.DataFrame({
            "Date": dates,
            "Open": closes, "High": closes * 1.01,
            "Low": closes * 0.99, "Close": closes,
            "Volume": np.full(len(dates), 1_000_000, dtype=int),
        })
        cache[t] = df
    return cache


def _make_snapshots(n_snaps: int = 8, n_tickers: int = 60, horizon: int = 3,
                    signal_factor: str = "cmf", signal_strength: float = 0.04,
                    noise: float = 0.005, seed: int = 7
                    ) -> tuple[list[tuple[datetime, pd.DataFrame]],
                               list[tuple[datetime, dict[str, float]]]]:
    """Synthesise ``n_snaps`` snapshots with random factor values.  Forward
    returns are driven primarily by ``signal_factor`` so the optimal weights
    concentrate on it.
    """
    rng = np.random.default_rng(seed)
    snapshots: list[tuple[datetime, pd.DataFrame]] = []
    targets: list[tuple[datetime, dict[str, float]]] = []
    base_ts = datetime(2025, 11, 1, 12, 0, 0)
    for i in range(n_snaps):
        ts = base_ts + timedelta(days=i * 2)
        tickers = [f"S{i:02d}T{j:03d}" for j in range(n_tickers)]
        factor_data = {k: rng.normal(0, 1, n_tickers) for k in FACTOR_KEYS}
        df = pd.DataFrame({"ticker": tickers, **factor_data})
        # Also include a dummy ``score`` column like the real snapshot would
        df["score"] = df[signal_factor] * 10
        snapshots.append((ts, df))
        # Forward returns: proportional to signal factor + noise
        signal_vals = factor_data[signal_factor]
        rets = signal_strength * signal_vals + rng.normal(0, noise, n_tickers)
        targets.append((ts, dict(zip(tickers, rets))))
    return snapshots, targets


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

def test_collect_eval_tables_with_explicit_cache():
    snapshots, targets = _make_snapshots(n_snaps=4, n_tickers=60, horizon=3)
    cache = _build_cache_for_targets(targets, horizon=3)
    for ts, df in snapshots:
        save_snapshot(df, ts=ts)
    tables = collect_eval_tables(horizon=3, cache=cache)
    assert len(tables) >= 3
    # every table must have finite fwd_ret values
    for t in tables:
        assert t["fwd_ret"].notna().all()
        assert t["fwd_ret"].abs().sum() > 0
    # each table should have fwd_ret and z_ columns
    for t in tables:
        assert "fwd_ret" in t.columns
        for k in FACTOR_KEYS:
            assert f"z_{k}" in t.columns


def test_ic_signal_factor_is_predictive():
    snapshots, targets = _make_snapshots(n_snaps=6, n_tickers=80, horizon=3,
                                         signal_factor="cmf",
                                         signal_strength=0.05, noise=0.003)
    cache = _build_cache_for_targets(targets, horizon=3)
    for ts, df in snapshots:
        save_snapshot(df, ts=ts)
    tables = collect_eval_tables(horizon=3, cache=cache)

    # Pure-CMF weights should have high IC
    cmf_only = {k: (1.0 if k == "cmf" else 0.0) for k in FACTOR_KEYS}
    ic_cmf, _, n_cmf = _metrics_for_weights(tables, cmf_only)
    assert n_cmf >= 5
    assert ic_cmf > 0.5, f"signal factor IC should be high, got {ic_cmf}"

    # Equal weights should have lower IC (it dilutes the signal)
    equal = {k: 0.2 for k in FACTOR_KEYS}
    ic_equal, _, _ = _metrics_for_weights(tables, equal)
    assert ic_cmf > ic_equal


def test_optimize_finds_signal_factor():
    snapshots, targets = _make_snapshots(n_snaps=8, n_tickers=100, horizon=3,
                                         signal_factor="obv_slope",
                                         signal_strength=0.05, noise=0.003,
                                         seed=42)
    cache = _build_cache_for_targets(targets, horizon=3)
    for ts, df in snapshots:
        save_snapshot(df, ts=ts)
    tables = collect_eval_tables(horizon=3, cache=cache)
    assert len(tables) >= MIN_SNAPSHOTS_FOR_TUNING

    base = dict(DEFAULT_WEIGHTS)
    base_ic, _, _ = _metrics_for_weights(tables, base)
    learned, metrics = optimize_weights(base, horizon=3, tables=tables,
                                        n_random=200, seed=11)

    # Learned weights must sum to 1
    assert abs(sum(learned.values()) - 1.0) < 1e-6
    # The optimizer should have improved IC and marked it as tuned
    assert metrics["mean_ic"] >= base_ic
    if metrics["tuned"]:
        assert metrics["mean_ic"] > base_ic
        # Most weight should land on the true signal factor
        sorted_keys = sorted(learned.items(), key=lambda kv: kv[1], reverse=True)
        top_factor = sorted_keys[0][0]
        # With a small random search the absolute top might miss; allow top-2
        assert top_factor in {"obv_slope"} or sorted_keys[1][0] == "obv_slope"


def test_optimize_without_enough_history_returns_baseline():
    # Only 2 snapshots → below MIN_SNAPSHOTS_FOR_TUNING
    snapshots, targets = _make_snapshots(n_snaps=2, n_tickers=40, horizon=3)
    cache = _build_cache_for_targets(targets, horizon=3)
    for ts, df in snapshots:
        save_snapshot(df, ts=ts)
    tables = collect_eval_tables(horizon=3, cache=cache)
    learned, metrics = optimize_weights(DEFAULT_WEIGHTS, horizon=3, tables=tables)
    assert learned == dict(DEFAULT_WEIGHTS)
    assert metrics["tuned"] is False


def test_load_learned_weights_round_trip():
    w = {"cmf": 0.5, "obv_slope": 0.2, "big_bar_ratio": 0.1,
         "vwap_dev": 0.1, "rvol_signed": 0.1}
    assert save_learned_weights(w, metrics={"mean_ic": 0.123})
    loaded = load_learned_weights()
    assert loaded is not None
    # Check the keys we actually saved round-trip; new factor keys
    # added to FACTOR_KEYS later are not required to be present.
    for k in w:
        assert loaded[k] == pytest.approx(w[k])
    meta = load_learned_meta()
    assert meta["metrics"]["mean_ic"] == pytest.approx(0.123)


def test_load_learned_weights_missing():
    assert load_learned_weights() is None
    assert load_learned_meta() == {}


def test_performance_log_grows():
    append_performance_log({"mean_ic": 0.1, "n_periods": 5}, DEFAULT_WEIGHTS)
    append_performance_log({"mean_ic": 0.2, "n_periods": 6}, DEFAULT_WEIGHTS)
    log = load_performance_log()
    assert len(log) == 2
    assert "mean_ic" in log.columns
    assert "w_cmf" in log.columns


def test_auto_improve_end_to_end_with_signal():
    # Persist a cache so collect_eval_tables(_cache_load()) returns data
    snapshots, targets = _make_snapshots(n_snaps=8, n_tickers=80, horizon=3,
                                         signal_factor="vwap_dev",
                                         signal_strength=0.05, noise=0.003,
                                         seed=99)
    cache = _build_cache_for_targets(targets, horizon=3)
    _cache_save(cache)
    for ts, df in snapshots:
        save_snapshot(df, ts=ts)

    result = auto_improve(DEFAULT_WEIGHTS, horizon=3)
    assert "weights" in result and "metrics" in result
    assert isinstance(result["weights"], dict)
    log = load_performance_log()
    assert len(log) == 1


def test_auto_improve_handles_no_data():
    # No snapshots, no cache
    result = auto_improve(DEFAULT_WEIGHTS, horizon=3)
    assert result["metrics"]["tuned"] is False
    assert result["weights"] == dict(DEFAULT_WEIGHTS)


def test_evaluate_returns_metrics_dict():
    snapshots, targets = _make_snapshots(n_snaps=6, n_tickers=50, horizon=3,
                                         seed=5)
    cache = _build_cache_for_targets(targets, horizon=3)
    _cache_save(cache)
    for ts, df in snapshots:
        save_snapshot(df, ts=ts)
    metrics = evaluate(DEFAULT_WEIGHTS, horizon=3)
    for k in ("mean_ic", "hit_rate", "n_periods", "horizon_days"):
        assert k in metrics
    assert metrics["horizon_days"] == 3