import pytest
import stim

from tqecd.exceptions import TQECDException
from tqecd.pauli import PauliString, pauli_literal_to_bools, pauli_product


def test_pauli_string_construction() -> None:
    ps1 = PauliString({0: "X", 1: "Y", 2: "Z"})
    ps2 = PauliString({0: "X", 2: "Z", 1: "Y"})
    empty = PauliString({})
    assert bool(ps1)
    assert not empty
    assert len(ps1) == 3
    assert ps1 == ps2
    assert len({ps1, ps2}) == 1
    assert ps1.overlaps(ps2)
    assert not ps1.overlaps(PauliString({3: "X"}))
    with pytest.raises(TQECDException, match=r"^Invalid Pauli operator.*"):
        PauliString({0: "W"})  # type: ignore


def test_pauli_string_interop_with_stim() -> None:
    stim_pauli_string = stim.PauliString.random(num_qubits=23)
    pauli_string = PauliString.from_stim_pauli_string(stim_pauli_string)
    assert (
        pauli_string.to_stim_pauli_string(length=23) * stim_pauli_string
    ).weight == 0

    pauli_string = PauliString.from_stim_pauli_string(stim.PauliString("_XYZ"))
    assert pauli_string.after(
        stim.Tableau.from_named_gate("CZ"), targets=[0, 1]
    ) == PauliString.from_stim_pauli_string(stim.PauliString("+ZXYZ"))

    assert PauliString.from_stim_pauli_string(stim.PauliString("IXYZ")) == PauliString(
        {0: "I", 1: "X", 2: "Y", 3: "Z"}
    )

    with pytest.raises(
        TQECDException, match=r"^The length specified 2 <= the maximum qubit index.*"
    ):
        pauli_string.to_stim_pauli_string(2)


def test_pauli_string_mul() -> None:
    a = PauliString({q: p for q, p in enumerate("IIIIXXXXYYYYZZZZ")})  # type:ignore
    b = PauliString({q: p for q, p in enumerate("IXYZ" * 4)})  # type:ignore
    c = PauliString({q: p for q, p in enumerate("IXYZXIZYYZIXZYXI")})  # type:ignore
    assert a * b == c


def test_pauli_string_commutation() -> None:
    a = PauliString({0: "X", 1: "Y"})
    b = PauliString({0: "Y", 1: "Z"})
    c = PauliString({0: "Z", 1: "Y"})
    assert a.commutes(b)
    assert b.commutes(a)
    assert a.anticommutes(c)
    assert c.anticommutes(a)
    assert b.commutes(c)
    assert c.commutes(b)


def test_pauli_string_collapse_by() -> None:
    X0Z1 = PauliString({0: "X", 1: "Z"})
    Z0 = PauliString({0: "Z"})
    Z1 = PauliString({1: "Z"})
    X0 = PauliString({0: "X"})
    assert X0Z1.collapse_by([X0]) == Z1
    assert X0Z1.collapse_by([X0, Z1]) == PauliString({})
    with pytest.raises(TQECDException):
        X0Z1.collapse_by([Z0])


def test_pauli_string_weight() -> None:
    X0Z1 = PauliString({0: "X", 1: "Z", 2: "I"})
    Z0 = PauliString({0: "Z"})
    I0to20 = PauliString({i: "I" for i in range(20)})
    assert X0Z1.non_trivial_pauli_count == 2
    assert Z0.non_trivial_pauli_count == 1
    assert I0to20.non_trivial_pauli_count == 0


def test_pauli_string_qubit() -> None:
    X0Z1 = PauliString({0: "X", 1: "Z", 2: "I"})
    Z0 = PauliString({0: "Z"})
    I0to20 = PauliString({i: "I" for i in range(20)})
    assert Z0.qubit == 0
    with pytest.raises(TQECDException):
        X0Z1.qubit
    with pytest.raises(TQECDException):
        I0to20.qubit


def test_pauli_string_indexing() -> None:
    X0Z1 = PauliString({0: "X", 1: "Z", 2: "I"})
    assert X0Z1[0] == "X"
    assert X0Z1[1] == "Z"
    assert X0Z1[2] == "I"
    assert X0Z1[3] == "I"


def test_pauli_literals_to_bool() -> None:
    assert pauli_literal_to_bools("I") == (False, False)
    assert pauli_literal_to_bools("X") == (True, False)
    assert pauli_literal_to_bools("Y") == (True, True)
    assert pauli_literal_to_bools("Z") == (False, True)


def test_pauli_product() -> None:
    X0Z1 = PauliString({0: "X", 1: "Z", 2: "I"})
    Z0 = PauliString({0: "Z"})
    I0to20 = PauliString({i: "I" for i in range(20)})

    assert pauli_product([I0to20 for _ in range(20)]) == I0to20
    assert pauli_product([X0Z1 for _ in range(20)]) == I0to20
    assert pauli_product([X0Z1 for _ in range(21)]) == X0Z1
    assert pauli_product([X0Z1, Z0]) == PauliString({0: "Y", 1: "Z"})
