#!/usr/bin/env python3
"""Tests for density evolution optimizer."""

import numpy as np
import pytest
import sys
import os

sys.path.insert(0, os.path.dirname(__file__))


class TestDensityEvolution:
    """Tests for the Monte Carlo DE engine."""

    def test_de_known_good_converges(self):
        """DE with original staircase profile at lam_s=10 should converge easily."""
        from density_evolution import run_de, ORIGINAL_STAIRCASE_PROFILE
        np.random.seed(42)
        converged, error_frac = run_de(
            ORIGINAL_STAIRCASE_PROFILE, lam_s=10.0, lam_b=0.1,
            z_pop=10000, max_iter=50
        )
        assert converged, f"DE should converge at lam_s=10, error_frac={error_frac}"

    def test_de_known_bad_fails(self):
        """DE at very low lam_s=0.1 should not converge."""
        from density_evolution import run_de, ORIGINAL_STAIRCASE_PROFILE
        np.random.seed(42)
        converged, error_frac = run_de(
            ORIGINAL_STAIRCASE_PROFILE, lam_s=0.1, lam_b=0.1,
            z_pop=10000, max_iter=50
        )
        assert not converged, f"DE should NOT converge at lam_s=0.1, error_frac={error_frac}"

    def test_de_population_shape(self):
        """Verify belief arrays have correct shapes after one step."""
        from density_evolution import de_channel_init, density_evolution_step
        np.random.seed(42)
        n_base = 8
        m_base = 7
        z_pop = 1000

        # Original staircase H_base profile
        from density_evolution import ORIGINAL_STAIRCASE_PROFILE
        beliefs, msg_memory = de_channel_init(ORIGINAL_STAIRCASE_PROFILE, z_pop, lam_s=5.0, lam_b=0.1)

        # beliefs should be (n_base, z_pop)
        assert beliefs.shape == (n_base, z_pop), f"Expected ({n_base}, {z_pop}), got {beliefs.shape}"

        # Run one step
        beliefs = density_evolution_step(beliefs, msg_memory, ORIGINAL_STAIRCASE_PROFILE, z_pop)
        assert beliefs.shape == (n_base, z_pop), f"Shape changed after step: {beliefs.shape}"


class TestThresholdComputation:
    """Tests for threshold binary search."""

    def test_threshold_original_staircase(self):
        """Threshold for original staircase [7,2,2,2,2,2,2,1] should be ~3-6 photons."""
        from density_evolution import compute_threshold_for_profile
        np.random.seed(42)
        threshold = compute_threshold_for_profile(
            [7, 2, 2, 2, 2, 2, 2, 1], m_base=7, lam_b=0.1,
            z_pop=10000, tol=0.5
        )
        assert 2.0 < threshold < 8.0, f"Expected threshold ~3-6, got {threshold}"

    def test_threshold_peg_ring(self):
        """PEG ring [7,3,3,3,2,2,2,2] should have lower or equal threshold than original."""
        from density_evolution import compute_threshold_for_profile
        np.random.seed(42)
        thresh_orig = compute_threshold_for_profile(
            [7, 2, 2, 2, 2, 2, 2, 1], m_base=7, lam_b=0.1,
            z_pop=15000, tol=0.25
        )
        np.random.seed(123)
        thresh_peg = compute_threshold_for_profile(
            [7, 3, 3, 3, 2, 2, 2, 2], m_base=7, lam_b=0.1,
            z_pop=15000, tol=0.25
        )
        assert thresh_peg <= thresh_orig, (
            f"PEG threshold {thresh_peg} should be <= original {thresh_orig}"
        )

    def test_profile_to_hbase(self):
        """build_de_profile should produce valid profile with correct column degrees."""
        from density_evolution import build_de_profile
        profile = build_de_profile([7, 3, 2, 2, 2, 2, 2, 2], m_base=7)
        assert profile['n_base'] == 8
        assert profile['m_base'] == 7
        assert profile['vn_degrees'] == [7, 3, 2, 2, 2, 2, 2, 2]
        # Every row should have at least 2 connections
        for r, conns in enumerate(profile['connections']):
            assert len(conns) >= 2, f"Row {r} has only {len(conns)} connections"