"""Tests for Guppy code generation from SLR programs."""

from pecos.slr import CReg, If, Main, QReg, Repeat, SlrConverter
from pecos.slr.qeclib import qubit as qb


def test_simple_circuit() -> None:
    """Test simple quantum circuit generation."""
    prog = Main(
        q := QReg("q", 2),
        c := CReg("c", 2),
        # Prepare Bell state
        qb.H(q[0]),
        qb.CX(q[0], q[1]),
        # Measure
        qb.Measure(q) > c,
    )

    # Generate Guppy code
    guppy_code = SlrConverter(prog).guppy()

    # Basic assertions - AST codegen uses array parameters
    assert "@guppy" in guppy_code
    assert "def main(" in guppy_code
    assert "array[qubit," in guppy_code  # Array parameter type
    assert "quantum.h(" in guppy_code
    assert "quantum.cx(" in guppy_code


def test_conditional_logic() -> None:
    """Test conditional logic generation."""
    prog = Main(
        q := QReg("q", 1),
        c := CReg("c", 1),
        qb.H(q[0]),
        qb.Measure(q[0]) > c[0],
        If(c[0] == 1).Then(
            qb.X(q[0]),
        ),
    )

    guppy_code = SlrConverter(prog).guppy()

    # Check conditional structure - AST codegen preserves comparison expression
    assert "if (c_0 == 1):" in guppy_code
    assert "quantum.x(q[0])" in guppy_code


def test_repeat_loop() -> None:
    """Test repeat loop generation."""
    prog = Main(
        q := QReg("q", 1),
        # Apply H gate 3 times
        Repeat(3).block(
            qb.H(q[0]),
        ),
    )

    guppy_code = SlrConverter(prog).guppy()

    # Check loop structure
    assert "for _ in range(3):" in guppy_code
    assert "quantum.h(q[0])" in guppy_code


def test_steane_snippet() -> None:
    """Test a simple Steane code snippet."""
    from pecos.slr.qeclib.steane.steane_class import Steane

    prog = Main(
        # Create two logical qubits
        s1 := Steane("s1"),
        s2 := Steane("s2"),
        # Prepare logical |0>
        s1.pz(),
        # Apply logical Hadamard
        s1.h(),
        # Logical CNOT
        s1.cx(s2),
        # QEC cycle
        s1.qec(),
        s2.qec(),
    )

    guppy_code = SlrConverter(prog).guppy()

    # Check that Steane registers are declared as array parameters
    assert "s1_d: array[qubit, 7] @owned" in guppy_code
    assert "s2_d: array[qubit, 7] @owned" in guppy_code
    # Check that some quantum operations are present
    assert "quantum.h(" in guppy_code
    assert "quantum.cx(" in guppy_code


def test_measurement_handling() -> None:
    """Test different measurement patterns."""
    prog = Main(
        q := QReg("q", 3),
        c := CReg("c", 3),
        # Individual qubit measurement
        qb.Measure(q[0]) > c[0],
        # Full register measurement
        qb.Measure(q) > c,
    )

    guppy_code = SlrConverter(prog).guppy()

    # Check measurement generation - AST codegen unpacks registers
    assert "c_0 = quantum.measure(" in guppy_code
    # Full register measurement is expanded to individual measurements
    assert "c_1 = quantum.measure(" in guppy_code
    assert "c_2 = quantum.measure(" in guppy_code


def test_various_gates() -> None:
    """Test generation of various quantum gates."""
    prog = Main(
        q := QReg("q", 2),
        # Single-qubit gates
        qb.H(q[0]),
        qb.X(q[0]),
        qb.Y(q[0]),
        qb.Z(q[0]),
        qb.SZ(q[0]),  # S gate
        qb.SZdg(q[0]),  # Sdg gate
        qb.T(q[0]),
        qb.Tdg(q[0]),
        # Two-qubit gates
        qb.CX(q[0], q[1]),
        qb.CY(q[0], q[1]),
        qb.CZ(q[0], q[1]),
    )

    guppy_code = SlrConverter(prog).guppy()

    # Check all gates are present - AST codegen uses sz/szdg names
    gates = ["h", "x", "y", "z", "sz", "szdg", "t", "tdg", "cx", "cy", "cz"]
    for gate in gates:
        assert f"quantum.{gate}(" in guppy_code


def test_bitwise_operations() -> None:
    """Test generation of bitwise operations."""
    prog = Main(
        c := CReg("c", 8),
        # Initialize some bits
        c[0].set(1),
        c[1].set(0),
        c[2].set(1),
        # Test bitwise in assignments
        c[3].set(c[0] ^ c[1]),
        c[4].set(c[0] & c[2]),
        c[5].set(c[1] | c[2]),
        # Test NOT operation
        c[6].set(~c[0]),
        c[7].set((c[0] | c[1]) & ~c[2]),
    )

    guppy_code = SlrConverter(prog).guppy()

    # Check bitwise operations - AST codegen uses underscore naming for expressions
    # XOR uses ^ operator
    assert "c[3] = " in guppy_code
    assert "^" in guppy_code
    # AND uses 'and' operator
    assert "c[4] = " in guppy_code
    assert "and" in guppy_code
    # OR uses 'or' operator
    assert "c[5] = " in guppy_code
    assert "or" in guppy_code
    # NOT uses 'not' operator
    assert "c[6] = " in guppy_code
    assert "not" in guppy_code
    # Complex expression with multiple operators
    assert "c[7] = " in guppy_code


