# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

"""Pauli operator representation for quantum spin systems."""

from collections.abc import Sequence

try:
    import cirq
except Exception as ex:
    raise ImportError(
        "qdk.magnets.models requires the cirq extras. Install with 'pip install \"qdk[cirq]\"'."
    ) from ex


class Pauli:
    """Single-qubit Pauli term tied to an explicit qubit index.

    ``Pauli`` stores a Pauli identifier and the qubit it acts on. The Pauli
    identifier can be provided either as an integer code or a label:

    - ``0`` / ``"I"``
    - ``1`` / ``"X"``
    - ``2`` / ``"Z"``
    - ``3`` / ``"Y"``

    Note:
        The integer mapping follows the internal QDK convention where ``2`` is
        ``Z`` and ``3`` is ``Y``.

    Example:

    .. code-block:: python
        >>> p = Pauli("Y", qubit=2)
        >>> p.op
        3
        >>> p.qubit
        2
    """

    _VALID_INTS = {0, 1, 2, 3}
    _STR_TO_INT = {"I": 0, "X": 1, "Z": 2, "Y": 3}

    def __init__(self, value: int | str, qubit: int = 0) -> None:
        """Initialize a Pauli operator.

        Args:
            value: An integer 0-3 or one of 'I', 'X', 'Y', 'Z' (case-insensitive).
            qubit: The index of the qubit this operator acts on. Defaults to 0.

        Raises:
            ValueError: If ``value`` is not a valid integer/string Pauli identifier.
        """
        if isinstance(value, int):
            if value not in self._VALID_INTS:
                raise ValueError(f"Integer value must be 0-3, got {value}.")
            self._op = value
        elif isinstance(value, str):
            key = value.upper()
            if key not in self._STR_TO_INT:
                raise ValueError(
                    f"String value must be one of 'I', 'X', 'Y', 'Z', got '{value}'."
                )
            self._op = self._STR_TO_INT[key]
        else:
            raise ValueError(f"Expected int or str, got {type(value).__name__}.")
        self.qubit: int = qubit

    @property
    def op(self) -> int:
        """Integer encoding of this Pauli term.

        Returns:
            ``0`` for ``I``, ``1`` for ``X``, ``2`` for ``Z``, ``3`` for ``Y``.
        """
        return self._op

    def __str__(self) -> str:
        labels = {0: "I", 1: "X", 2: "Z", 3: "Y"}
        return f"{labels[self._op]}({self.qubit})"

    def __repr__(self) -> str:
        labels = {0: "I", 1: "X", 2: "Z", 3: "Y"}
        return f"Pauli('{labels[self._op]}', qubit={self.qubit})"

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Pauli):
            return NotImplemented
        return self._op == other._op and self.qubit == other.qubit

    def __hash__(self) -> int:
        return hash((self._op, self.qubit))

    @property
    def cirq(self):
        """Return this Pauli term as a Cirq gate operation on ``LineQubit``.

        Returns:
            A Cirq operation equivalent to
            ``cirq.{I|X|Z|Y}.on(cirq.LineQubit(self.qubit))``.
        """
        _INT_TO_CIRQ = (cirq.I, cirq.X, cirq.Z, cirq.Y)
        return _INT_TO_CIRQ[self._op].on(cirq.LineQubit(self.qubit))


def PauliX(qubit: int) -> Pauli:
    """Create a Pauli-X operator on the given qubit."""
    return Pauli("X", qubit)


def PauliY(qubit: int) -> Pauli:
    """Create a Pauli-Y operator on the given qubit."""
    return Pauli("Y", qubit)


def PauliZ(qubit: int) -> Pauli:
    """Create a Pauli-Z operator on the given qubit."""
    return Pauli("Z", qubit)


