"""
Tests for strategy generator — genetic programming for strategy evolution.

Verifies:
  - Genome creation and representation
  - Crossover produces valid offspring
  - Mutation produces valid variants
  - Tournament selection preferentially selects fitter genomes
  - Population initialization includes baseline
  - Evolution improves fitness over generations
  - Successful strategies are added to pool
  - Diverse strategies are returned
  - Hardcoded baseline is never replaced
"""
import random
import pytest
from malkhut.state import FulfilmentPolicyParams
from malkhut.training.generator import (
    StrategyGenome, StrategyType, GeneticOperators,
    StrategyEvaluator, StrategyGenerator, GeneratorConfig,
)
from malkhut.training.cma_trainer import CMAParameterCodec, SelfPlayPool
from malkhut.training.registry import PolicyRegistry


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. STRATEGY GENOME
# ══════════════════════════════════════════════════════════════════════════════

class TestStrategyGenome:
    def test_construction(self):
        genome = StrategyGenome(
            strategy_type=StrategyType.SM_MCTS,
            params=_baseline(),
        )
        assert genome.strategy_type == StrategyType.SM_MCTS
        assert genome.generation == 0
        assert genome.fitness == 0.0

    def test_genome_id(self):
        genome = StrategyGenome(
            strategy_type=StrategyType.SM_MCTS,
            params=_baseline(version="v1"),
            generation=3,
        )
        assert "SM_MCTS" in genome.genome_id
        assert "v1" in genome.genome_id
        assert "3" in genome.genome_id

    def test_frozen(self):
        genome = StrategyGenome(
            strategy_type=StrategyType.SM_MCTS,
            params=_baseline(),
        )
        with pytest.raises(AttributeError):
            genome.fitness = 10.0


# ══════════════════════════════════════════════════════════════════════════════
# 2. GENETIC OPERATORS
# ══════════════════════════════════════════════════════════════════════════════

class TestGeneticOperators:
    def _ops(self):
        return GeneticOperators(codec=CMAParameterCodec())

    def test_crossover_produces_valid_child(self):
        ops = self._ops()
        p1 = StrategyGenome(strategy_type=StrategyType.SM_MCTS, params=_baseline(version="p1"))
        p2 = StrategyGenome(strategy_type=StrategyType.UCB1, params=_baseline(version="p2"))
        child = ops.crossover(p1, p2, random.Random(42))
        assert isinstance(child, StrategyGenome)
        assert child.strategy_type in (StrategyType.SM_MCTS, StrategyType.UCB1)
        assert child.generation == 1
        assert len(child.parent_ids) == 2

    def test_crossover_inherits_fitter_parent_type(self):
        ops = self._ops()
        p1 = StrategyGenome(strategy_type=StrategyType.SM_MCTS, params=_baseline(), fitness=10.0)
        p2 = StrategyGenome(strategy_type=StrategyType.UCB1, params=_baseline(), fitness=5.0)
        child = ops.crossover(p1, p2, random.Random(42))
        assert child.strategy_type == StrategyType.SM_MCTS

    def test_mutation_produces_valid_child(self):
        ops = self._ops()
        parent = StrategyGenome(strategy_type=StrategyType.SM_MCTS, params=_baseline())
        child = ops.mutate(parent, random.Random(42))
        assert isinstance(child, StrategyGenome)
        assert child.generation == 1
        assert len(child.parent_ids) == 1

    def test_mutation_different_params(self):
        ops = self._ops()
        parent = StrategyGenome(strategy_type=StrategyType.SM_MCTS, params=_baseline())
        child = ops.mutate(parent, random.Random(42))
        # With high mutation rate, params should differ
        assert child.params.ucb_c != parent.params.ucb_c or child.params.root_temperature != parent.params.root_temperature

    def test_tournament_select_prefers_fitter(self):
        ops = self._ops()
        pop = [
            StrategyGenome(strategy_type=StrategyType.SM_MCTS, params=_baseline(), fitness=1.0),
            StrategyGenome(strategy_type=StrategyType.UCB1, params=_baseline(), fitness=10.0),
            StrategyGenome(strategy_type=StrategyType.GREEDY, params=_baseline(), fitness=5.0),
        ]
        wins = 0
        for i in range(200):
            winner = ops.tournament_select(pop, tournament_size=2, rng=random.Random(i))
            if winner.fitness == 10.0:
                wins += 1
        # With 3 genomes and tournament_size=2, fitter wins ~50% (vs 25% random)
        assert wins > 60

    def test_random_genome_valid(self):
        ops = self._ops()
        genome = ops.random_genome(rng=random.Random(42))
        assert isinstance(genome, StrategyGenome)
        assert genome.strategy_type in list(StrategyType)
        assert genome.generation == 0

    def test_random_genome_respects_type(self):
        ops = self._ops()
        genome = ops.random_genome(strategy_type=StrategyType.THOMPSON_SAMPLING, rng=random.Random(42))
        assert genome.strategy_type == StrategyType.THOMPSON_SAMPLING


