"""
Tests for watkins_nn.kyber — WatkinsKyberSampler module.

Phase 85.1 | CSID: KYBER-WATKINS-001
"""

import math
import torch
import pytest


# =============================================================================
# WatkinsSimplex2D Tests
# =============================================================================

class TestWatkinsSimplex2D:

    def test_golden_init_values(self):
        from watkins_nn.kyber import WatkinsSimplex2D
        from watkins_nn.constants import LAM_STAR

        simplex = WatkinsSimplex2D(init='golden')
        p = simplex()
        assert abs(p[0].item() - LAM_STAR) < 1e-4
        assert abs(p[1].item() - (1 - LAM_STAR)) < 1e-4

    def test_forward_sums_to_one(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        for init in ['golden', 'uniform', 'random']:
            simplex = WatkinsSimplex2D(init=init)
            p = simplex()
            assert abs(p.sum().item() - 1.0) < 1e-6, f"init={init} failed conservation"

    def test_uniform_init(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        simplex = WatkinsSimplex2D(init='uniform')
        p = simplex()
        assert abs(p[0].item() - 0.5) < 1e-5
        assert abs(p[1].item() - 0.5) < 1e-5

    def test_random_init(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        simplex = WatkinsSimplex2D(init='random')
        p = simplex()
        assert p[0].item() > 0 and p[0].item() < 1
        assert p[1].item() > 0 and p[1].item() < 1

    def test_invalid_init_raises(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        with pytest.raises(ValueError, match="Unknown init"):
            WatkinsSimplex2D(init='bogus')

    def test_get_coherence(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        simplex = WatkinsSimplex2D(init='golden')
        lam = simplex.get_coherence()
        p = simplex()
        assert abs(lam.item() - p[0].item()) < 1e-10

    def test_get_noise(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        simplex = WatkinsSimplex2D(init='golden')
        eta = simplex.get_noise()
        p = simplex()
        assert abs(eta.item() - p[1].item()) < 1e-10

    def test_conservation_coherence_plus_noise(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        simplex = WatkinsSimplex2D(init='golden')
        lam = simplex.get_coherence().item()
        eta = simplex.get_noise().item()
        assert abs(lam + eta - 1.0) < 1e-7

    def test_dim_attribute(self):
        from watkins_nn.kyber import WatkinsSimplex2D

        simplex = WatkinsSimplex2D()
        assert simplex.dim == 2


# =============================================================================
# WatkinsKyberSampler Tests
# =============================================================================

class TestWatkinsKyberSampler:

    def test_construction_eta2(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=256)
        assert sampler.eta == 2
        assert sampler.dim == 256

    def test_construction_eta3(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=3, dimension=512)
        assert sampler.eta == 3
        assert sampler.dim == 512

    def test_centered_binomial_shape(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=64)
        noise = sampler._centered_binomial((4, 64))
        assert noise.shape == (4, 64)

    def test_centered_binomial_support_eta2(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=256)
        noise = sampler._centered_binomial((1000, 256))
        assert noise.min().item() >= -2
        assert noise.max().item() <= 2

    def test_centered_binomial_support_eta3(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=3, dimension=256)
        noise = sampler._centered_binomial((1000, 256))
        assert noise.min().item() >= -3
        assert noise.max().item() <= 3

    def test_centered_binomial_mean_approx_zero(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=256)
        noise = sampler._centered_binomial((10000, 256))
        mean = noise.mean().item()
        assert abs(mean) < 0.05, f"CBD mean should be ~0, got {mean}"

    def test_sample_shape(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=64, steps=2)
        center = torch.zeros(64)
        result = sampler.sample(center, batch_size=3)
        assert result.shape == (3, 64)

    def test_sample_single(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32, steps=2)
        center = torch.zeros(32)
        result = sampler.sample(center)
        assert result.shape == (1, 32)

    def test_coherence_valid_range(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32, steps=4)
        center = torch.zeros(32)
        sampler.sample(center)
        lam = sampler.simplex.get_coherence().item()
        assert 0 < lam < 1

    def test_get_phi_positive(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32)
        phi = sampler.get_phi()
        assert phi > 0

    def test_security_margin_positive(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32)
        margin = sampler.security_margin()
        assert margin > 0

    def test_get_state_keys(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32)
        state = sampler.get_state()
        expected_keys = {
            'coherence_lambda', 'noise_eta', 'phi',
            'security_margin', 'in_habitable_wedge', 'samples_generated',
        }
        assert set(state.keys()) == expected_keys

    def test_get_state_conservation(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32)
        state = sampler.get_state()
        assert abs(state['coherence_lambda'] + state['noise_eta'] - 1.0) < 1e-7

    def test_reset_restores_golden(self):
        from watkins_nn.kyber import WatkinsKyberSampler
        from watkins_nn.constants import LAM_STAR

        sampler = WatkinsKyberSampler(eta=2, dimension=32, steps=4)
        center = torch.zeros(32)

        # Run a few samples to drift from equilibrium
        for _ in range(3):
            sampler.sample(center)

        # Reset
        sampler.reset()
        lam = sampler.simplex.get_coherence().item()
        assert abs(lam - LAM_STAR) < 1e-3
        assert len(sampler.phi_history) == 0
        assert len(sampler.coherence_history) == 0

    def test_history_tracking(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32, steps=2)
        center = torch.zeros(32)

        assert len(sampler.phi_history) == 0
        sampler.sample(center)
        assert len(sampler.phi_history) == 1
        sampler.sample(center)
        assert len(sampler.phi_history) == 2

    def test_simplex_conservation_after_sampling(self):
        from watkins_nn.kyber import WatkinsKyberSampler

        sampler = WatkinsKyberSampler(eta=2, dimension=32, steps=4)
        center = torch.zeros(32)

        for _ in range(5):
            sampler.sample(center)
            p = sampler.simplex()
            assert abs(p.sum().item() - 1.0) < 1e-6, "2-simplex conservation violated"


# =============================================================================
# Package Import Tests
# =============================================================================

class TestKyberPackageImport:

    def test_import_from_package(self):
        from watkins_nn import WatkinsSimplex2D, WatkinsKyberSampler
        assert WatkinsSimplex2D is not None
        assert WatkinsKyberSampler is not None

    def test_in_all(self):
        import watkins_nn
        assert "WatkinsSimplex2D" in watkins_nn.__all__
        assert "WatkinsKyberSampler" in watkins_nn.__all__

    def test_version_bumped(self):
        import watkins_nn
        assert watkins_nn.__version__ == "0.6.0"
