##############################################################################
# Copyright 2016-2018 Rigetti Computing
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
##############################################################################
"""Module for working with Pauli algebras."""

import copy
import re
import warnings
from collections import OrderedDict
from collections.abc import Hashable, Iterable, Iterator, Sequence
from functools import reduce
from itertools import product
from numbers import Complex, Number
from typing import (
    Callable,
    Optional,
    Union,
    cast,
)

import numpy as np
import quil.instructions as quil_rs
from numpy.typing import NDArray
from scipy.linalg import expm

from pyquil.quilatom import (
    Expression,
    ExpressionDesignator,
    FormalArgument,
    QubitDesignator,
    QubitPlaceholder,
    _convert_to_py_expression,
    _convert_to_py_qubits,
    _convert_to_rs_expression,
)

from .gates import CNOT, PHASE, QUANTUM_GATES, RX, RZ, H, X
from .quil import Program

PauliTargetDesignator = QubitDesignator
PauliDesignator = Union["PauliTerm", "PauliSum"]

PAULI_OPS = ["X", "Y", "Z", "I"]
PAULI_PROD = {
    "ZZ": "I",
    "YY": "I",
    "XX": "I",
    "II": "I",
    "XY": "Z",
    "XZ": "Y",
    "YX": "Z",
    "YZ": "X",
    "ZX": "Y",
    "ZY": "X",
    "IX": "X",
    "IY": "Y",
    "IZ": "Z",
    "ZI": "Z",
    "YI": "Y",
    "XI": "X",
    "X": "X",
    "Y": "Y",
    "Z": "Z",
    "I": "I",
}
PAULI_COEFF = {
    "ZZ": 1.0,
    "YY": 1.0,
    "XX": 1.0,
    "II": 1.0,
    "XY": 1.0j,
    "XZ": -1.0j,
    "YX": -1.0j,
    "YZ": 1.0j,
    "ZX": 1.0j,
    "ZY": -1.0j,
    "IX": 1.0,
    "IY": 1.0,
    "IZ": 1.0,
    "ZI": 1.0,
    "YI": 1.0,
    "XI": 1.0,
    "X": 1.0,
    "Y": 1.0,
    "Z": 1.0,
    "I": 1.0,
}


class UnequalLengthWarning(Warning):
    """Warning issued when multiplying PauliTerms of different lengths."""

    def __init__(self, *args: object, **kwargs: object):
        """Initialize the warning."""
        super().__init__(*args, **kwargs)


integer_types = (int, np.int64, np.int32, np.int16, np.int8)
"""Explicitly include numpy integer dtypes (for python 3)."""

HASH_PRECISION = 1e6
"""The precision used when hashing terms to check equality. The simplify() method
uses np.isclose() for coefficient comparisons to 0 which has its own default precision. We
can't use np.isclose() for hashing terms though.
"""


def _valid_qubit(index: Optional[Union[PauliTargetDesignator, QubitPlaceholder]]) -> bool:
    return (
        (isinstance(index, integer_types) and index >= 0)
        or isinstance(index, QubitPlaceholder)
        or isinstance(index, FormalArgument)
    )