# ══════════════════════════════════════════════════════════════════════════════
# 3. STRATEGY GENERATOR
# ══════════════════════════════════════════════════════════════════════════════

class TestStrategyGenerator:
    def test_initialize_population(self):
        gen = StrategyGenerator(config=GeneratorConfig(population_size=5, generations=1))
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.initialize_population(_baseline(), scenarios)
        assert gen.population_size == 5
        # First should be baseline
        assert gen.population[0].strategy_type == StrategyType.SM_MCTS

    def test_baseline_always_first(self):
        gen = StrategyGenerator(config=GeneratorConfig(population_size=10, generations=1))
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.initialize_population(_baseline(), scenarios)
        # Baseline should be in population
        types = [g.strategy_type for g in gen.population]
        assert StrategyType.SM_MCTS in types

    def test_evolve_returns_population(self):
        config = GeneratorConfig(population_size=5, generations=2)
        gen = StrategyGenerator(config=config)
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        result = gen.evolve(_baseline(), scenarios)
        # Population may be slightly larger due to baseline preservation
        assert len(result) >= 5
        assert all(isinstance(g, StrategyGenome) for g in result)

    def test_evolve_improves_over_generations(self):
        config = GeneratorConfig(population_size=8, generations=3, elitism_count=2)
        gen = StrategyGenerator(config=config)
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.evolve(_baseline(), scenarios)
        # Population should have diverse fitness
        fitnesses = [g.fitness for g in gen.population]
        assert len(set(fitnesses)) > 1  # not all same

    def test_get_successful_strategies(self):
        gen = StrategyGenerator(config=GeneratorConfig(population_size=5, generations=1))
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.evolve(_baseline(), scenarios)
        successful = gen.get_successful_strategies(min_fitness=-1000)
        assert len(successful) > 0

    def test_get_diverse_strategies(self):
        gen = StrategyGenerator(config=GeneratorConfig(population_size=10, generations=2))
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.evolve(_baseline(), scenarios)
        diverse = gen.get_diverse_strategies(n=3)
        assert len(diverse) == 3

    def test_add_to_pool(self):
        gen = StrategyGenerator()
        genome = StrategyGenome(
            strategy_type=StrategyType.SM_MCTS,
            params=_baseline(version="test"),
            fitness=10.0,
        )
        gen.add_to_pool(genome)
        assert len(gen._pool.policies()) == 1

    def test_baseline_never_replaced(self):
        """The hardcoded baseline should always be in the population."""
        gen = StrategyGenerator(config=GeneratorConfig(population_size=5, generations=3))
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.evolve(_baseline(), scenarios)
        # At least one SM_MCTS should exist (the baseline)
        assert any(g.strategy_type == StrategyType.SM_MCTS for g in gen.population)

    def test_best_fitness_tracked(self):
        gen = StrategyGenerator(config=GeneratorConfig(population_size=5, generations=1))
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.evolve(_baseline(), scenarios)
        assert gen.best_fitness > -float("inf")

    def test_history_tracked(self):
        gen = StrategyGenerator(config=GeneratorConfig(population_size=5, generations=2))
        from malkhut.training.cma_trainer import ScenarioFactory
        factory = ScenarioFactory()
        scenarios = factory.build_suite(symbols=("BTCUSDT",), steps_per_scenario=3)
        gen.evolve(_baseline(), scenarios)
        assert len(gen._history) > 0