"""
Exhaustive training subsystem tests.

Categories:
  1. CMA Parameter Codec (encode/decode/bounds)
  2. SelfPlayPool (add/evict/diversity)
  3. Scenario Factory (suite generation)
  4. Policy Evaluator (real episodes, metrics)
  5. EpisodeResult (fields, computation)
  6. Bootstrap CI
  7. CMAESTrainer (end-to-end)
  8. Promotion logic
"""
import dataclasses
import math
import pytest
from malkhut.state import (
    AccountState, ExecutionIntent, FulfilmentPolicyParams, IntentKind,
    MarketWorldState, Mode, OrderBookState, PriceLevel, Side, VenueRules,
)
from malkhut.actions import ActionKind, FulfilmentAction, OrderType
from malkhut.cwm.core import MinimalCryptoLOBCWM
from malkhut.counterparties import default_counterparty_ecology
from malkhut.training.cma_trainer import (
    CMAParameterCodec, SelfPlayPool, PolicySnapshot,
    ScenarioFactory, PolicyEvaluator, EpisodeResult,
    CMAESTrainer, bootstrap_ci,
)


def _baseline(**kw):
    d = dict(
        version="baseline", ucb_c=1.414, max_sims=256, max_depth=3,
        rollout_depth=3, root_temperature=0.5, min_root_entropy=0.25,
        quote_offsets_ticks=(0, 1, 2), quote_size_fractions=(0.1, 0.25, 0.5),
        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,
    )
    d.update(kw)
    return FulfilmentPolicyParams(**d)


# ══════════════════════════════════════════════════════════════════════════════
# 1. CMA PARAMETER CODEC
# ══════════════════════════════════════════════════════════════════════════════