class PauliTerm:
    """A term is a product of Pauli operators operating on different qubits."""

    def __init__(
        self,
        op: str,
        index: Optional[PauliTargetDesignator],
        coefficient: ExpressionDesignator = 1.0,
    ):
        """Create a new Pauli Term with a Pauli operator at a particular index and a leading coefficient.

        :param op: The Pauli operator as a string "X", "Y", "Z", or "I"
        :param index: The qubit index that that operator is applied to.
        :param coefficient: The coefficient multiplying the operator, e.g. 1.5 * Z_1
        """
        if op not in PAULI_OPS:
            raise ValueError(f"{op} is not a valid Pauli operator")

        self._ops: dict[PauliTargetDesignator, str] = OrderedDict()
        if op != "I":
            if not _valid_qubit(index):
                raise ValueError(f"{index} is not a valid qubit")
            if index is None:
                raise ValueError("PauliTerm index cannot be None")
            self._ops[index] = op

        if isinstance(coefficient, Number):
            self.coefficient: Union[complex, Expression] = complex(coefficient)
        else:
            self.coefficient = coefficient

    @classmethod
    def _from_rs_pauli_term(cls, term: quil_rs.PauliTerm) -> "PauliTerm":
        term_list = [(str(gate), FormalArgument(arg)) for (gate, arg) in term.arguments]
        coefficient = _convert_to_py_expression(term.expression.into_simplified())

        return cls.from_list(term_list, coefficient)

    def _to_rs_pauli_term(self) -> quil_rs.PauliTerm:
        arguments = [(quil_rs.PauliGate.parse(gate), str(arg)) for arg, gate in self._ops.items()]
        return quil_rs.PauliTerm(arguments, _convert_to_rs_expression(self.coefficient))

    def id(self, sort_ops: bool = True) -> str:
        """Return an identifier string for the PauliTerm (ignoring the coefficient).

        Don't use this to compare terms. This function will not work with qubits that
        aren't sortable.

        :param sort_ops: Whether to sort operations by qubit. This is True by default for
            backwards compatibility but will change in a future version. Callers should never rely
            on comparing id's for testing equality. See ``operations_as_set`` instead.
        :return: A string representation of this term's operations.
        """
        if len(self._ops) == 0 and not sort_ops:
            # This is nefariously backwards-compatibility breaking. There's potentially
            # lots of code floating around that says is_identity = term.id() == ''
            # Please use `is_identity(term)`!
            # Therefore, we only return 'I' when sort_ops is set to False, which is the newer
            # way of calling this function and implies the user knows what they're doing.
            return "I"

        if sort_ops and len(self._ops) > 1:
            warnings.warn(
                "`PauliTerm.id()` will not work on PauliTerms where the qubits are not "
                "sortable and should be avoided in favor of `operations_as_set`.",
                FutureWarning,
                stacklevel=2,
            )
            return "".join(f"{self._ops[q]}{q}" for q in sorted(self._ops.keys()))  # type: ignore
        else:
            return "".join(f"{p}{q}" for q, p in self._ops.items())

    def operations_as_set(self) -> frozenset[tuple[PauliTargetDesignator, str]]:
        """Return a frozenset of operations in this term.

        Use this in place of :py:func:`id` if the order of operations in the term does not
        matter.

        :return: frozenset of (qubit, op_str) representing Pauli operations
        """
        return frozenset(self._ops.items())

    def __eq__(self, other: object) -> bool:
        """Return True if two PauliTerms are equal."""
        if not isinstance(other, (PauliTerm, PauliSum)):
            raise TypeError(f"Can't compare PauliTerm with object of type {type(other)}.")
        elif isinstance(other, PauliSum):
            return other == self
        else:
            return self.operations_as_set() == other.operations_as_set() and np.allclose(
                self.coefficient, other.coefficient
            )

    def __hash__(self) -> int:
        """Return a hash of the PauliTerm."""
        if not isinstance(self.coefficient, Complex):
            raise ValueError("PauliTerm cannot be hashed if the coefficient is not a complex number.")
        return hash(
            (
                round(self.coefficient.real * HASH_PRECISION),
                round(self.coefficient.imag * HASH_PRECISION),
                self.operations_as_set(),
            )
        )

    def __len__(self) -> int:
        """Length of the PauliTerm is the number of Pauli operators in the term.

        A term that consists of only a scalar has a length of zero.
        """
        return len(self._ops)

    def copy(self) -> "PauliTerm":
        """Create a new PauliTerm, with a completely new dictionary of operators."""
        new_term = PauliTerm("I", 0, 1.0)  # create new object
        # manually copy all attributes over
        for key in self.__dict__.keys():
            val = self.__dict__[key]
            if isinstance(val, (dict, list, set)):  # mutable types
                new_term.__dict__[key] = copy.copy(val)
            else:  # immutable types
                new_term.__dict__[key] = val

        return new_term

    @property
    def program(self) -> Program:
        """Create Program from the PauliTerm."""
        return Program([QUANTUM_GATES[gate](q) for q, gate in self])

    def get_qubits(self) -> list[PauliTargetDesignator]:
        """Get all the qubits that this PauliTerm operates on."""
        return list(self._ops.keys())

    def __getitem__(self, i: PauliTargetDesignator) -> str:
        """Get the Pauli operator for a particular target."""
        return self._ops.get(i, "I")

    def __iter__(self) -> Iterator[tuple[PauliTargetDesignator, str]]:
        """Iterate over the quibit indices and Pauli operators in the term."""
        for i in self.get_qubits():
            yield i, self[i]

    def _multiply_factor(self, factor: str, index: PauliTargetDesignator) -> "PauliTerm":
        new_term = PauliTerm("I", 0)
        new_coeff = self.coefficient
        new_ops = self._ops.copy()

        ops = self[index] + factor
        new_op = PAULI_PROD[ops]
        if new_op != "I":
            new_ops[index] = new_op
        else:
            del new_ops[index]
        new_coeff *= PAULI_COEFF[ops]

        new_term._ops = new_ops
        new_term.coefficient = new_coeff

        return new_term

    def __mul__(self, term: Union[PauliDesignator, ExpressionDesignator]) -> PauliDesignator:
        """Multiply this Pauli Term with another PauliTerm, PauliSum, or number according to the Pauli algebra rules.

        :param term: (PauliTerm or PauliSum or Number) A term to multiply by.
        :returns: The product of this PauliTerm and term.
        """
        if isinstance(term, PauliSum):
            return (PauliSum([self]) * term).simplify()
        elif isinstance(term, PauliTerm):
            new_term = PauliTerm("I", 0, 1.0)
            new_term._ops = self._ops.copy()
            new_coeff = self.coefficient * term.coefficient
            for index, op in term:
                new_term = new_term._multiply_factor(op, index)

            return term_with_coeff(new_term, new_term.coefficient * new_coeff)
        return term_with_coeff(self, self.coefficient * term)

    def __rmul__(self, other: ExpressionDesignator) -> "PauliTerm":
        """Multiplies this PauliTerm with another object, probably a number.

        :param other: A number or PauliTerm to multiply by
        :returns: A new PauliTerm
        """
        p = self * other
        if not isinstance(p, PauliTerm):
            raise ValueError(f"PauliTerm multiplication should return a PauliTerm, got {type(other)}")
        return p

    def __pow__(self, power: int) -> "PauliTerm":
        """Raise the PauliTerm to power.

        :param power: The power to raise this PauliTerm to.
        :return: The power-fold product of power.
        """
        if not isinstance(power, int) or power < 0:
            raise ValueError("The power must be a non-negative integer.")

        if len(self.get_qubits()) == 0:
            # There weren't any nontrivial operators
            return term_with_coeff(self, 1)

        result = ID()
        for _ in range(power):
            result = cast(PauliTerm, result * self)
        return result

    def __add__(self, other: Union[PauliDesignator, ExpressionDesignator]) -> "PauliSum":
        """Add this PauliTerm with another one.

        :param other: A PauliTerm object, a PauliSum object, or a Number
        :returns: A PauliSum object representing the sum of this PauliTerm and other
        """
        if isinstance(other, PauliSum):
            return other + self
        elif isinstance(other, PauliTerm):
            new_sum = PauliSum([self, other])
            return new_sum.simplify()
        else:  # is a Number
            return self + PauliTerm("I", 0, other)

    def __radd__(self, other: ExpressionDesignator) -> "PauliSum":
        """Add this PauliTerm with a Number.

        :param other: A Number
        :returns: A new PauliSum
        """
        return PauliTerm("I", 0, other) + self

    def __sub__(self, other: Union["PauliTerm", ExpressionDesignator]) -> "PauliSum":
        """Subtracts a PauliTerm from this one.

        :param other: A PauliTerm object, a number, or an Expression
        :returns: A PauliSum object representing the difference of this PauliTerm and term
        """
        return self + -1.0 * other

    def __rsub__(self, other: Union["PauliTerm", ExpressionDesignator]) -> "PauliSum":
        """Subtracts this PauliTerm from a Number or PauliTerm.

        :param other: A PauliTerm object or a Number
        :returns: A PauliSum object representing the difference of this PauliTerm and term
        """
        return other + -1.0 * self

    def __repr__(self) -> str:
        """Get a string representation of the PauliTerm."""
        term_strs = []
        for index in self._ops.keys():
            term_strs.append(f"{self[index]}{index}")

        if len(term_strs) == 0:
            term_strs.append("I")
        out = f"{self.coefficient}*{'*'.join(term_strs)}"
        return out

    def compact_str(self) -> str:
        """Return string representation of the Pauli term that is more compact than ``str(term)``.

        >>> term = 2.0 * sX(1) * sZ(2)
        >>> str(term)
        '(2+0j)*X1*Z2'
        >>> term.compact_str()
        '(2+0j)*X1Z2'
        """
        return f"{self.coefficient}*{self.id(sort_ops=False)}"

    @classmethod
    def from_list(
        cls, terms_list: Sequence[tuple[str, PauliTargetDesignator]], coefficient: ExpressionDesignator = 1.0
    ) -> "PauliTerm":
        """Allocate a PauliTerm from a list of operators and indices.

        This is more efficient than multiplying together individual terms.

        :param list terms_list: A list of tuples, e.g. [("X", 0), ("Y", 1)]
        :return: PauliTerm
        """
        if not all([isinstance(op, tuple) for op in terms_list]):
            raise TypeError(
                "The type of terms_list should be a list of (name, index) " "tuples suitable for PauliTerm()."
            )

        pterm = PauliTerm("I", 0)
        if not all([op[0] in PAULI_OPS for op in terms_list]):
            raise ValueError("All operators must be in ['X', 'Y', 'Z', 'I']")

        indices = [op[1] for op in terms_list]
        if not all(_valid_qubit(index) for index in indices):
            raise ValueError("All qubit indices must be valid.")

        # this is because from_list doesn't call simplify in order to be more efficient.
        if len(set(indices)) != len(indices):
            raise ValueError(
                "Elements of PauliTerm that are allocated using from_list must "
                "be on disjoint qubits. Use PauliTerm multiplication to simplify "
                "terms instead."
            )

        for op, index in terms_list:
            if op != "I":
                pterm._ops[index] = op
        if isinstance(coefficient, Number):
            pterm.coefficient = complex(coefficient)
        else:
            pterm.coefficient = coefficient
        return pterm

    @classmethod
    def from_compact_str(cls, str_pauli_term: str) -> "PauliTerm":
        """Construct a PauliTerm from the result of str(pauli_term)."""
        # split into str_coef, str_op at first '*'' outside parenthesis
        try:
            str_coef, str_op = re.split(r"\*(?![^(]*\))", str_pauli_term, maxsplit=1)
        except ValueError as e:
            raise ValueError(
                "Could not separate the pauli string into "
                f"coefficient and operator. {str_pauli_term} does"
                " not match <coefficient>*<operator>"
            ) from e

        # parse the coefficient into either a float or complex
        str_coef = str_coef.replace(" ", "")
        try:
            coef: Union[float, complex] = float(str_coef)
        except ValueError:
            try:
                coef = complex(str_coef)
            except ValueError as e:
                raise ValueError(f"Could not parse the coefficient {str_coef}") from e

        op = sI() * coef
        if str_op == "I":
            if not isinstance(op, PauliTerm):
                raise ValueError(f"Expected operation to be PauliTerm, got {type(op)}.")
            return op

        # parse the operator
        str_op = re.sub(r"\*", "", str_op)
        if not re.match(r"^(([XYZ])(\d+))+$", str_op):
            raise ValueError(rf"Could not parse operator string {str_op}. It should match ^(([XYZ])(\d+))+$")

        for factor in re.finditer(r"([XYZ])(\d+)", str_op):
            op *= cls(factor.group(1), int(factor.group(2)))

        if not isinstance(op, PauliTerm):
            raise ValueError(f"Expected operation to be PauliTerm, got {type(op)}.")
        return op

    def pauli_string(self, qubits: Optional[Iterable[int]] = None) -> str:
        """Return a string representation of this PauliTerm without its coefficient and with implicit qubit indices.

        If an iterable of qubits is provided, each character in the resulting string represents
        a Pauli operator on the corresponding qubit.

        >>> p = PauliTerm("X", 0) * PauliTerm("Y", 1, 1.0j)
        >>> p.pauli_string()
        'XY'
        >>> p.pauli_string(qubits=[0])
        'X'
        >>> p.pauli_string(qubits=[0, 2])
        'XI'

        :param iterable of qubits: The iterable of qubits to represent, given as ints.
        :return: The string representation of this PauliTerm, sans coefficient
        """
        if qubits is None:
            return "".join(self._ops.values())
        return "".join(self[q] for q in qubits)


