"""Tests for fce.stabilizer_codes -- Steane code with real syndrome measurement."""

import numpy as np
import pytest

import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from fce.stabilizer_codes import (
    steane_stabilizers, steane_encode, steane_logical_operators,
    measure_syndrome, build_steane_decoder, decode_and_correct,
    apply_single_qubit_error, estimate_logical_error_rate,
    pauli_on_qubit, X, Y, Z, I2,
)


@pytest.fixture
def stabilizers():
    return steane_stabilizers()


@pytest.fixture
def decoder():
    return build_steane_decoder()


@pytest.fixture
def encoded_zero():
    return steane_encode(np.array([1.0, 0.0], dtype=complex))


@pytest.fixture
def encoded_plus():
    return steane_encode(np.array([1.0, 1.0], dtype=complex) / np.sqrt(2))


class TestSteaneStabilizers:
    def test_six_stabilizers(self, stabilizers):
        """Steane code has 6 stabilizer generators."""
        assert len(stabilizers) == 6

    def test_stabilizers_are_unitary(self, stabilizers):
        """Each stabilizer should be unitary (S S^dag = I)."""
        for i, S in enumerate(stabilizers):
            product = S @ S.conj().T
            assert np.allclose(product, np.eye(128), atol=1e-10), \
                f"Stabilizer {i} is not unitary"

    def test_stabilizers_are_hermitian(self, stabilizers):
        """Pauli stabilizers are Hermitian (S = S^dag)."""
        for i, S in enumerate(stabilizers):
            assert np.allclose(S, S.conj().T, atol=1e-10), \
                f"Stabilizer {i} is not Hermitian"

    def test_stabilizers_commute(self, stabilizers):
        """All stabilizers should commute: [S_i, S_j] = 0."""
        for i in range(len(stabilizers)):
            for j in range(i + 1, len(stabilizers)):
                comm = stabilizers[i] @ stabilizers[j] - stabilizers[j] @ stabilizers[i]
                assert np.allclose(comm, 0, atol=1e-10), \
                    f"Stabilizers {i} and {j} don't commute"


class TestSteaneEncode:
    def test_encoded_state_normalized(self, encoded_zero):
        """Encoded state should be normalized."""
        assert abs(np.linalg.norm(encoded_zero) - 1.0) < 1e-10

    def test_encoded_state_dimension(self, encoded_zero):
        """Encoded state should be 128-dimensional (2^7)."""
        assert encoded_zero.shape == (128,)

    def test_encoded_zero_is_stabilizer_eigenstate(self, encoded_zero, stabilizers):
        """Encoded |0_L> should be +1 eigenstate of all stabilizers."""
        for i, S in enumerate(stabilizers):
            expectation = np.real(encoded_zero.conj() @ S @ encoded_zero)
            assert abs(expectation - 1.0) < 1e-8, \
                f"|0_L> is not +1 eigenstate of stabilizer {i}: got {expectation}"

    def test_encoded_plus_is_stabilizer_eigenstate(self, encoded_plus, stabilizers):
        """Encoded |+_L> should also be +1 eigenstate of all stabilizers."""
        for i, S in enumerate(stabilizers):
            expectation = np.real(encoded_plus.conj() @ S @ encoded_plus)
            assert abs(expectation - 1.0) < 1e-8

    def test_logical_z_eigenvalue(self, encoded_zero):
        """Encoded |0_L> should be +1 eigenstate of Z_L."""
        _, Z_L = steane_logical_operators()
        expectation = np.real(encoded_zero.conj() @ Z_L @ encoded_zero)
        assert abs(expectation - 1.0) < 1e-8

    def test_logical_z_eigenvalue_one(self):
        """Encoded |1_L> should be -1 eigenstate of Z_L."""
        encoded_one = steane_encode(np.array([0.0, 1.0], dtype=complex))
        _, Z_L = steane_logical_operators()
        expectation = np.real(encoded_one.conj() @ Z_L @ encoded_one)
        assert abs(expectation - (-1.0)) < 1e-8