class TestCMAParameterCodec:
    def test_initial_vector_length(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        assert len(x0) == len(codec.SPECS)

    def test_bounds_length(self):
        codec = CMAParameterCodec()
        lows, highs = codec.bounds()
        assert len(lows) == len(codec.SPECS)
        assert len(highs) == len(codec.SPECS)

    def test_bounds_ordered(self):
        codec = CMAParameterCodec()
        lows, highs = codec.bounds()
        assert all(l <= h for l, h in zip(lows, highs))

    def test_initial_vector_midpoint(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        lows, highs = codec.bounds()
        for v, lo, hi in zip(x0, lows, highs):
            assert lo <= v <= hi

    def test_decode_returns_params(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        p = codec.decode(x0, "v_test")
        assert isinstance(p, FulfilmentPolicyParams)

    def test_decode_version_preserved(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        p = codec.decode(x0, "my_version")
        assert p.version == "my_version"

    def test_decode_int_fields_are_integers(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        p = codec.decode(x0, "int_test")
        assert isinstance(p.max_depth, int)
        assert isinstance(p.passive_ttl_ms, int)
        assert isinstance(p.failed_recovery_cut_count, int)

    def test_decode_clips_above(self):
        codec = CMAParameterCodec()
        highs = [s.high for s in codec.SPECS]
        x_over = [h + 10.0 for h in highs]
        p = codec.decode(x_over, "over")
        for spec in codec.SPECS:
            val = getattr(p, spec.name)
            if spec.kind == "float":
                assert val <= spec.high + 1e-9

    def test_decode_clips_below(self):
        codec = CMAParameterCodec()
        lows = [s.low for s in codec.SPECS]
        x_under = [l - 10.0 for l in lows]
        p = codec.decode(x_under, "under")
        for spec in codec.SPECS:
            val = getattr(p, spec.name)
            if spec.kind == "float":
                assert val >= spec.low - 1e-9

    def test_decode_idempotent(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        p1 = codec.decode(x0, "a")
        p2 = codec.decode(x0, "b")
        assert p1.ucb_c == p2.ucb_c
        assert p1.max_depth == p2.max_depth

    def test_different_vectors_different_params(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        x1 = list(x0)
        x1[0] += 0.5
        p0 = codec.decode(x0, "a")
        p1 = codec.decode(x1, "b")
        assert p0.ucb_c != p1.ucb_c

    def test_decode_midpoint_values(self):
        codec = CMAParameterCodec()
        x0 = codec.initial_vector(_baseline())
        p = codec.decode(x0, "mid")
        # ucb_c midpoint = (0.2 + 3.0) / 2 = 1.6
        assert p.ucb_c == pytest.approx(1.6, abs=0.01)


# ══════════════════════════════════════════════════════════════════════════════
# 2. SELF-PLAY POOL
# ══════════════════════════════════════════════════════════════════════════════

class TestSelfPlayPool:
    def test_add_and_retrieve(self):
        pool = SelfPlayPool(max_size=5)
        snap = PolicySnapshot(
            params=_baseline(version="v1"), score=10.0,
            created_ts_ns=1000, performance_vector=(1.0, 0.5, 0.8, 0.3, 0.1),
        )
        pool.maybe_add(snap)
        assert len(pool.policies()) == 1
        assert pool.policies()[0].version == "v1"

    def test_evict_lowest_when_full(self):
        pool = SelfPlayPool(max_size=3)
        for i in range(5):
            pool.maybe_add(PolicySnapshot(
                params=_baseline(version=f"v{i}"), score=float(i),
                created_ts_ns=1000 + i,
                performance_vector=(float(i), float(i) * 0.1, 0.5, 0.3, 0.1),
            ))
        assert len(pool.policies()) == 3
        # Should keep top 3 by score (diversity eviction keeps diverse ones)
        versions = [s.params.version for s in pool.snapshots]
        assert "v4" in versions  # highest score always kept

    def test_diversity_preservation(self):
        pool = SelfPlayPool(max_size=3)
        # Add 3 diverse policies
        pool.maybe_add(PolicySnapshot(
            params=_baseline(version="a"), score=10.0,
            created_ts_ns=1000, performance_vector=(1.0, 0.0, 0.5, 0.3, 0.1),
        ))
        pool.maybe_add(PolicySnapshot(
            params=_baseline(version="b"), score=9.0,
            created_ts_ns=1001, performance_vector=(0.0, 1.0, 0.5, 0.3, 0.1),
        ))
        pool.maybe_add(PolicySnapshot(
            params=_baseline(version="c"), score=8.0,
            created_ts_ns=1002, performance_vector=(0.5, 0.5, 0.5, 0.3, 0.1),
        ))
        # Add a similar one — should not evict diverse ones
        pool.maybe_add(PolicySnapshot(
            params=_baseline(version="d"), score=7.0,
            created_ts_ns=1003, performance_vector=(1.0, 0.0, 0.5, 0.3, 0.1),
        ))
        assert len(pool.policies()) == 3

    def test_policies_returns_tuple(self):
        pool = SelfPlayPool(max_size=5)
        pool.maybe_add(PolicySnapshot(
            params=_baseline(version="v1"), score=10.0,
            created_ts_ns=1000,
        ))
        assert isinstance(pool.policies(), tuple)

    def test_empty_pool(self):
        pool = SelfPlayPool(max_size=5)
        assert len(pool.policies()) == 0

    def test_cosine_similarity(self):
        a = (1.0, 0.0, 0.0)
        b = (1.0, 0.0, 0.0)
        assert SelfPlayPool._cosine_similarity(a, b) == pytest.approx(1.0)

    def test_cosine_orthogonal(self):
        a = (1.0, 0.0, 0.0)
        b = (0.0, 1.0, 0.0)
        assert SelfPlayPool._cosine_similarity(a, b) == pytest.approx(0.0)

    def test_cosine_empty(self):
        assert SelfPlayPool._cosine_similarity((), ()) == 0.0


# ══════════════════════════════════════════════════════════════════════════════
# 3. SCENARIO FACTORY
# ══════════════════════════════════════════════════════════════════════════════

class TestScenarioFactory:
    def test_build_suite_returns_scenarios(self):
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=10)
        assert len(suite) >= 30  # 30 real market scenarios per symbol

    def test_each_scenario_has_unique_id(self):
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=10)
        ids = [s.scenario_id for s in suite]
        assert len(set(ids)) == len(ids)

    def test_scenarios_have_different_tags(self):
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=10)
        all_tags = set()
        for s in suite:
            all_tags.update(s.tags)
        assert "normal" in all_tags
        assert "thin" in all_tags
        assert "toxic" in all_tags

    def test_scenario_has_valid_state(self):
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=10)
        for s in suite:
            assert s.initial_state.book.best_bid > 0
            assert s.initial_state.book.best_ask > 0
            assert s.initial_state.account.equity > 0

    def test_multi_symbol(self):
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=("BTCUSDT", "ETHUSDT"), steps_per_scenario=5)
        assert len(suite) >= 60  # 30 per symbol × 2 symbols

    def test_scenario_counterparties(self):
        factory = ScenarioFactory()
        suite = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=10)
        for s in suite:
            assert len(s.counterparties) > 0


# ══════════════════════════════════════════════════════════════════════════════
# 4. POLICY EVALUATOR
# ══════════════════════════════════════════════════════════════════════════════

class TestPolicyEvaluator:
    def _evaluator(self):
        return PolicyEvaluator(
            cwm_factory=lambda: MinimalCryptoLOBCWM(),
            counterparties=default_counterparty_ecology(),
        )

    def _scenario(self):
        factory = ScenarioFactory()
        return factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=5)[0]

    def test_evaluate_returns_score_and_results(self):
        eval_ = self._evaluator()
        scenario = self._scenario()
        score, results = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        assert isinstance(score, float)
        assert len(results) == 1
        assert isinstance(results[0], EpisodeResult)

    def test_episode_has_steps(self):
        eval_ = self._evaluator()
        scenario = self._scenario()
        _, results = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        assert results[0].steps > 0

    def test_episode_has_pnl(self):
        eval_ = self._evaluator()
        scenario = self._scenario()
        _, results = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        assert isinstance(results[0].pnl_bps, float)

    def test_episode_has_drawdown(self):
        eval_ = self._evaluator()
        scenario = self._scenario()
        _, results = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        assert results[0].max_drawdown_bps >= 0

    def test_episode_has_final_equity(self):
        eval_ = self._evaluator()
        scenario = self._scenario()
        _, results = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        assert results[0].final_equity > 0

    def test_deterministic_with_same_seed(self):
        eval_ = self._evaluator()
        scenario = self._scenario()
        _, r1 = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        _, r2 = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        assert r1[0].pnl_bps == r2[0].pnl_bps

    def test_different_seeds_different_results(self):
        eval_ = self._evaluator()
        scenario = self._scenario()
        _, r1 = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=42,
        )
        _, r2 = eval_.evaluate_candidate(
            params=_baseline(version="test"), scenarios=[scenario], rng_seed=99,
        )
        # Different seeds produce different planner distributions (entropy differs)
        assert r1[0].policy_entropy_avg != r2[0].policy_entropy_avg

    def test_performance_vector(self):
        results = [
            EpisodeResult(scenario_id="s1", policy_version="v1", seed=1, pnl_bps=10.0, max_drawdown_bps=5.0, fill_ratio=0.8, policy_entropy_avg=0.5, order_count=10, cancel_count=2),
            EpisodeResult(scenario_id="s2", policy_version="v1", seed=2, pnl_bps=-5.0, max_drawdown_bps=10.0, fill_ratio=0.6, policy_entropy_avg=0.3, order_count=8, cancel_count=1),
        ]
        vec = PolicyEvaluator.performance_vector(results)
        assert len(vec) == 5
        assert vec[0] == pytest.approx(2.5, abs=0.01)  # mean PnL