# For convenience, a shorthand for several operators.
def ID() -> PauliTerm:
    """Identity operator."""
    return PauliTerm("I", 0, 1)


def ZERO() -> PauliTerm:
    """Zero operator."""
    return PauliTerm("I", 0, 0)


def sI(q: Optional[int] = None) -> PauliTerm:
    """Return the identity operator, optionally on a particular qubit.

    This can be specified without a qubit.

    :param qubit_index: The optional index of a qubit.
    :returns: A PauliTerm object
    """
    return PauliTerm("I", q)


def sX(q: int) -> PauliTerm:
    """Return the sigma_X operator on a particular qubit.

    :param qubit_index: The index of the qubit
    :returns: A PauliTerm object
    """
    return PauliTerm("X", q)


def sY(q: int) -> PauliTerm:
    """Return the sigma_Y operator on a particular qubit.

    :param qubit_index: The index of the qubit
    :returns: A PauliTerm object
    """
    return PauliTerm("Y", q)


def sZ(q: int) -> PauliTerm:
    """Return the sigma_Z operator on a particular qubit.

    :param qubit_index: The index of the qubit
    :returns: A PauliTerm object
    """
    return PauliTerm("Z", q)


def term_with_coeff(term: PauliTerm, coeff: ExpressionDesignator) -> PauliTerm:
    """Change the coefficient of a PauliTerm.

    :param term: A PauliTerm object
    :param coeff: The coefficient to set on the PauliTerm
    :returns: A new PauliTerm that duplicates term but sets coeff
    """
    if not isinstance(coeff, Number):
        raise ValueError("coeff must be a Number")
    new_pauli = term.copy()
    # We cast to a complex number to ensure that internally the coefficients remain compatible.
    new_pauli.coefficient = complex(coeff)
    return new_pauli


