# Copyright 2021 Xanadu Quantum Technologies Inc.

# 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.

"""
This module implements the :class:`.Circuit` class which acts as a representation for quantum circuits.
"""

from __future__ import annotations

__all__ = ["Circuit"]

from typing import List, Optional, Tuple

from mrmustard import settings
from mrmustard.lab.abstract import State, Transformation
from mrmustard.utils.typing import RealMatrix, RealVector
from mrmustard.lab.circuit_drawer import circuit_text
from mrmustard.math.tensor_wrappers import XPMatrix, XPVector

import numpy as np


class Circuit(Transformation):
    """Represents a quantum circuit: a set of operations to be applied on quantum states.

    Args:
        ops (list or none): A list of operations comprising the circuit.
    """

    def __init__(self, ops: Optional[List] = None):
        self._ops = list(ops) if ops is not None else []
        super().__init__(name="Circuit")
        self.reset()

    @property
    def ops(self) -> Optional[List]:
        r"""
        The list of operations comprising the circuit.
        """
        return self._ops

    def reset(self):
        """Resets the state of the circuit clearing the list of modes and setting the compiled flag to false."""
        self._compiled: bool = False
        self._modes: List[int] = []

    @property
    def num_modes(self) -> int:
        all_modes = {mode for op in self._ops for mode in op.modes}
        return len(all_modes)

    def primal(self, state: State) -> State:
        for op in self._ops:
            state = op.primal(state)
        return state

    def dual(self, state: State) -> State:
        for op in reversed(self._ops):
            state = op.dual(state)
        return state

    def XYd(
        self,
        allow_none: bool = True,
    ) -> Tuple[
        RealMatrix, RealMatrix, RealVector
    ]:  # NOTE: Overriding Transformation.XYd for efficiency
        X = XPMatrix(like_1=True)
        Y = XPMatrix(like_0=True)
        d = XPVector()
        for op in self._ops:
            opx, opy, opd = op.XYd(allow_none)
            opX = XPMatrix.from_xxpp(opx, modes=(op.modes, op.modes), like_1=True)
            opY = XPMatrix.from_xxpp(opy, modes=(op.modes, op.modes), like_0=True)
            opd = XPVector.from_xxpp(opd, modes=op.modes)
            if opX.shape is not None and opX.shape[-1] == 1 and len(op.modes) > 1:
                opX = opX.clone(len(op.modes), modes=(op.modes, op.modes))
            if opY.shape is not None and opY.shape[-1] == 1 and len(op.modes) > 1:
                opY = opY.clone(len(op.modes), modes=(op.modes, op.modes))
            if opd.shape is not None and opd.shape[-1] == 1 and len(op.modes) > 1:
                opd = opd.clone(len(op.modes), modes=op.modes)
            X = opX @ X
            Y = opX @ Y @ opX.T + opY
            d = opX @ d + opd
        return X.to_xxpp(), Y.to_xxpp(), d.to_xxpp()

    @property
    def is_gaussian(self):
        """Returns `true` if all operations in the circuit are Gaussian."""
        return all(op.is_gaussian for op in self._ops)

    @property
    def is_unitary(self):
        """Returns `true` if all operations in the circuit are unitary."""
        return all(op.is_unitary for op in self._ops)

    def value(self, shape: Tuple[int]):
        raise NotImplementedError

    def __len__(self):
        return len(self._ops)

    _repr_markdown_ = None

    def __repr__(self) -> str:
        """String to display the object on the command line."""
        # ops_repr = [repr(op) for op in self._ops]
        # return " >> ".join(ops_repr)
        return circuit_text(self._ops, decimals=settings.CIRCUIT_DECIMALS)

    def __str__(self):
        """String representation of the circuit."""
        ops_repr = [repr(op) for op in self._ops]
        return " >> ".join(ops_repr)

    # pylint: disable=too-many-branches,too-many-return-statements
    def __eq__(self, other):
        r"""Returns ``True`` if the two transformations are equal."""
        if not isinstance(other, Circuit):
            return False
        if not (self.is_gaussian and other.is_gaussian):
            return np.allclose(
                self.choi(cutoffs=[settings.EQ_TRANSFORMATION_CUTOFF] * 4 * self.num_modes),
                other.choi(cutoffs=[settings.EQ_TRANSFORMATION_CUTOFF] * 4 * self.num_modes),
                rtol=settings.EQ_TRANSFORMATION_RTOL_FOCK,
            )

        sX, sY, sd = self.XYd(allow_none=False)
        oX, oY, od = other.XYd(allow_none=False)
        return np.allclose(sX, oX) and np.allclose(sY, oY) and np.allclose(sd, od)