class TestSyndromeMeasurement:
    def test_no_error_zero_syndrome(self, encoded_zero, stabilizers):
        """No error -> all-zero syndrome."""
        syndrome = measure_syndrome(encoded_zero, stabilizers)
        assert np.all(syndrome == 0), \
            f"Expected zero syndrome for error-free state, got {syndrome}"

    def test_x_error_detected(self, encoded_zero, stabilizers):
        """Single X error should produce nonzero syndrome."""
        for qubit in range(7):
            errored = apply_single_qubit_error(encoded_zero, 'X', qubit)
            syndrome = measure_syndrome(errored, stabilizers)
            assert np.any(syndrome != 0), \
                f"X error on qubit {qubit} not detected"

    def test_z_error_detected(self, encoded_zero, stabilizers):
        """Single Z error should produce nonzero syndrome."""
        for qubit in range(7):
            errored = apply_single_qubit_error(encoded_zero, 'Z', qubit)
            syndrome = measure_syndrome(errored, stabilizers)
            assert np.any(syndrome != 0), \
                f"Z error on qubit {qubit} not detected"

    def test_y_error_detected(self, encoded_zero, stabilizers):
        """Single Y error should produce nonzero syndrome."""
        for qubit in range(7):
            errored = apply_single_qubit_error(encoded_zero, 'Y', qubit)
            syndrome = measure_syndrome(errored, stabilizers)
            assert np.any(syndrome != 0), \
                f"Y error on qubit {qubit} not detected"


class TestDecodeAndCorrect:
    def test_x_error_corrected(self, encoded_zero, stabilizers, decoder):
        """Single X error should be corrected."""
        for qubit in range(7):
            errored = apply_single_qubit_error(encoded_zero, 'X', qubit)
            syndrome = measure_syndrome(errored, stabilizers)
            corrected = decode_and_correct(errored, syndrome, decoder)

            # Should recover original (up to global phase)
            overlap = abs(encoded_zero.conj() @ corrected) ** 2
            assert overlap > 0.99, \
                f"X correction on qubit {qubit} failed: overlap={overlap}"

    def test_z_error_corrected(self, encoded_zero, stabilizers, decoder):
        """Single Z error should be corrected."""
        for qubit in range(7):
            errored = apply_single_qubit_error(encoded_zero, 'Z', qubit)
            syndrome = measure_syndrome(errored, stabilizers)
            corrected = decode_and_correct(errored, syndrome, decoder)

            overlap = abs(encoded_zero.conj() @ corrected) ** 2
            assert overlap > 0.99, \
                f"Z correction on qubit {qubit} failed: overlap={overlap}"

    def test_no_error_no_change(self, encoded_zero, stabilizers, decoder):
        """Zero syndrome should not modify the state."""
        syndrome = measure_syndrome(encoded_zero, stabilizers)
        corrected = decode_and_correct(encoded_zero, syndrome, decoder)
        assert np.allclose(corrected, encoded_zero, atol=1e-10)


class TestLogicalErrorRate:
    def test_zero_physical_error_zero_logical(self):
        """Zero physical error rate -> zero logical error rate."""
        rate = estimate_logical_error_rate(
            physical_error_rate=0.0, n_trials=100,
            rng=np.random.default_rng(42),
        )
        assert rate == 0.0

    def test_low_error_below_threshold(self):
        """For p << threshold, logical error rate should be very low."""
        rate = estimate_logical_error_rate(
            physical_error_rate=0.001, n_trials=500,
            rng=np.random.default_rng(42),
        )
        # At p=0.001, logical error rate should be very small
        assert rate < 0.05, f"Logical error rate too high at p=0.001: {rate}"

    def test_high_error_above_threshold(self):
        """For p >> threshold, logical error rate should be high."""
        rate = estimate_logical_error_rate(
            physical_error_rate=0.3, n_trials=200,
            rng=np.random.default_rng(42),
        )
        # At p=0.3, code can't help much
        assert rate > 0.01, f"Logical error rate suspiciously low at p=0.3: {rate}"

    def test_error_rate_is_computed(self):
        """Different physical rates should give different logical rates."""
        rate_low = estimate_logical_error_rate(
            physical_error_rate=0.001, n_trials=500,
            rng=np.random.default_rng(42),
        )
        rate_high = estimate_logical_error_rate(
            physical_error_rate=0.1, n_trials=500,
            rng=np.random.default_rng(42),
        )
        assert rate_high >= rate_low, \
            "Higher physical error should give higher logical error"