class PauliSum:
    """A sum of one or more PauliTerms."""

    def __init__(self, terms: Sequence[PauliTerm]):
        """Initialize a PauliSum with a list of PauliTerms.

        :param Sequence terms: A Sequence of PauliTerms.
        """
        if not (isinstance(terms, Sequence) and all([isinstance(term, PauliTerm) for term in terms])):
            raise ValueError("PauliSum's are currently constructed from Sequences of PauliTerms.")
        self.terms: Sequence[PauliTerm]
        if len(terms) == 0:
            self.terms = [0.0 * ID()]
        else:
            self.terms = terms

    @classmethod
    def _from_rs_pauli_sum(cls, pauli_sum: quil_rs.PauliSum) -> "PauliSum":
        return cls([PauliTerm._from_rs_pauli_term(term) for term in pauli_sum.terms])

    def _to_rs_pauli_sum(self, arguments: Optional[list[PauliTargetDesignator]] = None) -> quil_rs.PauliSum:
        rs_arguments: list[str]
        if arguments is None:
            argument_set: dict[str, None] = {}
            for term_arguments in [term.get_qubits() for term in self.terms]:
                argument_set.update({str(arg): None for arg in term_arguments})
            rs_arguments = list(argument_set.keys())
        else:
            rs_arguments = [str(arg) for arg in arguments]
        terms = [term._to_rs_pauli_term() for term in self.terms]
        return quil_rs.PauliSum(rs_arguments, terms)

    def __eq__(self, other: object) -> bool:
        """Equality testing to see if two PauliSum's are equivalent.

        :param other: The PauliSum to compare this PauliSum with.
        :return: True if other is equivalent to this PauliSum, False otherwise.
        """
        if not isinstance(other, (PauliTerm, PauliSum)):
            raise TypeError(f"Can't compare PauliSum with object of type {type(other)}.")
        elif isinstance(other, PauliTerm):
            return self == PauliSum([other])
        elif len(self.terms) != len(other.terms):
            return False

        return set(self.terms) == set(other.terms)

    def __hash__(self) -> int:
        """Return a hash value for the sum."""
        return hash(frozenset(self.terms))

    def __repr__(self) -> str:
        """Return a string representation for the sum."""
        return " + ".join([str(term) for term in self.terms])

    def __len__(self) -> int:
        """Return the number of PauliTerms in the sum."""
        return len(self.terms)

    def __getitem__(self, item: int) -> PauliTerm:
        """Get the PauliTerm at the given position the sum.

        :param item: The index of the term in the sum to return
        :return: The PauliTerm at the index-th position in the PauliSum
        """
        return self.terms[item]

    def __iter__(self) -> Iterator[PauliTerm]:
        """Iterate over the PauliTerms in the sum."""
        return self.terms.__iter__()

    def __mul__(self, other: Union[PauliDesignator, ExpressionDesignator]) -> "PauliSum":
        """Multiply and simplify this PauliSum with another PauliSum, PauliTerm or Number object.

        :param other: a PauliSum, PauliTerm or Number object
        :return: A new PauliSum object given by the multiplication.
        """
        if not isinstance(other, (Expression, Number, PauliTerm, PauliSum)):
            raise ValueError("Cannot multiply PauliSum by term that is not a Number, PauliTerm, or PauliSum")

        other_terms: list[Union[PauliTerm, ExpressionDesignator]] = []
        if isinstance(other, PauliSum):
            other_terms += other.terms
        else:
            other_terms += [other]

        new_terms = [lterm * rterm for lterm, rterm in product(self.terms, other_terms)]
        new_sum = PauliSum(cast(list[PauliTerm], new_terms))
        return new_sum.simplify()

    def __rmul__(self, other: ExpressionDesignator) -> "PauliSum":
        """Multiply and simplify this PauliSum with another PauliSum, PauliTerm or Number object.

        :param other: a PauliSum, PauliTerm or Number object
        :return: A new PauliSum object given by the multiplication.
        """
        if not isinstance(other, Number):
            raise TypeError(f"Expected a Number object, got {type(other)}")
        new_terms = [term.copy() for term in self.terms]
        for term in new_terms:
            term.coefficient *= other
        return PauliSum(new_terms).simplify()

    def __pow__(self, power: int) -> "PauliSum":
        """Raise this PauliSum to power.

        :param power: The power to raise this PauliSum to.
        :return: The power-th power of this PauliSum.
        """
        if not isinstance(power, int) or power < 0:
            raise ValueError("The power must be a non-negative integer.")
        result = PauliSum([ID()])

        if not self.get_qubits():
            # There aren't any nontrivial operators
            terms = [term_with_coeff(term, 1) for term in self.terms]
            for term in terms:
                result *= term
        else:
            for term in self.terms:
                for qubit_id in term.get_qubits():
                    result *= PauliTerm("I", qubit_id)

        for _ in range(power):
            result *= self
        return result

    def __add__(self, other: Union[PauliDesignator, ExpressionDesignator]) -> "PauliSum":
        """Add and simplify this PauliSum with another PauliSum, PauliTerm or Number objects.

        :param other: a PauliSum, PauliTerm or Number object
        :return: A new PauliSum object given by the addition.
        """
        if isinstance(other, PauliTerm):
            other_sum = PauliSum([other])
        elif isinstance(other, (Expression, Number, complex)):
            other_sum = PauliSum([other * ID()])
        elif isinstance(other, PauliSum):
            other_sum = other
        else:
            raise TypeError(f"{type(other)} is not a valid PauliDesignator or ExpressionDesignator")
        new_terms = [term.copy() for term in self.terms]
        new_terms.extend(other_sum.terms)
        new_sum = PauliSum(new_terms)
        return new_sum.simplify()

    def __radd__(self, other: ExpressionDesignator) -> "PauliSum":
        """Add and simplify this PauliSum with another PauliSum, PauliTerm or Number object.

        :param other: A Number
        :return: A new PauliSum object given by the addition.
        """
        if not isinstance(other, Number):
            raise TypeError(f"Expected a Number object, got {type(other)}")
        return self + other

    def __sub__(self, other: Union[PauliDesignator, ExpressionDesignator]) -> "PauliSum":
        """Subtract and simplify this PauliSum with another PauliSum, PauliTerm or Number object.

        :param other: a PauliSum, PauliTerm or Number object
        :return: A new PauliSum object given by the subtraction.
        """
        return self + -1.0 * other

    def __rsub__(self, other: Union[PauliDesignator, ExpressionDesignator]) -> "PauliSum":
        """Subtract and simplify this PauliSum with another PauliSum, PauliTerm or Number object.

        :param other: a PauliSum, PauliTerm or Number object
        :return: A new PauliSum object given by the subtraction.
        """
        return other + -1.0 * self

    def get_qubits(self) -> list[PauliTargetDesignator]:
        """Get the qubits that this PauliSum operates on.

        :returns: A list of all the qubits in the sum of terms.
        """
        all_qubits: set[PauliTargetDesignator] = set()
        for term in self.terms:
            all_qubits.update(term.get_qubits())
        return _convert_to_py_qubits(set(all_qubits))

    def simplify(self) -> "PauliSum":
        """Simplify the sum of Pauli operators according to Pauli algebra rules."""
        return simplify_pauli_sum(self)

    def get_programs(self) -> tuple[list[Program], np.ndarray]:
        """Get a Pyquil Program corresponding to each term in the PauliSum and a coefficient for each program.

        :return: (programs, coefficients)
        """
        programs = [term.program for term in self.terms]
        coefficients = np.array([term.coefficient for term in self.terms])
        return programs, coefficients

    def compact_str(self) -> str:
        """Return a string representation of the PauliSum that is more compact than ``str(pauli_sum)``.

        >>> pauli_sum = 2.0 * sX(1) * sZ(2) + 1.5 * sY(2)
        >>> str(pauli_sum)
        '(2+0j)*X1*Z2 + (1.5+0j)*Y2'
        >>> pauli_sum.compact_str()
        '(2+0j)*X1Z2+(1.5+0j)*Y2'
        """
        return "+".join([term.compact_str() for term in self.terms])

    @classmethod
    def from_compact_str(cls, str_pauli_sum: str) -> "PauliSum":
        """Construct a PauliSum from the result of str(pauli_sum)."""
        # split str_pauli_sum only at "+" outside of parenthesis to allow
        # e.g. "0.5*X0 + (0.5+0j)*Z2"
        str_terms = re.split(r"\+(?![^(]*\))", str_pauli_sum)
        str_terms = [s.strip() for s in str_terms]
        terms = [PauliTerm.from_compact_str(term) for term in str_terms]
        return cls(terms).simplify()