class PauliString:
    """Ordered tensor product of single-qubit ``Pauli`` terms with a coefficient.

    ``PauliString`` stores:

    - an ordered tuple of :class:`Pauli` objects (including each term's qubit), and
    - a complex scalar coefficient.

    Construction options:

    - pass a sequence of :class:`Pauli` objects to ``PauliString(...)``
    - use :meth:`from_qubits` to pair qubit indices with Pauli labels/codes

    Example:

    .. code-block:: python
        >>> ps = PauliString([PauliX(0), PauliZ(1)], coefficient=-1j)
        >>> ps.qubits
        (0, 1)
        >>> ps2 = PauliString.from_qubits((0, 1), "XZ", coefficient=-1j)
        >>> ps == ps2
        True
    """

    def __init__(self, paulis: Sequence[Pauli], coefficient: complex = 1.0) -> None:
        """Initialize a PauliString from a sequence of Pauli operators.

        Args:
            paulis: A sequence of :class:`Pauli` instances, each with its
                own qubit index.
            coefficient: Complex coefficient multiplying the Pauli string (default 1.0).

        Raises:
            TypeError: If any element is not a Pauli instance.
        """
        for p in paulis:
            if not isinstance(p, Pauli):
                raise TypeError(
                    f"Expected Pauli instance, got {type(p).__name__}. "
                    "Use PauliString.from_qubits() for int/str values."
                )
        self._paulis: tuple[Pauli, ...] = tuple(paulis)
        self._coefficient: complex = coefficient

    @classmethod
    def from_qubits(
        cls,
        qubits: tuple[int, ...],
        values: Sequence[int | str] | str,
        coefficient: complex = 1.0,
    ) -> "PauliString":
        """Create a PauliString from qubit indices and Pauli labels.

        Args:
            qubits: Tuple of qubit indices.
            values: Sequence of Pauli identifiers (integers 0-3 or strings
                'I', 'X', 'Y', 'Z'). A plain string like ``"XZI"`` is also
                accepted and treated as individual characters.
            coefficient: Complex coefficient multiplying the Pauli string.

        Returns:
            A new PauliString instance.

        Raises:
            ValueError: If qubits and values have different lengths, or if
                any value is not a valid Pauli identifier.
        """
        if len(qubits) != len(values):
            raise ValueError(
                f"Length mismatch: {len(qubits)} qubits vs {len(values)} values."
            )
        paulis = [Pauli(v, q) for q, v in zip(qubits, values)]
        return cls(paulis, coefficient=coefficient)

    @property
    def qubits(self) -> tuple[int, ...]:
        """Tuple of qubit indices in the same order as the stored Pauli terms.

        Returns:
            Tuple of qubit indices, one per Pauli operator.
        """
        return tuple(p.qubit for p in self._paulis)

    @property
    def coefficient(self) -> complex:
        """Complex coefficient multiplying this Pauli string."""
        return self._coefficient

    @property
    def paulis(self) -> str:
        """String of Pauli labels in the same order as the stored Pauli terms.

        Returns:
            String of Pauli labels ('I', 'X', 'Z', 'Y'), one per Pauli operator.
        """
        labels = {0: "I", 1: "X", 2: "Z", 3: "Y"}
        return "".join(labels[p.op] for p in self._paulis)

    def __iter__(self):
        """Iterate over Pauli terms in stored order.

        Yields:
            :class:`Pauli` instances in order.
        """
        return iter(self._paulis)

    def __len__(self) -> int:
        return len(self._paulis)

    def __getitem__(self, index: int) -> Pauli:
        return self._paulis[index]

    def __mul__(self, scalar: complex) -> "PauliString":
        """Scale the coefficient of this PauliString by a complex scalar."""
        return PauliString(self._paulis, coefficient=self._coefficient * scalar)

    def __str__(self) -> str:
        labels = {0: "I", 1: "X", 2: "Z", 3: "Y"}
        s = "".join(map(str, self._paulis))
        return f"{self._coefficient} * {s}"

    def __repr__(self) -> str:
        labels = {0: "I", 1: "X", 2: "Z", 3: "Y"}
        s = "".join(labels[p.op] for p in self._paulis)
        return f"PauliString(qubits={self.qubits}, ops='{s}', coefficient={self._coefficient})"

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, PauliString):
            return NotImplemented
        return self._paulis == other._paulis and self._coefficient == other._coefficient

    def __hash__(self) -> int:
        return hash((self._paulis, self._coefficient))

    @property
    def cirq(self):
        """Return the corresponding Cirq ``PauliString``.

        Constructs a ``cirq.PauliString`` by applying each single-qubit
        Pauli to its corresponding ``cirq.LineQubit``.

        Returns:
            A ``cirq.PauliString`` on ``cirq.LineQubit`` instances with
            ``self._coefficient`` as its coefficient.
        """
        _INT_TO_CIRQ = (cirq.I, cirq.X, cirq.Z, cirq.Y)
        return cirq.PauliString(
            {cirq.LineQubit(p.qubit): _INT_TO_CIRQ[p.op] for p in self._paulis},
            coefficient=self._coefficient,
        )