def test_register_operations() -> None:
    """Test operations on full quantum registers."""
    prog = Main(
        q := QReg("q", 4),
        _c := CReg("c", 4),
        # Apply gate to full register
        qb.H(q),
        # Apply gate to specific qubits
        qb.X(q[0]),
        qb.X(q[2]),
        # Two-qubit gates on specific pairs
        qb.CX(q[0], q[1]),
        qb.CX(q[2], q[3]),
    )

    guppy_code = SlrConverter(prog).guppy()

    # AST codegen unrolls register-wide operations
    assert "quantum.h(q[0])" in guppy_code
    assert "quantum.h(q[1])" in guppy_code
    assert "quantum.h(q[2])" in guppy_code
    assert "quantum.h(q[3])" in guppy_code

    # Check individual operations
    assert "quantum.x(q[0])" in guppy_code
    assert "quantum.x(q[2])" in guppy_code
    assert "quantum.cx(q[0], q[1])" in guppy_code


def test_steane_encoding_circuit_pattern() -> None:
    """Test the specific multi-pair CX pattern from Steane encoding circuit."""
    prog = Main(
        q := QReg("q", 7),
        # Prepare first 6 qubits
        qb.Prep(q[0], q[1], q[2], q[3], q[4], q[5]),
        # Single gates
        qb.CX(q[6], q[5]),
        qb.H(q[1]),
        qb.CX(q[1], q[0]),
        qb.H(q[2]),
        qb.CX(q[2], q[4]),
        qb.H(q[3]),
        # Multi-pair CX operations (Steane encoding pattern)
        qb.CX(
            (q[3], q[5]),
            (q[2], q[0]),
            (q[6], q[4]),
        ),
        qb.CX(
            (q[2], q[6]),
            (q[3], q[4]),
            (q[1], q[5]),
        ),
        qb.CX(
            (q[1], q[6]),
            (q[3], q[0]),
        ),
    )

    guppy_code = SlrConverter(prog).guppy()

    # AST codegen uses reset for Prep operations (unrolled, not looped)
    assert "quantum.reset(q[0])" in guppy_code
    assert "quantum.reset(q[1])" in guppy_code
    assert "quantum.reset(q[5])" in guppy_code

    # Check single CX operations
    assert "quantum.cx(q[6], q[5])" in guppy_code
    assert "quantum.cx(q[1], q[0])" in guppy_code
    assert "quantum.cx(q[2], q[4])" in guppy_code

    # Check multi-pair CX operations are expanded correctly
    assert "quantum.cx(q[3], q[5])" in guppy_code
    assert "quantum.cx(q[2], q[0])" in guppy_code
    assert "quantum.cx(q[6], q[4])" in guppy_code
    assert "quantum.cx(q[2], q[6])" in guppy_code
    assert "quantum.cx(q[3], q[4])" in guppy_code
    assert "quantum.cx(q[1], q[5])" in guppy_code
    assert "quantum.cx(q[1], q[6])" in guppy_code
    assert "quantum.cx(q[3], q[0])" in guppy_code


def test_reset_operations() -> None:
    """Test that Prep operations generate proper reset calls."""
    prog = Main(
        q := QReg("q", 3),
        _c := CReg("c", 3),
        # Reset single qubit
        qb.Prep(q[0]),
        # Apply some operations
        qb.H(q[0]),
        qb.CX(q[0], q[1]),
        # Reset multiple qubits
        qb.Prep(q[1], q[2]),
        # More operations
        qb.X(q[0]),
        qb.Y(q[1]),
        qb.Z(q[2]),
    )

    guppy_code = SlrConverter(prog).guppy()

    # AST codegen uses reset for Prep operations (not fresh allocation)
    assert "quantum.reset(q[0])" in guppy_code
    assert "quantum.reset(q[1])" in guppy_code
    assert "quantum.reset(q[2])" in guppy_code

    # Count reset occurrences - should be 3 (q[0], q[1], q[2])
    reset_count = guppy_code.count("quantum.reset(")
    assert reset_count == 3


def test_permute_operations() -> None:
    """Test that Permute operations generate proper swaps."""
    from pecos.slr.misc import Permute

    prog = Main(
        a := QReg("a", 3),
        b := QReg("b", 3),
        c := CReg("c", 2),
        d := CReg("d", 2),
        # Individual element permutation
        Permute([a[0], b[1]], [b[1], a[0]]),
        # Multiple element permutation (rotation)
        Permute([a[0], a[1], a[2]], [a[2], a[0], a[1]]),
        # Whole register permutation
        Permute(c, d),
        # Apply some gates to verify permutation
        qb.H(a[0]),
        qb.X(b[1]),
    )

    guppy_code = SlrConverter(prog).guppy()

    # Check that permutation operations are present
    # Note: The exact syntax may vary based on implementation
    assert "# Permute" in guppy_code or "swap" in guppy_code.lower() or ("a[0]" in guppy_code and "b[1]" in guppy_code)

    # Check that gates work after permutation
    assert "quantum.h(" in guppy_code
    assert "quantum.x(" in guppy_code