def simplify_pauli_sum(pauli_sum: PauliSum) -> PauliSum:
    """Simplify the sum of Pauli operators according to Pauli algebra rules.

    Warning: The simplified expression may re-order pauli operations, and may
    impact the observed performance when running on the QPU.
    """
    # You might want to use a defaultdict(list) here, but don't because
    # we want to do our best to preserve the order of terms.
    like_terms: dict[Hashable, list[PauliTerm]] = OrderedDict()
    for term in pauli_sum.terms:
        key = term.operations_as_set()
        if key in like_terms:
            like_terms[key].append(term)
        else:
            like_terms[key] = [term]

    terms = []
    for term_list in like_terms.values():
        first_term = term_list[0]
        if len(term_list) == 1 and not np.isclose(first_term.coefficient, 0.0):
            terms.append(first_term)
        else:
            coeff = sum(t.coefficient for t in term_list)
            if not np.isclose(coeff, 0.0):
                terms.append(term_with_coeff(term_list[0], coeff))
    return PauliSum(terms)


def check_commutation(pauli_list: Sequence[PauliTerm], pauli_two: PauliTerm) -> bool:
    """Check if commuting a PauliTerm commutes with a list of other terms by natural calculation.

    Uses the result in Section 3 of arXiv:1405.5749v2, modified slightly here to check for the
    number of anti-coincidences (which must always be even for commuting PauliTerms)
    instead of the no. of coincidences, as in the paper.

    :param pauli_list: A list of PauliTerm objects
    :param pauli_two_term: A PauliTerm object
    :returns: True if pauli_two object commutes with pauli_list, False otherwise
    """

    def coincident_parity(p1: PauliTerm, p2: PauliTerm) -> bool:
        non_similar = 0
        p1_indices = set(p1._ops.keys())
        p2_indices = set(p2._ops.keys())
        for idx in p1_indices.intersection(p2_indices):
            if p1[idx] != p2[idx]:
                non_similar += 1
        return non_similar % 2 == 0

    for term in pauli_list:
        if not coincident_parity(term, pauli_two):
            return False
    return True


