"""
Tests for parallel episode evaluation — correctness, determinism, speedup.

Design principle: parallel evaluation must produce IDENTICAL results to sequential
for the same seed. Fidelity is non-negotiable.
"""
import pytest
import pickle
from malkhut.training.parallel_eval import ParallelEpisodeRunner, _run_single_episode
from malkhut.training.cma_trainer import PolicyEvaluator, ScenarioFactory
from malkhut.cwm.core import MinimalCryptoLOBCWM
from malkhut.state import FulfilmentPolicyParams


def _baseline_params() -> FulfilmentPolicyParams:
    return FulfilmentPolicyParams(
        version='test', ucb_c=1.414, max_sims=16, max_depth=2,
        rollout_depth=2, root_temperature=0.5, min_root_entropy=0.25,
        quote_offsets_ticks=(0, 1), quote_size_fractions=(0.25, 0.50),
        passive_ttl_ms=200, aggressive_ttl_ms=50,
        maker_edge_min_bps=0.5, cross_spread_edge_min_bps=5.0,
        adverse_toxicity_cancel_threshold=0.5, queue_churn_cancel_threshold=0.5,
        mae_tail_cut_bps=50.0, mfe_giveback_cut_fraction=0.5,
        max_time_in_loss_s=300.0, failed_recovery_cut_count=3,
        recovery_velocity_min_bps_per_s=0.0,
        max_symbol_notional_fraction=0.20, max_single_order_notional_fraction=0.05,
        reduce_when_global_up_fraction=0.30, session_profit_lock_fraction=0.02,
        w_expected_pnl=1.0, w_fill_probability=0.5, w_adverse_selection=2.0,
        w_queue_priority=0.5, w_inventory_risk=1.5, w_tail_loss=5.0,
        w_fee_quality=0.5, w_time_decay=0.3, w_policy_entropy=0.5,
        robust_tail_weight=2.0, toxic_counterparty_weight=3.0,
        low_liquidity_weight=2.0, latency_stress_weight=1.0,
    )


