import random
import pytest
import numpy as np
import jax.numpy as jnp
import pyzx_param as pyzx
from pytest_lazyfixture import lazy_fixture as lf
import tensorcircuit as tc
from tensorcircuit.zx.stabilizertcircuit import StabilizerTCircuit
from tensorcircuit.zx.scalar_graph import (
    find_stab,
    compile_scalar_graphs,
)
from tensorcircuit.zx.converter import circuit_to_zx, is_pauli

# --- Helpers ---


def get_tc_unitary(c):
    return tc.backend.numpy(c.matrix())


def get_zx_unitary(g, n):
    if hasattr(g, "graph"):
        g = g.graph
    zx_tensor = g.to_tensor()
    return zx_tensor.reshape((2**n, 2**n))


def assert_unitary_match(g, c, atol=1e-5):
    n = c._nqubits
    zx_mat = get_zx_unitary(g, n)
    tc_mat = get_tc_unitary(c)

    mask = np.abs(tc_mat) > 1e-5
    if not np.any(mask):
        np.testing.assert_allclose(zx_mat, tc_mat, atol=atol)
        return

    first_idx = np.where(mask)[0][0], np.where(mask)[1][0]
    ratio = zx_mat[first_idx] / tc_mat[first_idx]

    # Check proportionality
    np.testing.assert_allclose(zx_mat, ratio * tc_mat, atol=atol)
    # Check if ratio magnitude is a power of sqrt(2) (unnormalized ZX) or 1
    mag = np.abs(ratio)
    log_ratio = 2 * np.log2(mag)
    np.testing.assert_allclose(
        log_ratio,
        np.round(log_ratio),
        atol=1e-5,
        rtol=0.0,
        err_msg=f"Ratio {ratio} is not a power of sqrt(2)",
    )