def commuting_sets(pauli_terms: PauliSum) -> list[list[PauliTerm]]:
    """Gather the Pauli terms of pauli_terms variable into commuting sets.

    Uses algorithm defined in (Raeisi, Wiebe, Sanders, arXiv:1108.4318, 2011)
    to find commuting sets. Except uses commutation check from arXiv:1405.5749v2

    :param pauli_terms: A PauliSum object
    :returns: List of lists where each list contains a commuting set
    """
    m_terms = len(pauli_terms.terms)
    m_s = 1
    groups = []
    groups.append([pauli_terms.terms[0]])
    for j in range(1, m_terms):
        isAssigned_bool = False
        for p in range(m_s):  # check if it commutes with each group
            if isAssigned_bool is False:
                if check_commutation(groups[p], pauli_terms.terms[j]):
                    isAssigned_bool = True
                    groups[p].append(pauli_terms.terms[j])
        if isAssigned_bool is False:
            m_s += 1
            groups.append([pauli_terms.terms[j]])
    return groups


def is_identity(term: PauliDesignator) -> bool:
    """Return True if PauliTerm or PauliSum is a scalar multiple of identity, False otherwise.

    :param term: Either a PauliTerm or PauliSum
    :returns: True if the PauliTerm or PauliSum is a scalar multiple of identity, False otherwise
    """
    if isinstance(term, PauliTerm):
        return (len(term) == 0) and (not np.isclose(term.coefficient, 0))
    elif isinstance(term, PauliSum):
        return (len(term.terms) == 1) and (len(term.terms[0]) == 0) and (not np.isclose(term.terms[0].coefficient, 0))
    else:
        raise TypeError("is_identity only checks PauliTerms and PauliSum objects!")


def exponentiate(term: PauliTerm) -> Program:
    """Create a pyQuil program that simulates the unitary evolution exp(-1j * term).

    :param term: A pauli term to exponentiate
    :returns: A Program object
    """
    return exponential_map(term)(1.0)


def exponential_map(term: PauliTerm) -> Callable[[float], Program]:
    """Return a function f(alpha) that constructs the Program corresponding to exp(-1j*alpha*term).

    :param term: A pauli term to exponentiate
    :returns: A function that takes an angle parameter and returns a program.
    """
    if not isinstance(term.coefficient, Complex):
        raise TypeError(f"PauliTerm coefficient type must be complex, got {type(term.coefficient)}")

    if not np.isclose(np.imag(term.coefficient), 0.0):
        raise TypeError("PauliTerm coefficient must be real")

    coeff = term.coefficient.real
    term.coefficient = term.coefficient.real

    def exp_wrap(param: float) -> Program:
        prog = Program()
        if is_identity(term):
            prog.inst(X(0))
            prog.inst(PHASE(-param * coeff, 0))
            prog.inst(X(0))
            prog.inst(PHASE(-param * coeff, 0))
        elif is_zero(term):
            pass
        else:
            prog += _exponentiate_general_case(term, param)
        return prog

    return exp_wrap