# ══════════════════════════════════════════════════════════════════════════════
# 5. BOOTSTRAP CI
# ══════════════════════════════════════════════════════════════════════════════

class TestBootstrapCI:
    def test_ci_returns_tuple(self):
        result = bootstrap_ci([1.0, 2.0, 3.0, 4.0, 5.0])
        assert len(result) == 3

    def test_ci_mean_correct(self):
        scores = [1.0, 2.0, 3.0, 4.0, 5.0]
        mean, lo, hi = bootstrap_ci(scores, n_bootstrap=100)
        assert mean == pytest.approx(3.0, abs=0.01)

    def test_ci_contains_mean(self):
        scores = [1.0, 2.0, 3.0, 4.0, 5.0]
        mean, lo, hi = bootstrap_ci(scores, n_bootstrap=100)
        assert lo <= mean <= hi

    def test_ci_wider_with_fewer_samples(self):
        scores = [1.0, 2.0, 3.0]
        _, lo1, hi1 = bootstrap_ci(scores, n_bootstrap=100)
        scores2 = [1.0] * 100
        _, lo2, hi2 = bootstrap_ci(scores2, n_bootstrap=100)
        # More uniform data → tighter CI
        assert (hi1 - lo1) > (hi2 - lo2)

    def test_ci_empty(self):
        mean, lo, hi = bootstrap_ci([])
        assert mean == 0.0

    def test_ci_single_value(self):
        mean, lo, hi = bootstrap_ci([5.0])
        assert mean == 5.0


