from tools.train_mlp import (
    aggregate_fold_score,
    robust_objective_score,
    threshold_fragility_count,
)
from tools.worker_utils import resolve_worker_count


def test_aggregate_fold_score_modes():
    calmars = [1.0, 0.0, 3.0]

    assert aggregate_fold_score(calmars, "mean") == 4.0 / 3.0
    assert aggregate_fold_score(calmars, "min") == 0.0
    assert aggregate_fold_score(calmars, "mean_min") == 0.7 * (4.0 / 3.0)


def test_threshold_fragility_count_uses_nearest_threshold():
    scores = [-1.0, -0.01, 0.04, 0.9, 2.0]
    thresholds = [0.0, 1.0, 10.0]

    assert threshold_fragility_count(scores, thresholds, 0.05) == 2
    assert threshold_fragility_count(scores, thresholds, 0.11) == 3


def test_robust_objective_penalizes_drawdown_pnl_dd_trades_and_fragility():
    calmars = [1.0, 0.5]
    clean_metrics = [
        {"Total Trades": 20, "Max Drawdown %": -30.0, "P&L/DD Ratio": 2.0, "Calmar Ratio": 1.0},
        {"Total Trades": 20, "Max Drawdown %": -35.0, "P&L/DD Ratio": 1.0, "Calmar Ratio": 0.5},
    ]
    weak_metrics = [
        {"Total Trades": 5, "Max Drawdown %": -70.0, "P&L/DD Ratio": 0.0, "Calmar Ratio": -0.4},
        {"Total Trades": 8, "Max Drawdown %": -60.0, "P&L/DD Ratio": 0.0, "Calmar Ratio": 0.5},
    ]

    clean_score, clean_penalties = robust_objective_score(
        calmars, clean_metrics, [10.0, 20.0, 30.0], [0.0, 1.0, 2.0, 0.0], 10
    )
    weak_score, weak_penalties = robust_objective_score(
        calmars, weak_metrics, [-0.01, 0.5, 0.99], [0.0, 1.0, 2.0, 0.0], 10
    )

    assert weak_score < clean_score
    assert clean_penalties["pnl_dd_penalty"] == 0.0
    assert weak_penalties["pnl_dd_penalty"] == 1.0
    assert weak_penalties["drawdown_penalty"] > clean_penalties["drawdown_penalty"]
    assert weak_penalties["trade_penalty"] > clean_penalties["trade_penalty"]
    assert weak_penalties["fragility_penalty"] > clean_penalties["fragility_penalty"]


def test_resolve_worker_count_auto_uses_fraction_and_workload_cap():
    assert resolve_worker_count(0, workload_size=100, worker_fraction=0.5, cpu_count=16) == 8
    assert resolve_worker_count(0, workload_size=3, worker_fraction=0.5, cpu_count=16) == 3
    assert resolve_worker_count(0, workload_size=100, worker_fraction=0.25, cpu_count=16) == 4


def test_resolve_worker_count_manual_overrides_auto_and_caps_to_workload():
    assert resolve_worker_count(6, workload_size=100, cpu_count=16) == 6
    assert resolve_worker_count(6, workload_size=4, cpu_count=16) == 4
    assert resolve_worker_count(1, workload_size=100, cpu_count=16) == 1