def exponentiate_commuting_pauli_sum(
    pauli_sum: PauliSum,
) -> Callable[[float], Program]:
    """Return a function that maps all substituent PauliTerms and sums them into a program.

    NOTE: Use this function with care. Substituent PauliTerms should commute.

    :param pauli_sum: PauliSum to exponentiate.
    :returns: A function that parametrizes the exponential.
    """
    if not isinstance(pauli_sum, PauliSum):
        raise TypeError("Argument 'pauli_sum' must be a PauliSum.")

    fns = [exponential_map(term) for term in pauli_sum]

    def combined_exp_wrap(param: float) -> Program:
        return Program([f(param) for f in fns])

    return combined_exp_wrap


def exponentiate_pauli_sum(
    pauli_sum: Union[PauliSum, PauliTerm],
) -> NDArray[np.complex_]:
    r"""Exponentiate a sequence of PauliTerms, which may or may not commute.

    The Pauliterms must have fixed (non-parametric) coefficients. The coefficients are interpreted in cycles rather than
    radians or degrees.

    e^{-i\pi\sum_i \theta_i P_i}

    Thus, 1.0 corresponds to one full rotation.

    To produce a CZ gate:

    >>> from numpy import pi
    >>> phi = pi
    >>> coeff = phi / (-4 * pi)  # -0.25
    >>> hamiltonian = PauliTerm("Z", 0) * PauliTerm("Z", 1) - 1 * PauliTerm("Z", 0) - 1 * PauliTerm("Z", 1)
    >>> exponentiate_pauli_sum(coeff * hamiltonian)
    array([[...]])

    To produce the Quil XY(theta) gate, you can use:

    >>> theta = pi / 2
    >>> coeff = theta / (-4 * pi)
    >>> hamiltonian = PauliTerm("X", 0) * PauliTerm("X", 1) + PauliTerm("Y", 0) * PauliTerm("Y", 1)
    >>> exponentiate_pauli_sum(coeff * hamiltonian)
    array([[...]])

    A global phase is applied to the unitary such that the [0,0] element is always real.

    :param pauli_sum: PauliSum to exponentiate.
    :returns: The matrix exponent of the PauliSum
    """
    if isinstance(pauli_sum, PauliTerm):
        pauli_sum = PauliSum([pauli_sum])

    pauli_matrices = {
        "X": np.array([[0.0, 1.0], [1.0, 0.0]]),
        "Y": np.array([[0.0, 0.0 - 1.0j], [0.0 + 1.0j, 0.0]]),
        "Z": np.array([[1.0, 0.0], [0.0, -1.0]]),
        "I": np.array([[1.0, 0.0], [0.0, 1.0]]),
    }

    qubits = pauli_sum.get_qubits()
    for q in qubits:
        if not isinstance(q, int):
            raise ValueError("PauliSum must be defined on qubits, not placeholders.")
    qubits.sort()

    matrices = []
    for term in pauli_sum.terms:
        coeff = term.coefficient
        if not isinstance(coeff, Number):
            raise ValueError("PauliSum must have fixed coefficients.")
        qubit_paulis = {qubit: pauli for qubit, pauli in term.operations_as_set()}
        paulis = [qubit_paulis[q] if q in qubit_paulis else "I" for q in qubits]
        matrix = float(np.real(coeff)) * reduce(np.kron, [pauli_matrices[p] for p in paulis])
        matrices.append(matrix)
    generated_unitary = expm(-1j * np.pi * sum(matrices))
    phase = np.exp(-1j * np.angle(generated_unitary[0, 0]))
    return np.asarray(phase * generated_unitary, dtype=np.complex_)


def _exponentiate_general_case(pauli_term: PauliTerm, param: float) -> Program:
    """Return a Quil (Program()) object corresponding to the exponential of the pauli_term object.

    For example, exp[-1.0j * param * pauli_term]

    :param pauli_term: A PauliTerm to exponentiate
    :param param: scalar, non-complex, value
    :returns: A Quil program object
    """

    def reverse_hack(p: Program) -> Program:
        # A hack to produce a *temporary* program which reverses p.
        revp = Program()
        revp.inst(list(reversed(p.instructions)))
        return revp

    quil_prog = Program()
    change_to_z_basis = Program()
    change_to_original_basis = Program()
    cnot_seq = Program()
    prev_index = None
    highest_target_index = None

    for index, op in pauli_term:
        if not isinstance(index, (int, QubitPlaceholder)):
            raise ValueError("PauliTerm indices must be qubits or qubit placeholders, got {type(index)}.")
        if "X" == op:
            change_to_z_basis.inst(H(index))
            change_to_original_basis.inst(H(index))
        elif "Y" == op:
            change_to_z_basis.inst(RX(np.pi / 2.0, index))
            change_to_original_basis.inst(RX(-np.pi / 2.0, index))
        elif "I" == op:
            continue

        if prev_index is not None:
            cnot_seq.inst(CNOT(prev_index, index))

        prev_index = index
        highest_target_index = index

    # building rotation circuit
    quil_prog += change_to_z_basis
    quil_prog += cnot_seq
    if not isinstance(pauli_term.coefficient, Complex) or highest_target_index is None:
        raise ValueError("PauliTerm coefficient must be complex and highest_target_index must be set.")
    quil_prog.inst(RZ(2.0 * pauli_term.coefficient * param, highest_target_index))
    quil_prog += reverse_hack(cnot_seq)
    quil_prog += change_to_original_basis

    return quil_prog


