#!/usr/bin/env python3
"""
Tests for the frame synchronization prototype (frame_sync.py).

Tests cover stream generation (length, offsets) and syndrome screening
(correct offset identification, wrong offset rejection).

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

import numpy as np
import pytest

from ldpc_sim import N, build_full_h_matrix
from frame_sync import generate_stream, syndrome_screen


# =============================================================================
# Fixtures
# =============================================================================

@pytest.fixture(scope="module")
def H():
    """Full expanded H matrix, built once per module."""
    return build_full_h_matrix()


# =============================================================================
# TestStreamGeneration
# =============================================================================

class TestStreamGeneration:
    """Validate continuous LLR stream generation."""

    def test_stream_length(self, H):
        """Stream should be n_frames * N + offset long."""
        np.random.seed(100)
        n_frames = 5
        offset = 37
        stream_llr, true_offset, info_list = generate_stream(
            H, n_frames=n_frames, lam_s=5.0, lam_b=0.1, offset=offset
        )
        assert len(stream_llr) == n_frames * N + offset
        assert true_offset == offset
        assert len(info_list) == n_frames

    def test_stream_zero_offset(self, H):
        """offset=0 should work, producing stream of length n_frames * N."""
        np.random.seed(101)
        n_frames = 3
        stream_llr, true_offset, info_list = generate_stream(
            H, n_frames=n_frames, lam_s=5.0, lam_b=0.1, offset=0
        )
        assert len(stream_llr) == n_frames * N
        assert true_offset == 0
        assert len(info_list) == n_frames

    def test_stream_random_offset(self, H):
        """offset=None should pick a random offset in [0, N-1]."""
        np.random.seed(102)
        stream_llr, true_offset, info_list = generate_stream(
            H, n_frames=4, lam_s=5.0, lam_b=0.1, offset=None
        )
        assert 0 <= true_offset < N
        assert len(stream_llr) == 4 * N + true_offset


# =============================================================================
# TestSyndromeScreen
# =============================================================================

class TestSyndromeScreen:
    """Validate syndrome-based offset screening."""

    def test_correct_offset_low_syndrome(self, H):
        """At high SNR (lam_s=10), the best offset should match the true offset."""
        np.random.seed(200)
        n_frames = 3
        offset = 42
        stream_llr, true_offset, _ = generate_stream(
            H, n_frames=n_frames, lam_s=10.0, lam_b=0.1, offset=offset
        )
        scores = syndrome_screen(stream_llr, n_offsets=N)
        # The true offset should have the lowest (or near-lowest) syndrome weight
        best_offset = min(scores, key=scores.get)
        assert best_offset == true_offset, (
            f"Best offset {best_offset} (sw={scores[best_offset]}) "
            f"!= true offset {true_offset} (sw={scores[true_offset]})"
        )

    def test_wrong_offsets_high_syndrome(self, H):
        """Wrong offsets should have average syndrome weight > 80."""
        np.random.seed(201)
        n_frames = 3
        offset = 10
        stream_llr, true_offset, _ = generate_stream(
            H, n_frames=n_frames, lam_s=10.0, lam_b=0.1, offset=offset
        )
        scores = syndrome_screen(stream_llr, n_offsets=N)
        # Collect syndrome weights for wrong offsets
        wrong_weights = [
            scores[o] for o in scores if o != true_offset
        ]
        avg_wrong = np.mean(wrong_weights)
        assert avg_wrong > 80, (
            f"Average wrong-offset syndrome weight {avg_wrong:.1f} is not > 80"
        )