#!/usr/bin/env python3
"""
Tests for LDPC code analysis module (ldpc_analysis.py).

Covers IRA staircase code construction and rate sweep functionality.

Run:
    python3 -m pytest model/test_ldpc_analysis.py -v
"""

import numpy as np
import pytest

from ldpc_analysis import (
    build_ira_staircase,
    ira_encode,
    generic_decode,
    rate_sweep_single,
)
from ldpc_sim import poisson_channel, quantize_llr


# =============================================================================
# TestIRAStaircase
# =============================================================================

class TestIRAStaircase:
    """Validate generic IRA staircase code construction for various rates."""

    def test_rate_1_2(self):
        """m_base=1 -> rate 1/2: (1,2) base matrix, (32,64) full matrix, rank=32."""
        H_base, H_full = build_ira_staircase(m_base=1, z=32)
        assert H_base.shape == (1, 2), f"Base shape: {H_base.shape}"
        assert H_full.shape == (32, 64), f"Full shape: {H_full.shape}"
        rank = np.linalg.matrix_rank(H_full.astype(float))
        assert rank == 32, f"Rank: {rank}, expected 32"

    def test_rate_1_4(self):
        """m_base=3 -> rate 1/4: (3,4) base matrix, (96,128) full matrix, rank=96."""
        H_base, H_full = build_ira_staircase(m_base=3, z=32)
        assert H_base.shape == (3, 4), f"Base shape: {H_base.shape}"
        assert H_full.shape == (96, 128), f"Full shape: {H_full.shape}"
        rank = np.linalg.matrix_rank(H_full.astype(float))
        assert rank == 96, f"Rank: {rank}, expected 96"

    def test_rate_1_8_matches_existing(self):
        """m_base=7 -> rate 1/8: (7,8) base matrix, (224,256) full matrix."""
        H_base, H_full = build_ira_staircase(m_base=7, z=32)
        assert H_base.shape == (7, 8), f"Base shape: {H_base.shape}"
        assert H_full.shape == (224, 256), f"Full shape: {H_full.shape}"

    def test_encode_decode_roundtrip(self):
        """For m_base in [1,2,3,5], encode random info, verify syndrome=0."""
        np.random.seed(42)
        for m_base in [1, 2, 3, 5]:
            H_base, H_full = build_ira_staircase(m_base=m_base, z=32)
            z = 32
            k = z  # info bits = Z
            info = np.random.randint(0, 2, k).astype(np.int8)
            codeword = ira_encode(info, H_base, H_full, z=z)
            syndrome = H_full @ codeword % 2
            assert np.all(syndrome == 0), (
                f"m_base={m_base}: syndrome weight = {syndrome.sum()}"
            )


# =============================================================================
# TestRateSweep
# =============================================================================

class TestRateSweep:
    """Validate rate sweep simulation."""

    def test_high_snr_all_rates_decode(self):
        """At lam_s=20, all rates should achieve FER=0 with 10 frames."""
        np.random.seed(42)
        for m_base in [1, 2, 3, 5, 7]:
            result = rate_sweep_single(
                m_base=m_base, z=32, lam_s=20.0, lam_b=0.1,
                n_frames=10, max_iter=30, q_bits=6,
            )
            assert result['fer'] == 0.0, (
                f"m_base={m_base} (rate 1/{m_base+1}): FER={result['fer']} at lam_s=20, "
                f"expected 0.0"
            )