tests/test_factor_weight_optimization.py

import pandas as pd

from scripts.optimize_factor_weights import (
    WeightConfig,
    apply_weight_config,
    compute_event_metrics,
    load_stock_codes_file,
    optimize_weights,
    pass_decision,
)


def _events(rows):
    return pd.DataFrame(rows)


def test_apply_weight_config_can_flip_to_buy_with_bull_factor():
    events = _events(
        [
            {
                "stock": "2330",
                "date": "2025-04-01",
                "y_true": 1,
                "base_pred": 0,
                "p_buy": 0.42,
                "p_hold": 0.46,
                "p_sell": 0.12,
                "bull_score": 3.0,
                "bear_score": 0.0,
            }
        ]
    )

    pred = apply_weight_config(events, WeightConfig(buy_bull_weight=0.02))

    assert pred.tolist() == [1]


def test_apply_weight_config_bear_gate_blocks_risky_buy():
    events = _events(
        [
            {
                "stock": "2330",
                "date": "2025-04-01",
                "y_true": 0,
                "base_pred": 1,
                "p_buy": 0.62,
                "p_hold": 0.31,
                "p_sell": 0.07,
                "bull_score": 0.0,
                "bear_score": 2.0,
            }
        ]
    )

    pred = apply_weight_config(events, WeightConfig(buy_bear_gate=2.0))

    assert pred.tolist() == [0]


def test_compute_event_metrics_reports_precision_and_coverage():
    events = _events(
        [
            {"y_true": 1},
            {"y_true": 0},
            {"y_true": -1},
            {"y_true": 1},
        ]
    )

    metrics = compute_event_metrics(events, pred=pd.Series([1, 0, -1, -1]).to_numpy())

    assert metrics["accuracy"] == 75.0
    assert metrics["direction_accuracy"] == 66.6667
    assert metrics["buy_precision"] == 100.0
    assert metrics["sell_precision"] == 50.0
    assert metrics["coverage"] == 0.75


def test_optimize_weights_uses_earlier_events_then_validates_later_events():
    rows = []
    for i in range(6):
        risky_false_buy = i < 4
        rows.append(
            {
                "stock": "2330",
                "date": f"2025-04-0{i + 1}",
                "y_true": 0 if risky_false_buy else 1,
                "base_pred": 1,
                "p_buy": 0.62 if risky_false_buy else 0.58,
                "p_hold": 0.31,
                "p_sell": 0.07 if risky_false_buy else 0.11,
                "bull_score": 0.0 if risky_false_buy else 3.0,
                "bear_score": 2.0 if risky_false_buy else 0.0,
            }
        )
    for i in range(4):
        risky_false_buy = i < 2
        rows.append(
            {
                "stock": "2330",
                "date": f"2025-04-1{i + 1}",
                "y_true": 0 if risky_false_buy else 1,
                "base_pred": 1,
                "p_buy": 0.62 if risky_false_buy else 0.58,
                "p_hold": 0.31,
                "p_sell": 0.07 if risky_false_buy else 0.11,
                "bull_score": 0.0 if risky_false_buy else 3.0,
                "bear_score": 2.0 if risky_false_buy else 0.0,
            }
        )

    result = optimize_weights(
        _events(rows),
        optimize_ratio=0.6,
        min_signal_ratio=0.3,
        grid={
            "buy_bear_gate": [None, 2.0],
            "buy_bull_weight": [0.0],
            "buy_bear_weight": [0.0],
            "sell_bear_weight": [0.0],
            "sell_bull_weight": [0.0],
            "sell_bull_gate": [None],
            "hold_bias": [0.0],
        },
    )

    assert result["best_config"]["buy_bear_gate"] == 2.0
    assert result["split"]["tune_events"] == 6
    assert result["split"]["validation_events"] == 4
    assert result["validation"]["deltas"]["accuracy_delta_pp"] == 50.0
    assert result["validation"]["deltas"]["buy_precision_delta_pp"] == 50.0


def test_pass_decision_requires_validation_improvement():
    result = {
        "validation": {
            "baseline": {"signal_count": 4},
            "candidate": {"signal_count": 4},
            "deltas": {"accuracy_delta_pp": 3.0, "buy_precision_delta_pp": 0.0},
        }
    }

    decision = pass_decision(result, min_accuracy_delta_pp=3.0, min_signal_ratio=0.6)

    assert decision["passed"] is True


def test_precision_first_rejects_buy_count_expansion_and_prefers_precision():
    rows = []
    for i in range(10):
        rows.append(
            {
                "stock": "2330",
                "date": f"2025-04-{i + 1:02d}",
                "y_true": 1 if i in {0, 1, 2, 3} else 0,
                "base_pred": 1 if i in {0, 1} else 0,
                "p_buy": 0.52 if i in {0, 1} else 0.48,
                "p_hold": 0.50 if i not in {0, 1} else 0.44,
                "p_sell": 0.02,
                "bull_score": 3.0 if i in {0, 1, 2, 3, 4, 5} else 0.0,
                "bear_score": 2.0 if i in {4, 5} else 0.0,
            }
        )

    result = optimize_weights(
        _events(rows),
        optimize_ratio=0.6,
        min_signal_ratio=0.3,
        objective="precision-first",
        max_buy_count_ratio=1.3,
        grid={
            "buy_bear_gate": [None, 2.0],
            "buy_bull_weight": [0.0, 0.02],
            "buy_bear_weight": [0.0],
            "sell_bear_weight": [0.0],
            "sell_bull_weight": [0.0],
            "sell_bull_gate": [None],
            "hold_bias": [0.0],
        },
    )

    assert result["objective"] == "precision-first"
    assert result["max_buy_count_ratio"] == 1.3
    assert result["best_config"]["buy_bull_weight"] == 0.0


def test_pass_decision_can_enforce_buy_count_ratio():
    result = {
        "validation": {
            "baseline": {"signal_count": 10, "buy_count": 10},
            "candidate": {"signal_count": 10, "buy_count": 14},
            "deltas": {"accuracy_delta_pp": 5.0, "buy_precision_delta_pp": 0.0},
        }
    }

    decision = pass_decision(
        result,
        min_accuracy_delta_pp=5.0,
        min_signal_ratio=0.7,
        max_buy_count_ratio=1.3,
    )

    assert decision["passed"] is False
    assert decision["checks"]["buy_count_ratio"] is False
    assert decision["buy_count_ratio"] == 1.4


def test_load_stock_codes_file_accepts_popular_stock_payload(tmp_path):
    path = tmp_path / "popular.json"
    path.write_text(
        """
        {
          "codes": ["2330", "0050", "2330"],
          "stocks": [{"code": "ignored"}]
        }
        """
    )

    assert load_stock_codes_file(path) == ["2330", "0050"]