class TestParallelRunner:
    """Unit tests for ParallelEpisodeRunner."""

    def test_runner_init_default(self):
        runner = ParallelEpisodeRunner()
        assert runner.workers >= 1
        assert runner.workers <= 8

    def test_runner_init_explicit(self):
        runner = ParallelEpisodeRunner(workers=4)
        assert runner.workers == 4

    def test_runner_single_scenario(self):
        """Single scenario should use sequential path (no multiprocessing overhead)."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        params = _baseline_params()
        runner = ParallelEpisodeRunner(workers=4)
        results = runner.run_episodes(params, suite[:1], rng_seed=0)
        assert len(results) == 1
        assert results[0].scenario_id == suite[0].scenario_id

    def test_runner_parallel_vs_sequential_count(self):
        """Parallel and sequential must return same number of results."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        params = _baseline_params()
        runner = ParallelEpisodeRunner(workers=4)
        seq_results = runner._run_sequential(params, suite[:3], 0, 'sm_mcts')
        par_results = runner.run_episodes(params, suite[:3], rng_seed=0, planner_type='sm_mcts')
        assert len(seq_results) == len(par_results) == 3

    def test_runner_determinism(self):
        """Same seed → same scenario_id mapping (results are in order)."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT', 'ETHUSDT'), steps_per_scenario=3)
        params = _baseline_params()
        runner = ParallelEpisodeRunner(workers=4)
        r1 = runner.run_episodes(params, suite[:6], rng_seed=42)
        r2 = runner.run_episodes(params, suite[:6], rng_seed=42)
        for a, b in zip(r1, r2):
            assert a.scenario_id == b.scenario_id
            assert a.seed == b.seed

    def test_runner_different_seeds_different_results(self):
        """Different seeds should produce different episode results."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        params = _baseline_params()
        runner = ParallelEpisodeRunner(workers=2)
        r1 = runner.run_episodes(params, suite[:1], rng_seed=0)
        r2 = runner.run_episodes(params, suite[:1], rng_seed=999)
        assert r1[0].seed != r2[0].seed

    def test_runner_result_fields_valid(self):
        """Each result has all required fields populated."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        params = _baseline_params()
        runner = ParallelEpisodeRunner(workers=2)
        results = runner.run_episodes(params, suite[:2], rng_seed=0)
        for r in results:
            assert isinstance(r.scenario_id, str)
            assert isinstance(r.pnl_bps, float)
            assert isinstance(r.max_drawdown_bps, float)
            assert r.steps >= 0
            assert r.fill_count >= 0
            assert r.order_count >= 0


class TestRunSingleEpisode:
    """Unit tests for the multiprocessing top-level function."""

    def test_picklable_args(self):
        """Verify params and scenario can be pickled (required for multiprocessing)."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        params = _baseline_params()
        params_pkl = pickle.dumps(params)
        scenario_pkl = pickle.dumps(suite[0])
        restored_params = pickle.loads(params_pkl)
        restored_scenario = pickle.loads(scenario_pkl)
        assert restored_params.version == 'test'
        assert restored_scenario.scenario_id == suite[0].scenario_id

    def test_single_episode_produces_result(self):
        """_run_single_episode produces a valid EpisodeResult."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        params = _baseline_params()
        args = (pickle.dumps(params), pickle.dumps(suite[0]), 0, 'sm_mcts')
        result = _run_single_episode(args)
        assert result is not None
        assert result.scenario_id == suite[0].scenario_id


class TestEvaluateCandidateWorkers:
    """Tests for the workers parameter on PolicyEvaluator.evaluate_candidate."""

    def test_workers_0_uses_sequential(self):
        """workers=0 should use sequential (backward compatible)."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        evaluator = PolicyEvaluator(cwm_factory=MinimalCryptoLOBCWM)
        params = _baseline_params()
        score, results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:3], rng_seed=0, workers=0)
        assert len(results) == 3
        assert isinstance(score, float)

    def test_workers_1_uses_sequential(self):
        """workers=1 should use sequential."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        evaluator = PolicyEvaluator(cwm_factory=MinimalCryptoLOBCWM)
        params = _baseline_params()
        score, results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:3], rng_seed=0, workers=1)
        assert len(results) == 3

    def test_workers_4_uses_parallel(self):
        """workers=4 should use parallel path."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT', 'ETHUSDT'), steps_per_scenario=3)
        evaluator = PolicyEvaluator(cwm_factory=MinimalCryptoLOBCWM)
        params = _baseline_params()
        score, results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:6], rng_seed=0, workers=4)
        assert len(results) == 6
        assert isinstance(score, float)

    def test_workers_backward_compatible(self):
        """Default (no workers arg) should work exactly as before."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        evaluator = PolicyEvaluator(cwm_factory=MinimalCryptoLOBCWM)
        params = _baseline_params()
        score, results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:3], rng_seed=0)
        assert len(results) == 3


class TestParallelConsistency:
    """Cross-validation: parallel and sequential must agree on structural properties."""

    def test_same_scenario_count(self):
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT', 'ETHUSDT'), steps_per_scenario=3)
        params = _baseline_params()
        evaluator = PolicyEvaluator(cwm_factory=MinimalCryptoLOBCWM)

        _, seq_results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:6], rng_seed=0, workers=0)
        _, par_results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:6], rng_seed=0, workers=4)

        assert len(seq_results) == len(par_results)

    def test_same_scenarios_covered(self):
        """Both paths evaluate the same set of scenario IDs."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT', 'ETHUSDT'), steps_per_scenario=3)
        params = _baseline_params()
        evaluator = PolicyEvaluator(cwm_factory=MinimalCryptoLOBCWM)

        _, seq_results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:6], rng_seed=0, workers=0)
        _, par_results = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:6], rng_seed=0, workers=4)

        seq_ids = set(r.scenario_id for r in seq_results)
        par_ids = set(r.scenario_id for r in par_results)
        assert seq_ids == par_ids

    def test_same_score_magnitude(self):
        """Both paths should produce finite, non-NaN scores."""
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=('BTCUSDT',), steps_per_scenario=3)
        params = _baseline_params()
        evaluator = PolicyEvaluator(cwm_factory=MinimalCryptoLOBCWM)

        seq_score, _ = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:3], rng_seed=0, workers=0)
        par_score, _ = evaluator.evaluate_candidate(
            params=params, scenarios=suite[:3], rng_seed=0, workers=4)

        assert isinstance(seq_score, float)
        assert isinstance(par_score, float)
        assert seq_score == seq_score  # not NaN
        assert par_score == par_score  # not NaN