def suzuki_trotter(trotter_order: int, trotter_steps: int) -> list[tuple[float, int]]:
    """Generate trotterization coefficients for a given number of Trotter steps.

    U = exp(A + B) is approximated as exp(w1*o1)exp(w2*o2)... This method returns
    a list [(w1, o1), (w2, o2), ... , (wm, om)] of tuples where o=0 corresponds
    to the A operator, o=1 corresponds to the B operator, and w is the
    coefficient in the exponential. For example, a second order Suzuki-Trotter
    approximation to exp(A + B) results in the following
    [(0.5/trotter_steps, 0), (1/trotter_steps, 1),
    (0.5/trotter_steps, 0)] * trotter_steps.

    :param trotter_order: order of Suzuki-Trotter approximation
    :param trotter_steps: number of steps in the approximation
    :returns: List of tuples corresponding to the coefficient and operator
              type: o=0 is A and o=1 is B.
    """
    p1 = p2 = p4 = p5 = 1.0 / (4 - (4 ** (1.0 / 3)))
    p3 = 1 - 4 * p1
    trotter_dict: dict[int, list[tuple[float, int]]] = {
        1: [(1, 0), (1, 1)],
        2: [(0.5, 0), (1, 1), (0.5, 0)],
        3: [
            (7.0 / 24, 0),
            (2.0 / 3.0, 1),
            (3.0 / 4.0, 0),
            (-2.0 / 3.0, 1),
            (-1.0 / 24, 0),
            (1.0, 1),
        ],
        4: [
            (p5 / 2, 0),
            (p5, 1),
            (p5 / 2, 0),
            (p4 / 2, 0),
            (p4, 1),
            (p4 / 2, 0),
            (p3 / 2, 0),
            (p3, 1),
            (p3 / 2, 0),
            (p2 / 2, 0),
            (p2, 1),
            (p2 / 2, 0),
            (p1 / 2, 0),
            (p1, 1),
            (p1 / 2, 0),
        ],
    }

    order_slices = [(x0 / trotter_steps, x1) for x0, x1 in trotter_dict[trotter_order]]
    order_slices = order_slices * trotter_steps
    return order_slices


def is_zero(pauli_object: PauliDesignator) -> bool:
    """Return True if a PauliTerm or PauliSum is zero, False otherwise.

    :param pauli_object: Either a PauliTerm or PauliSum
    :returns: True if PauliTerm is zero, False otherwise
    """
    if isinstance(pauli_object, PauliTerm):
        if not isinstance(pauli_object.coefficient, Complex):
            raise ValueError("PauliTerm coefficient must be fixed and complex.")
        return bool(np.isclose(pauli_object.coefficient, 0))
    elif isinstance(pauli_object, PauliSum):
        if not isinstance(pauli_object.terms[0].coefficient, Complex):
            raise ValueError("PauliTerm coefficient must be fixed and complex.")
        return len(pauli_object.terms) == 1 and bool(np.isclose(pauli_object.terms[0].coefficient, 0))
    else:
        raise TypeError("is_zero only checks PauliTerms and PauliSum objects!")


def trotterize(
    first_pauli_term: PauliTerm,
    second_pauli_term: PauliTerm,
    trotter_order: int = 1,
    trotter_steps: int = 1,
) -> Program:
    """Create a Quil program that approximates exp( (A + B)t) where A and B are PauliTerm operators.

    :param first_pauli_term: PauliTerm denoted `A`
    :param second_pauli_term: PauliTerm denoted `B`
    :param trotter_order: Optional argument indicating the Suzuki-Trotter
                          approximation order--only accepts orders 1, 2, 3, 4.
    :param trotter_steps: Optional argument indicating the number of products
                          to decompose the exponential into.

    :return: Quil program
    """
    if not (1 <= trotter_order < 5):
        raise ValueError("trotterize only accepts trotter_order in {1, 2, 3, 4}.")

    commutator = (first_pauli_term * second_pauli_term) + (-1 * second_pauli_term * first_pauli_term)

    prog = Program()
    if is_zero(commutator):
        param_exp_prog_one = exponential_map(first_pauli_term)
        exp_prog = param_exp_prog_one(1)
        prog += exp_prog
        param_exp_prog_two = exponential_map(second_pauli_term)
        exp_prog = param_exp_prog_two(1)
        prog += exp_prog
        return prog

    order_slices = suzuki_trotter(trotter_order, trotter_steps)
    for coeff, operator in order_slices:
        if operator == 0:
            param_prog = exponential_map(coeff * first_pauli_term)
            exp_prog = param_prog(1)
            prog += exp_prog
        else:
            param_prog = exponential_map(coeff * second_pauli_term)
            exp_prog = param_prog(1)
            prog += exp_prog
    return prog