# --- New simple sanity tests ---


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_simple_cases(backend):
    batch = 5000
    cases = [
        ("X_ERROR(0.1)", lambda c: c.x_error(0, 0.1), 0.1),
        ("X gate", lambda c: c.x(0), 1.0),
        ("Z gate + H", lambda c: (c.z(0), c.h(0)), 0.5),
        ("Z_ERROR(0.2) + H", lambda c: (c.z_error(0, 0.2), c.h(0)), 0.5),
        ("Y_ERROR(0.3)", lambda c: c.y_error(0, 0.3), 0.3),
        ("X + RESET", lambda c: (c.x(0), c.reset_instruction(0)), 0.0),
        ("H + M", lambda c: c.h(0), 0.5),
    ]

    for name, setup, expected in cases:
        c = StabilizerTCircuit(1)
        setup(c)
        c.measure_instruction(0)
        stc = StabilizerTCircuit.from_circuit(c)
        res = stc.sample_measurements(shots=batch)
        prob = np.mean(res)
        np.testing.assert_allclose(
            prob,
            expected,
            atol=0.05,
            rtol=0.0,
            err_msg=f"Case {name} failed: got {prob}, exp {expected}",
        )


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_detector_xor_simple(backend):
    batch = 5000
    c = StabilizerTCircuit(1)
    c.x_error(0, 0.1)
    c.measure_instruction(0)  # m0
    c.x_error(0, 0.2)
    c.measure_instruction(0)  # m1
    c.detector_instruction([-1, -2])  # m0 ^ m1
    stc = StabilizerTCircuit.from_circuit(c)
    res = stc.sample_detectors(shots=batch)
    prob = np.mean(res)
    np.testing.assert_allclose(prob, 0.2, atol=0.05, rtol=0.0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_depolarize2_simple(backend):
    batch = 5000
    c = StabilizerTCircuit(2)
    c.depolarizing2(0, 1, 0.1)
    c.measure_instruction(0)
    c.measure_instruction(1)
    stc = StabilizerTCircuit.from_circuit(c)
    res = stc.sample_measurements(shots=batch)
    p11 = np.mean(res[:, 0] & res[:, 1])
    # Theory: 3/15 * 0.1 = 0.02
    np.testing.assert_allclose(p11, 0.02, atol=0.02, rtol=0.0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_pauli_h_simple(backend):
    batch = 5000
    c = StabilizerTCircuit(1)
    c.pauli_instruction(0, pz=0.1)
    c.h(0)
    c.measure_instruction(0)
    stc = StabilizerTCircuit.from_circuit(c)
    res = stc.sample_measurements(shots=batch)
    prob = np.mean(res)
    np.testing.assert_allclose(prob, 0.5, atol=0.05, rtol=0.0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_outcome_probability(backend):
    # |0> state, prob(0) = 1.0, prob(1) = 0
    stc = StabilizerTCircuit(1)
    stc.measure_instruction(0)

    # outcome_probability expects a state matching the measurements
    # For a simple circuit with 1 measure, state is [0] or [1]
    p0 = stc.outcome_probability(jnp.array([0]))
    p1 = stc.outcome_probability(jnp.array([1]))
    assert np.allclose(p0, 1.0)
    assert np.allclose(p1, 0.0)

    # H gate, prob(0) = 0.5, prob(1) = 0.5
    stc = StabilizerTCircuit(1)
    stc.h(0)
    stc.measure_instruction(0)
    p0 = stc.outcome_probability(jnp.array([0]))
    p1 = stc.outcome_probability(jnp.array([1]))
    assert np.allclose(p0, 0.5)
    assert np.allclose(p1, 0.5)

    # X_ERROR(0.1), prob(1) = 0.1
    stc = StabilizerTCircuit(1)
    stc.x_error(0, 0.1)
    stc.measure_instruction(0)
    p1 = stc.outcome_probability(jnp.array([1]), shots=10)
    assert p1.shape == (10,)
    # For a specific error realization, it's either 0 or 1.
    # But since x_error adds a parameter 'e0', and outcome_probability
    # computes P(state | error_i), it should correctly reflect the error bit.
    # Actually, outcome_probability in tsim computes P(state | error).
    # If the error bit e0=1, then P(1|e0=1) = 1.0. If e0=0, P(1|e0=0) = 0.
    assert np.all(jnp.logical_or(jnp.isclose(p1, 0.0), jnp.isclose(p1, 1.0)))
    # Average should be around 0.1
    np.testing.assert_allclose(
        np.mean(p1), 0.1, atol=0.3, rtol=0.0
    )  # low shots, just checking shape/range


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_sample_detectors_advanced(backend):
    # Test separate_observables and use_reference
    c = StabilizerTCircuit(2)
    c.x(0)
    c.measure_instruction(0)  # m0 = 1
    c.measure_instruction(1)  # m1 = 0
    c.detector_instruction([0])  # det0 = m0 = 1
    c.observable_instruction([1], observable_index=0)  # obs0 = m1 = 0

    stc = StabilizerTCircuit.from_circuit(c)

    # 1. Raw
    samples = stc.sample_detectors(shots=10)
    assert samples.shape == (10, 2)
    assert np.all(samples[:, 0] == 1)
    assert np.all(samples[:, 1] == 0)

    # 2. Separate
    det, obs = stc.sample_detectors(shots=10, separate_observables=True)
    assert det.shape == (10, 1)
    assert obs.shape == (10, 1)
    assert np.all(det == 1)
    assert np.all(obs == 0)

    # 3. Reference (XOR with noiseless)
    # Noiseless outcome is [1, 0]. XORing [1, 0] with [1, 0] gives [0, 0]
    samples_ref = stc.sample_detectors(shots=10, use_reference=True)
    assert np.all(samples_ref == 0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_multi_component_sampling(backend):
    # Two disconnected qubits
    stc = StabilizerTCircuit(2)
    stc.x(0)
    stc.h(1)
    stc.measure_instruction(0)
    stc.measure_instruction(1)

    samples = stc.sample_measurements(shots=1000)
    assert samples.shape == (1000, 2)
    assert np.all(samples[:, 0] == 1)
    np.testing.assert_allclose(np.mean(samples[:, 1]), 0.5, atol=0.1, rtol=0.0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_noisy_measurements(backend):
    # M(p) is equivalent to X_ERROR(p) + M
    batch = 50000
    p = 0.1

    # Case 1: M(p)
    stc1 = StabilizerTCircuit(1)
    # Note: MR(p) or M(p) usually handled via parameters
    stc1._qir.append({"name": "M", "index": [0], "parameters": {"p": p}})
    res1 = stc1.sample_measurements(shots=batch)
    prob1 = np.mean(res1)
    np.testing.assert_allclose(prob1, p, atol=0.05, rtol=0.0)

    # Case 2: MR(p)
    stc2 = StabilizerTCircuit(1)
    stc2.x(0)
    stc2._qir.append({"name": "MR", "index": [0], "parameters": {"p": p}})
    stc2.measure_instruction(0)  # Should be 0 after MR
    res2 = stc2.sample_measurements(shots=batch)
    # res2 shape (batch, 2)
    m0 = np.mean(res2[:, 0])  # result of MR
    m1 = np.mean(res2[:, 1])  # result of following M
    # tsim's MR(p) applies noise twice: once in mr, once in m.
    # Total flip probability = p(1-p) + (1-p)p = 2p - 2p^2
    # For p=0.1, flip = 0.18. Since we start at 1 (X gate), m0 should be 1-0.18 = 0.82
    expected_m0 = 1 - (2 * p - 2 * p**2)
    np.testing.assert_allclose(m0, expected_m0, atol=0.05, rtol=0.0)

    np.testing.assert_allclose(m1, 0.0, atol=0.01, rtol=0.0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_batch_size_consistency(backend):
    stc = StabilizerTCircuit(1)
    stc.h(0)
    stc.measure_instruction(0)

    # Small batch size vs large batch size
    res1 = stc.sample_measurements(shots=100, batch_size=10)
    res2 = stc.sample_measurements(shots=100, batch_size=100)
    assert res1.shape == (100, 1)
    assert res2.shape == (100, 1)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_empty_circuit(backend):
    stc = StabilizerTCircuit(1)
    # No gates, no measurements
    # sample_measurements defaults to force_measure_all=True if no measurements exist
    res = stc.sample_measurements(shots=10)
    assert res.shape == (10, 1)
    assert np.all(res == 0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_observable_include_logical(backend):
    stc = StabilizerTCircuit(1)
    stc.x(0)
    stc.measure_instruction(0)
    stc.observable_instruction([0], observable_index=5)

    det, obs = stc.sample_detectors(shots=10, separate_observables=True)
    # No detectors, only 1 observable at index 5
    assert det.shape == (10, 0)
    assert obs.shape == (10, 1)
    assert np.all(obs == 1)


@pytest.mark.parametrize("backend", [lf("npb")])
def test_converter_all_single_qubit_gates(backend):
    # Test every single qubit gate supported
    for g_name in ["h", "s", "sd", "t", "td", "x", "y", "z", "i"]:
        c = tc.Circuit(1)
        getattr(c, g_name)(0)
        g = circuit_to_zx(c)
        assert_unitary_match(g, c), f"Gate {g_name} failed"


@pytest.mark.parametrize("backend", [lf("npb")])
def test_converter_all_two_qubit_gates(backend):
    for g_name in ["cnot", "cz", "swap"]:
        c = tc.Circuit(2)
        getattr(c, g_name)(0, 1)
        g = circuit_to_zx(c)
        assert_unitary_match(g, c), f"Gate {g_name} failed"


@pytest.mark.parametrize("backend", [lf("npb")])
def test_is_pauli(backend):
    _ = is_pauli(tc.gates.x().tensor)

    assert is_pauli(tc.gates.x().tensor) == "x"
    assert is_pauli(tc.gates.y().tensor) == "y"
    assert is_pauli(tc.gates.z().tensor) == "z"
    assert is_pauli(tc.gates.i().tensor) == "i"
    assert is_pauli(np.eye(2)) == "i"
    # Non-Pauli
    assert is_pauli(tc.gates.h().tensor) is None


@pytest.mark.parametrize("backend", [lf("npb")])
def test_converter_complex_circuit(backend):
    c = tc.Circuit(3)
    c.h(0)
    c.cnot(0, 1)
    c.t(1)
    c.s(2)
    c.cz(1, 2)
    c.h(2)
    g = circuit_to_zx(c)
    assert_unitary_match(g, c)


@pytest.mark.parametrize("backend", [lf("npb")])
def test_decomposition_scaling(backend):
    for tc_count in range(4):
        c = tc.Circuit(1)
        for _ in range(tc_count):
            c.t(0)
        g = circuit_to_zx(c)
        graphs = find_stab(g, strategy="cat5")
        assert len(graphs) >= 1
        for sg in graphs:
            assert all(
                sg.type(v)
                in [pyzx.VertexType.BOUNDARY, pyzx.VertexType.Z, pyzx.VertexType.X]
                for v in sg.vertices()
            )


@pytest.mark.parametrize("backend", [lf("npb")])
def test_scalar_ir_variable_mapping(backend):
    c = tc.Circuit(1)
    c.h(0)
    c.measure_instruction(0)
    g = circuit_to_zx(c)
    if hasattr(g, "graph"):
        g = g.graph
    # Initialize initial boundaries to X state (|0> state)
    for v in g.inputs():
        if g.type(v) == pyzx.VertexType.BOUNDARY:
            g.set_type(v, pyzx.VertexType.X)
            g.set_phase(v, 0)
    # Plug outputs
    g.apply_effect("0")  # or any value

    graphs = find_stab(g, strategy="cat5")

    params = ["rec[0]"]
    data = compile_scalar_graphs(graphs, params)
    assert data.a_param_bits is not None


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_stabilizert_circuit_methods(backend):
    stc = StabilizerTCircuit(2)
    stc.h(0)
    stc.cnot(0, 1)
    stc.t(1)
    assert len(stc._qir) == 3


@pytest.mark.parametrize("backend", [lf("npb")])
def test_zx_detector_chain(backend):
    c = StabilizerTCircuit(2)
    c.measure_instruction(0)
    c.measure_instruction(1)
    c.detector_instruction([0, 1])

    g = circuit_to_zx(c)
    if hasattr(g, "graph"):
        g = g.graph

    # Find the detector vertex
    det_v = -1
    for v in g.vertices():
        if "det[0]" in g._phaseVars.get(v, set()):
            det_v = v
            break
    assert det_v != -1
    # Should be connected to two record spiders
    neighbors = list(g.neighbors(det_v))
    rec_neighbors = 0
    for n in neighbors:
        if any(s.startswith("rec[") for s in g._phaseVars.get(n, set())):
            rec_neighbors += 1
    assert rec_neighbors == 2


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_from_circuit_with_noise(backend):
    c = tc.Circuit(1)
    c.h(0)
    c.depolarizing(0, px=0.1, py=0.1, pz=0.1)

    stc = StabilizerTCircuit.from_circuit(c)
    assert len(stc._qir) >= 2

    g = circuit_to_zx(stc)
    if hasattr(g, "graph"):
        g = g.graph

    has_ex = False
    has_ez = False
    for v in g.vertices():
        labels = g._phaseVars.get(v, set())
        if any(s.startswith("e") for s in labels):
            has_ex = True
            has_ez = True
    assert has_ex and has_ez


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_sample_detectors_fixed_gate_sequence(backend):
    stc = StabilizerTCircuit(2)
    stc.x(0)
    stc.measure_instruction(0)
    stc.measure_instruction(1)
    stc.detector_instruction([0, 1])

    samples = stc.sample_detectors(shots=32, seed=1234)
    assert samples.shape == (32, 1)
    np.testing.assert_array_equal(samples[:, 0], np.ones(32, dtype=bool))


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_sample_detectors_multiple_detectors_complex_sequence(backend):
    stc = StabilizerTCircuit(3)
    stc.x(0)
    stc.cnot(0, 1)
    stc.x(2)
    stc.measure_instruction(0)  # rec[0] = 1
    stc.measure_instruction(1)  # rec[1] = 1
    stc.measure_instruction(2)  # rec[2] = 1
    stc.detector_instruction([0, 1])  # 1 xor 1 = 0
    stc.detector_instruction([1, 2])  # 1 xor 1 = 0
    stc.detector_instruction([0, 2])  # 1 xor 1 = 0
    stc.detector_instruction([0])  # 1

    samples = stc.sample_detectors(shots=64, seed=101)
    assert samples.shape == (64, 4)
    expected = np.array([0, 0, 0, 1], dtype=bool)
    np.testing.assert_array_equal(samples, np.tile(expected, (64, 1)))


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_stabilizertcircuit_sample_measurements_final_state_deterministic(backend):
    stc = StabilizerTCircuit(2)
    stc.x(0)
    stc.cnot(0, 1)
    stc.measure_instruction(0)
    stc.measure_instruction(1)
    samples = stc.sample_measurements(shots=64, seed=9)
    assert samples.shape == (64, 2)
    expected = np.ones((64, 2), dtype=bool)
    np.testing.assert_array_equal(samples, expected)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_stabilizertcircuit_sample_with_t_gate_statistics(backend):
    stc = StabilizerTCircuit(1)
    stc.h(0)
    stc.t(0)
    stc.h(0)
    stc.measure_instruction(0)

    samples = stc.sample_measurements(shots=8000, seed=12345)
    assert samples.shape == (8000, 1)
    p1 = np.mean(samples[:, 0])
    p1_theory = 0.5 * (1.0 - np.cos(np.pi / 4.0))
    np.testing.assert_allclose(p1, p1_theory, atol=0.03, rtol=0.0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_stabilizertcircuit_sample_40q_few_t_scalable(backend):
    stc = StabilizerTCircuit(40)
    stc.h(0)
    stc.t(0)
    stc.h(0)
    stc.h(7)
    stc.t(7)
    stc.h(7)
    stc.h(21)
    stc.t(21)
    stc.h(21)
    for i in [0, 7, 21]:
        stc.measure_instruction(i)

    samples = stc.sample_measurements(shots=1000, seed=20260402)
    assert samples.shape == (1000, 3)

    p_theory = (2 - np.sqrt(2)) / 4
    for i in range(3):
        p_emp = np.mean(samples[:, i])
        stderr = np.sqrt(p_theory * (1 - p_theory) / 1000)
        assert abs(p_emp - p_theory) < 5 * stderr


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_stabilizertcircuit_mirror_all_zero(backend):
    n = 5
    stc = StabilizerTCircuit(n)
    # Forward
    stc.h(0)
    stc.cnot(0, 1)
    stc.t(1)
    stc.s(2)
    stc.cnot(2, 3)
    stc.h(4)
    # Backward (Adjoint)
    stc.h(4)
    stc.cnot(2, 3)
    stc.sd(2)
    stc.td(1)
    stc.cnot(0, 1)
    stc.h(0)
    for i in range(n):
        stc.measure_instruction(i)

    batch = 100
    samples = stc.sample_measurements(shots=batch, seed=42)
    assert samples.shape == (batch, n)
    assert np.all(samples == 0)


# --- Stim Import & Integration Tests (Borrowed from tsim) ---


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_from_stim_bell_state(backend):

    stim_str = """
    R 0 1
    H 0
    CNOT 0 1
    M 0 1
    DETECTOR rec[-1] rec[-2]
    """
    stc = StabilizerTCircuit.from_stim_str(stim_str)
    # Bell state: m0 and m1 should be correlated, so detector (m0^m1) should be 0
    res = stc.sample_detectors(shots=100)
    assert np.all(res == 0)

    m_res = stc.sample_measurements(shots=100)
    assert np.all(m_res[:, 0] == m_res[:, 1])


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_from_stim_t_gate(backend):
    # Tests T gate decomposition and sampling statistics
    stim_str = """
    RX 0
    S 0
    H 0
    M 0
    """
    # S on |+> gives |+i> = (|0>+i|1>)/sqrt(2)
    # H on |+i> gives (|0>+|1> + i|0>-i|1>)/2 = (1+i)/2 |0> + (1-i)/2 |1>
    # Prob(0) = |(1+i)/2|^2 = 2/4 = 0.5
    stc = StabilizerTCircuit.from_stim_str(stim_str)
    res = stc.sample_measurements(shots=1000)
    mean_val = np.mean(res)
    assert 0.4 < mean_val < 0.6


@pytest.mark.parametrize("basis", ["X", "Y", "Z"])
@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_from_stim_reset_after_state_change(backend, basis):
    reset_gate = "R" if basis == "Z" else f"R{basis}"
    measure_gate = "M" if basis == "Z" else f"M{basis}"
    stim_str = f"""
    H 0
    S 0
    {reset_gate} 0
    {measure_gate} 0
    """
    stc = StabilizerTCircuit.from_stim_str(stim_str)
    res = stc.sample_measurements(shots=100)
    assert np.all(res == 0)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_from_stim_singlet_state(backend):
    # Borrowed from tsim: test_singlet_state
    # High-level check for CNOT, H, X, Z
    stim_str = """
    R 0 1
    X 0
    H 1
    CNOT 1 0
    Z 0
    M 0 1
    """
    stc = StabilizerTCircuit.from_stim_str(stim_str)
    res = stc.sample_measurements(shots=100)
    # Singlet state (|01> - |10>)/sqrt(2) has anti-correlated Z outcomes
    assert np.all(res[:, 0] != res[:, 1])


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_from_stim_noisy_bell_detector(backend):
    # Surface code style check: detector should flip if error occurs
    stim_str = """
    H 0
    CNOT 0 1
    X_ERROR(0.1) 0
    M 0 1
    DETECTOR rec[-1] rec[-2]
    """
    stc = StabilizerTCircuit.from_stim_str(stim_str)
    res = stc.sample_detectors(shots=10000)
    mean_flip = np.mean(res)
    # Error on q0 with prob 0.1 -> m0 flips -> detector flips
    assert 0.07 < mean_flip < 0.13


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_vs_stim_random_clifford(backend):
    import stim

    num_qubits = 4
    num_gates = 20
    c = stim.Circuit()
    gates = ["H", "S", "CX", "CZ", "X", "Y", "Z", "SQRT_X", "SWAP"]
    for _ in range(num_gates):
        g = random.choice(gates)
        if g in ["CX", "CZ", "SWAP"]:
            q = random.sample(range(num_qubits), 2)
            c.append(g, q)
        else:
            q = random.randrange(num_qubits)
            c.append(g, q)
    c.append("M", range(num_qubits))

    # Reference from Stim
    sampler_stim = c.compile_sampler()
    batch = 10000
    samples_stim = sampler_stim.sample(batch)
    p_stim = np.mean(samples_stim, axis=0)

    # Result from StabilizerTCircuit
    stc = StabilizerTCircuit.from_stim_circuit(c)
    samples_tc = stc.sample_measurements(shots=batch)
    p_tc = np.mean(samples_tc, axis=0)

    # Statistical check
    diff = np.abs(p_stim - p_tc)
    threshold = 5 * np.sqrt(0.5 * 0.5 / batch)
    assert np.all(diff < threshold)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_vs_stim_noisy_circuit(backend):
    import stim

    stim_str = """
    H 0
    CNOT 0 1
    X_ERROR(0.1) 0
    Y_ERROR(0.1) 1
    DEPOLARIZE1(0.1) 0
    M 0 1
    """
    stim_circuit = stim.Circuit(stim_str)

    # Reference
    sampler_stim = stim_circuit.compile_sampler()
    batch = 10000
    samples_stim = sampler_stim.sample(batch)
    p_stim = np.mean(samples_stim, axis=0)

    # TC
    stc = StabilizerTCircuit.from_stim_circuit(stim_circuit)
    samples_tc = stc.sample_measurements(shots=batch)
    p_tc = np.mean(samples_tc, axis=0)

    # Statistical check
    diff = np.abs(p_stim - p_tc)
    threshold = 5 * np.sqrt(0.5 * 0.5 / batch)
    assert np.all(diff < threshold)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_outcome_probability_rigorous(backend):
    # Test complex amplitudes via outcome_probability
    stc = StabilizerTCircuit(2)
    stc.h(0)
    stc.cnot(0, 1)
    # Bell state |00>+|11> / sqrt(2)
    stc.measure_instruction(0)
    stc.measure_instruction(1)

    # P(00) = 0.5, P(11) = 0.5, P(01) = 0, P(10) = 0
    assert np.allclose(stc.outcome_probability(jnp.array([0, 0])), 0.5)
    assert np.allclose(stc.outcome_probability(jnp.array([1, 1])), 0.5)
    assert np.allclose(stc.outcome_probability(jnp.array([0, 1])), 0.0)
    assert np.allclose(stc.outcome_probability(jnp.array([1, 0])), 0.0)

    # Add T gate
    stc2 = StabilizerTCircuit(1)
    stc2.h(0)
    stc2.t(0)
    stc2.h(0)
    stc2.measure_instruction(0)
    expected_p0 = (2 + np.sqrt(2)) / 4
    assert np.allclose(stc2.outcome_probability(jnp.array([0])), expected_p0, atol=1e-5)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_outcome_probability_vs_exact_10q(backend):
    # 10 qubit circuit with some Clifford and T gates
    n = 10
    c = tc.Circuit(n)
    for i in range(n):
        c.h(i)
    for i in range(n - 1):
        c.cnot(i, i + 1)
    # Add some T gates to make it interesting
    c.t(2)
    c.t(5)
    c.cnot(2, 4)
    c.h(4)
    c.cnot(8, 5)
    c.t(8)
    c.s(8)

    # More Cliffords
    c.s(3)
    c.cz(4, 8)
    c.h(1)

    # We must add measurements to define the record indices for outcome_probability
    for i in range(n):
        c.measure_instruction(i)

    stc = StabilizerTCircuit.from_circuit(c)

    # Select a few bitstrings
    test_bitstrings = [
        [0] * n,
        [1] * n,
        [0, 1] * 5,
        [1, 0, 1, 0, 0, 1, 1, 0, 0, 1],
    ]

    for bs in test_bitstrings:
        state = jnp.array(bs)
        prob_zx = stc.outcome_probability(state)

        # Exact result from amplitude
        amp = c.amplitude("".join(map(str, bs)))
        prob_exact = tc.backend.abs(amp) ** 2

        np.testing.assert_allclose(prob_zx, prob_exact, atol=1e-5)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_high_precision(backend, highp):
    # Verify that complex128 and int64 are active
    assert tc.dtypestr == "complex128"
    assert tc.idtypestr == "int64"

    # Test with a case where high precision might matter
    stc = StabilizerTCircuit(1)
    stc.h(0)
    stc.t(0)
    stc.h(0)

    # Use outcome_probability for a more precise check
    p0 = stc.outcome_probability(jnp.array([0]))
    expected_p0 = (2 + np.sqrt(2)) / 4

    # Check that the result is indeed float64
    assert p0.dtype == jnp.float64
    np.testing.assert_allclose(p0, expected_p0, atol=1e-12)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_numerical_probability_sum(backend):
    """Verify that sum of outcome probabilities equals 1.0."""
    n = 2

    # --- Part 1: Noiseless (Strict Identity) ---
    stc_none = StabilizerTCircuit(n, seed=42)
    stc_none.h(0)
    stc_none.cnot(0, 1)
    stc_none.t(1)
    for i in range(n):
        stc_none.measure_instruction(i)

    total_p_none = 0.0
    for i in range(2**n):
        state = jnp.array([(i >> j) & 1 for j in range(n)])
        # Noiseless case should be deterministic, shots=1 is sufficient
        p = stc_none.outcome_probability(state, shots=1)
        total_p_none += jnp.sum(p)

    # Must be exactly 1.0 within float32 precision
    np.testing.assert_allclose(total_p_none, 1.0, atol=1e-6)

    # --- Part 2: Noisy (Statistical Consistency) ---
    stc_noise = StabilizerTCircuit(n, seed=42)
    stc_noise.h(0)
    stc_noise.cnot(0, 1)
    stc_noise.t(1)
    stc_noise.depolarizing(0, 0.05)
    stc_noise.x_error(1, 0.02)
    for i in range(n):
        stc_noise.measure_instruction(i)

    total_p_noise = 0.0
    shots = 4096  # Higher shots for better statistics
    for i in range(2**n):
        state = jnp.array([(i >> j) & 1 for j in range(n)])
        p = stc_noise.outcome_probability(state, shots=shots)
        total_p_noise += jnp.mean(p)

    # Statistical sum should be 1.0 within 3-sigma (approx 0.02 for 4096 shots)
    np.testing.assert_allclose(total_p_noise, 1.0, atol=0.03)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_autoregressive_consistency(backend):
    """Check consistency between sample_measurements and outcome_probability."""
    stc = StabilizerTCircuit(1)
    stc.h(0)
    stc.t(0)
    stc.h(0)
    stc.measure_instruction(0)

    # 1. Frequency from sampler
    shots = 20000
    samples = stc.sample_measurements(shots=shots, seed=42)
    freq1 = jnp.mean(samples)

    # 2. Probability from API
    prob1 = stc.outcome_probability(jnp.array([1]))

    # Compare
    stderr = jnp.sqrt(prob1 * (1 - prob1) / shots)
    assert abs(freq1 - prob1) < 5 * stderr


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_stabilizer_rank_6(backend):
    """Test a circuit with T-count=6 (decomposes to 7 stabilizer states in cat5)."""
    n = 3
    stc = StabilizerTCircuit(n)
    for i in range(6):
        stc.h(i % n)
        stc.t(i % n)
    stc.measure_instruction(0)
    stc.measure_instruction(1)
    stc.measure_instruction(2)

    # Compare with exact amplitude from statevector
    c = tc.Circuit(n)
    for i in range(6):
        c.h(i % n)
        c.t(i % n)

    state = jnp.array([1, 0, 1])
    prob_zx = stc.outcome_probability(state)

    amp = c.amplitude("101")
    prob_exact = tc.backend.abs(amp) ** 2

    np.testing.assert_allclose(prob_zx, prob_exact, atol=1e-5)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_exact_scalar_alignment(backend):
    """Test complex circuit that triggers recursive reduction in ExactScalarArray."""
    # A sequence of T gates and Hadamards can lead to large coefficients
    # and deep power reductions.
    n = 1
    stc = StabilizerTCircuit(n)
    for _ in range(12):
        stc.h(0)
        stc.t(0)
    stc.measure_instruction(0)

    p0 = stc.outcome_probability(jnp.array([0]))

    # Exact statevector reference
    c = tc.Circuit(n)
    for _ in range(12):
        c.h(0)
        c.t(0)
    p0_ref = tc.backend.abs(c.amplitude("0")) ** 2

    np.testing.assert_allclose(p0, p0_ref, atol=1e-6)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_amplitude_expectation_random(backend):
    n = 3
    for _ in range(4):
        c = tc.Circuit(n)
        for i in range(n):
            c.h(i)
        for i in range(n - 1):
            c.cnot(i, i + 1)
        # Add some T gates
        c.t(0)
        c.t(1)
        c.cnot(0, 2)
        c.s(2)
        c.t(2)

        stc = StabilizerTCircuit.from_circuit(c)

        # Test amplitude
        bitstrings = ["000", "011", "101", "110"]
        for bs in bitstrings:
            amp_zx = stc.amplitude(bs)
            amp_exact = c.amplitude(bs)
            np.testing.assert_allclose(amp_zx, amp_exact, atol=1e-5)

        # Test expectation_ps
        # Z0
        exp_zx = stc.expectation_ps(z=[0])
        exp_exact = c.expectation_ps(z=[0])
        np.testing.assert_allclose(exp_zx, exp_exact, atol=1e-5)

        # X1 Z2
        exp_zx = stc.expectation_ps(x=[1], z=[2])
        exp_exact = c.expectation_ps(x=[1], z=[2])
        np.testing.assert_allclose(exp_zx, exp_exact, atol=1e-5)

        # Y0 Y1
        exp_zx = stc.expectation_ps(y=[0, 1])
        exp_exact = c.expectation_ps(y=[0, 1])
        np.testing.assert_allclose(exp_zx, exp_exact, atol=1e-5)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_expectation_amplitude_clean_diverse(backend):
    # More comprehensive clean circuit tests against tc.Circuit
    n = 4
    c = tc.Circuit(n)
    c.x(0)
    c.h(1)
    c.y(2)
    c.z(3)
    c.t(0)
    c.cnot(0, 1)
    c.s(2)
    c.cz(2, 3)
    c.swap(0, 2)
    c.h(3)
    c.t(1)  # Non-Clifford gate
    c.cnot(1, 3)

    stc = StabilizerTCircuit.from_circuit(c)

    # 1. Test amplitudes
    for bs_int in range(2**n):
        bs = bin(bs_int)[2:].zfill(n)
        amp_zx = stc.amplitude(bs)
        amp_exact = c.amplitude(bs)
        np.testing.assert_allclose(amp_zx, amp_exact, atol=1e-5)

    # 2. Test diverse Pauli expectations
    test_cases = [
        ({"z": [0, 1, 2, 3]}, "ZZZZ"),
        ({"x": [0, 2]}, "X0 X2"),
        ({"y": [1, 3]}, "Y1 Y3"),
        ({"x": [0], "y": [1], "z": [2]}, "X0 Y1 Z2"),
        ({"ps": [1, 2, 3, 0]}, "XYZI"),
    ]

    for kwargs, _ in test_cases:
        exp_zx = stc.expectation_ps(**kwargs)
        exp_exact = c.expectation_ps(**kwargs)
        np.testing.assert_allclose(exp_zx, exp_exact, atol=1e-5)

    # Canonical ps API should agree with equivalent explicit x/y/z specification.
    ps_case = [1, 2, 3, 0]
    xyz_case = {"x": [0], "y": [1], "z": [2]}
    np.testing.assert_allclose(
        stc.expectation_ps(ps=ps_case),
        stc.expectation_ps(**xyz_case),
        atol=1e-5,
    )
    np.testing.assert_allclose(
        c.expectation_ps(ps=ps_case),
        c.expectation_ps(**xyz_case),
        atol=1e-5,
    )


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_random_clifford_expectation(backend):
    """Test random Clifford circuits against tc.Circuit."""
    n = 3
    for _ in range(5):
        c = tc.Circuit(n)
        for _ in range(10):
            gate_type = np.random.choice(["h", "s", "cnot", "cz", "x", "y", "z"])
            q1 = np.random.randint(n)
            if gate_type in ["cnot", "cz"]:
                q2 = np.random.randint(n)
                while q1 == q2:
                    q2 = np.random.randint(n)
                getattr(c, gate_type)(q1, q2)
            else:
                getattr(c, gate_type)(q1)

        stc = StabilizerTCircuit.from_circuit(c)

        # Random Pauli string
        ps = np.random.randint(0, 4, size=n)
        exp_zx = stc.expectation_ps(ps=list(ps))
        exp_exact = c.expectation_ps(ps=list(ps))
        np.testing.assert_allclose(exp_zx, exp_exact, atol=1e-5)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_noisy_expectation_ps_x_error_mc(backend):
    p = 0.2
    stc = StabilizerTCircuit(1, seed=123)
    stc.x_error(0, p)
    exp_mc = stc.expectation_ps(z=[0], nmc=20000)
    np.testing.assert_allclose(exp_mc, 1.0 - 2.0 * p, atol=0.03)


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_noisy_expectation_ps_matches_weighted_reference(backend):
    p_dep = 0.15
    p_x = 0.1
    pauli_spec = [1, 3]  # X0 Z1

    stc = StabilizerTCircuit(2, seed=321)
    stc.h(0)
    stc.cnot(0, 1)
    stc.s(1)
    stc.depolarizing(0, p_dep)
    stc.x_error(1, p_x)
    stc.h(1)

    ntraj = 5000
    exp_mc = stc.expectation_ps(ps=pauli_spec, nmc=ntraj)

    # Trajectory-level tc.Circuit baseline via built-in depolarizing noise sampling.
    ref_vals = []
    for _ in range(ntraj):
        c = tc.Circuit(2)
        c.h(0)
        c.cnot(0, 1)
        c.s(1)
        c.depolarizing(0, px=p_dep / 3, py=p_dep / 3, pz=p_dep / 3)
        c.depolarizing(1, px=p_x, py=0.0, pz=0.0)
        c.h(1)
        ref_vals.append(c.expectation_ps(ps=pauli_spec))
    ref = np.mean(np.asarray(ref_vals, dtype=np.complex128))

    np.testing.assert_allclose(exp_mc, ref, atol=0.04)


@pytest.mark.parametrize("backend", [lf("npb")])
def test_from_circuit_ordering(backend):
    c = tc.Circuit(1)
    c.h(0)
    # len(c._qir) is 1
    c.measure_instruction(0)
    c.reset_instruction(0)
    # len(c._qir) is still 1, but next gate will be at index 1
    c.x(0)
    # len(c._qir) is 2

    stc = StabilizerTCircuit.from_circuit(c)

    names = [d["name"].upper() for d in stc._qir]
    # translation._merge_extra_qir logic: extra instructions at pos i come BEFORE qir[i]
    expected = ["H", "MEASURE", "RESET", "X"]
    assert names == expected


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_mpp_basic(backend):
    """Test MPP (multi-Pauli product) measurements."""
    import stim

    # Test 1: MPP X0*X1 on Bell state |00⟩+|11⟩ should give deterministic result
    stim_circuit = stim.Circuit("""
        H 0
        CNOT 0 1
        MPP X0*X1
    """)

    stc = StabilizerTCircuit.from_stim_circuit(stim_circuit)
    samples = stc.sample_measurements(shots=100)
    # X0*X1 is a stabilizer of Bell state, eigenvalue +1
    assert np.all(samples == 0), "MPP X0*X1 on Bell state should be deterministic 0"

    # Test 2: MPP Z0*Z1 on Bell state should also be deterministic
    stim_circuit = stim.Circuit("""
        H 0
        CNOT 0 1
        MPP Z0*Z1
    """)

    stc = StabilizerTCircuit.from_stim_circuit(stim_circuit)
    samples = stc.sample_measurements(shots=100)
    # Z0*Z1 is a stabilizer of Bell state, eigenvalue +1
    assert np.all(samples == 0), "MPP Z0*Z1 on Bell state should be deterministic 0"

    # Test 3: MPP Y0*Y1 on Bell state
    stim_circuit = stim.Circuit("""
        H 0
        CNOT 0 1
        MPP Y0*Y1
    """)

    stc = StabilizerTCircuit.from_stim_circuit(stim_circuit)
    samples = stc.sample_measurements(shots=100)
    # Y0*Y1 = -X0*X1*Z0*Z1 on Bell state, eigenvalue -1
    assert np.all(samples == 1), "MPP Y0*Y1 on Bell state should be deterministic 1"


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_mpp_with_noise(backend):
    """Test MPP measurements with noise model."""
    import stim

    # X error should flip Z-basis MPP result
    stim_circuit = stim.Circuit("""
        H 0
        CNOT 0 1
        X_ERROR(0.2) 0
        MPP Z0*Z1
    """)

    stc = StabilizerTCircuit.from_stim_circuit(stim_circuit)
    samples = stc.sample_measurements(shots=5000)
    flip_rate = np.mean(samples)
    # X error on qubit 0 anticommutes with Z0, so flips the parity
    np.testing.assert_allclose(
        flip_rate,
        0.2,
        atol=0.05,
        rtol=0.0,
        err_msg=f"Expected ~0.2 flip rate, got {flip_rate}",
    )


@pytest.mark.parametrize("backend", [lf("jaxb")])
def test_zx_mpp_detector(backend):
    """Test MPP in detector context (syndrome extraction pattern)."""
    import stim

    # Repeated syndrome measurement pattern
    # Z error anticommutes with X0*X1, so it flips the parity
    stim_circuit = stim.Circuit("""
        H 0
        CNOT 0 1
        MPP X0*X1
        Z_ERROR(0.1) 0
        MPP X0*X1
        DETECTOR rec[-1] rec[-2]
    """)

    stc = StabilizerTCircuit.from_stim_circuit(stim_circuit)
    det_samples = stc.sample_detectors(shots=5000)
    det_rate = np.mean(det_samples)
    # Detector should fire when Z error occurs between measurements
    np.testing.assert_allclose(
        det_rate,
        0.1,
        atol=0.05,
        rtol=0.0,
        err_msg=f"Expected ~0.1 detector rate, got {det_rate}",
    )