# ══════════════════════════════════════════════════════════════════════════════
# 6. PROMOTION LOGIC
# ══════════════════════════════════════════════════════════════════════════════

class TestPromotion:
    def _trainer(self):
        codec = CMAParameterCodec()
        evaluator = PolicyEvaluator(
            cwm_factory=lambda: MinimalCryptoLOBCWM(),
            counterparties=default_counterparty_ecology(),
        )
        pool = SelfPlayPool(max_size=5)
        return CMAESTrainer(codec=codec, evaluator=evaluator, pool=pool)

    def test_promote_better_candidate(self):
        trainer = self._trainer()
        incumbent = PolicySnapshot(
            params=_baseline(version="inc"), score=10.0,
            created_ts_ns=1000, evaluation_summary={"mean_pnl": 5.0, "max_dd": 2.0},
            performance_vector=(5.0, 2.0, 0.8, 0.3, 0.1),
        )
        candidate = PolicySnapshot(
            params=_baseline(version="cand"), score=12.0,
            created_ts_ns=1001, evaluation_summary={"mean_pnl": 7.0, "max_dd": 1.5},
            performance_vector=(7.0, 1.5, 0.9, 0.4, 0.05),
        )
        ok, reason = trainer.promote(candidate, incumbent)
        assert ok
        assert reason == "promoted"

    def test_reject_worse_candidate(self):
        trainer = self._trainer()
        incumbent = PolicySnapshot(
            params=_baseline(version="inc"), score=10.0,
            created_ts_ns=1000, evaluation_summary={},
            performance_vector=(1.0,),
        )
        candidate = PolicySnapshot(
            params=_baseline(version="cand"), score=9.0,
            created_ts_ns=1001, evaluation_summary={},
            performance_vector=(0.5,),
        )
        ok, reason = trainer.promote(candidate, incumbent)
        assert not ok
        assert reason == "insufficient_edge"

    def test_reject_no_performance_vector(self):
        trainer = self._trainer()
        incumbent = PolicySnapshot(
            params=_baseline(version="inc"), score=10.0,
            created_ts_ns=1000, evaluation_summary={},
        )
        candidate = PolicySnapshot(
            params=_baseline(version="cand"), score=15.0,
            created_ts_ns=1001, evaluation_summary={},
        )
        ok, reason = trainer.promote(candidate, incumbent)
        assert not ok
        assert reason == "no_performance_vector"


# ══════════════════════════════════════════════════════════════════════════════
# 7. CMA-ES TRAINER (end-to-end)
# ══════════════════════════════════════════════════════════════════════════════

class TestCMAESTrainer:
    def test_train_returns_snapshot(self):
        codec = CMAParameterCodec()
        evaluator = PolicyEvaluator(
            cwm_factory=lambda: MinimalCryptoLOBCWM(),
            counterparties=default_counterparty_ecology(),
        )
        pool = SelfPlayPool(max_size=5)
        trainer = CMAESTrainer(codec=codec, evaluator=evaluator, pool=pool)

        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=5)

        result = trainer.train(
            incumbent=_baseline(version="init"),
            scenarios=scenarios,
            budget_evals=14,
            seed=42,
        )
        assert isinstance(result, PolicySnapshot)
        assert result.params is not None

    def test_pool_grows_during_training(self):
        codec = CMAParameterCodec()
        evaluator = PolicyEvaluator(
            cwm_factory=lambda: MinimalCryptoLOBCWM(),
            counterparties=default_counterparty_ecology(),
        )
        pool = SelfPlayPool(max_size=5)
        trainer = CMAESTrainer(codec=codec, evaluator=evaluator, pool=pool)

        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)

        trainer.train(
            incumbent=_baseline(version="init"),
            scenarios=scenarios,
            budget_evals=14,
            seed=42,
        )
        assert len(pool.policies()) > 0