# PyZX - Python library for quantum circuit rewriting
#        and optimization using the ZX-calculus
# Copyright (C) 2018 - Aleks Kissinger and John van de Wetering

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


from typing import Any, Dict, Iterable, List, Optional, Union
from pyzx_param.circuit import Circuit, Gate, gate_types, CNOT
from pyzx_param.linalg import Z2, Mat2, MatLike

import numpy as np

class CNOT_tracker(Circuit):
    """
    A circuit-like object that keeps track of row and column operations
    applied during Gauss elimination via CNOT gates.
    """

    matrix: Mat2
    """The qubit parity matrix computed from the CNOT gates."""

    row_perm: np.ndarray
    """The row permutation of the qubit parity matrix."""

    col_perm: np.ndarray
    """The column permutation of the qubit parity matrix."""

    def __init__(self, n_qubits: int, **kwargs):
        """
        Creates and returns a circuit-like object.

        :param n_qubits: The number of qubits
        :param **kwargs: Additional arguments passed
        """
        super().__init__(n_qubits, **kwargs)
        #self.matrix = Mat2(np.identity(n_qubits, dtype=np.int32).tolist())
        self.matrix = Mat2.id(n_qubits)
        self.row_perm = np.arange(n_qubits)
        self.col_perm = np.arange(n_qubits)
        self.n_qubits = n_qubits

    def count_cnots(self) -> int:
        """Returns the number of CNOT gates in the tracker."""
        return len(
            [g for g in self.gates if hasattr(g, "name") and g.name in ["CNOT", "CZ"]]
        )

    def cnot_depth(self) -> int:
        """Returns the CNOT/CZ depth of the tracked circuit."""
        depth = 1
        previous_gates: List[int] = []
        for g in self.gates:
            if hasattr(g, "name") and g.name in ["CNOT", "CZ"]:
                if g.control in previous_gates or g.target in previous_gates:  # type: ignore # Overlapping gate
                    # Start a new CNOT layer
                    previous_gates = []
                    depth += 1
                else:
                    previous_gates += [g.control, g.target]  # type: ignore
        return depth

    def row_add(self, q0: int, q1: int):
        """Track a row addition operation on the matrix."""
        self.add_gate("CNOT", q0, q1)
        self.matrix.row_add(q0, q1)

    def add_gate(self, gate: Union[Gate,str], *args, **kwargs):
        """ Adds a gate to the circuit, if it is a CNOT gate it will track a row addition operation on the matrix. """
        if isinstance(gate, CNOT):
            self.row_add(gate.control, gate.target)
        else:
            super().add_gate(gate, *args, **kwargs)

    def col_add(self, q0: int, q1: int):
        """Track a column addition operation on the matrix."""
        self.prepend_gate("CNOT", q1, q0)
        self.matrix.col_add(q0, q1)

    @staticmethod
    def get_metric_names() -> List[str]:
        """Metric names for the CNOT tracker."""
        return ["n_cnots", "depth"]

    def gather_metrics(self) -> Dict[str, int]:
        """Gather metrics for the CNOT tracker."""
        metrics = {}
        metrics["n_cnots"] = self.count_cnots()
        metrics["depth"] = self.cnot_depth()
        return metrics

    def prepend_gate(self, gate: Union[Gate, str], *args, **kwargs):
        """Adds a gate to the circuit. ``gate`` can either be
        an instance of a :class:`Gate`, or it can be the name of a gate,
        in which case additional arguments should be given.

        Example::

            circuit.add_gate("CNOT", 1, 4) # adds a CNOT gate with control 1 and target 4
            circuit.add_gate("ZPhase", 2, phase=Fraction(3,4)) # Adds a ZPhase gate on qubit 2 with phase 3/4
        """
        if isinstance(gate, str):
            gate_class = gate_types[gate]
            gate = gate_class(*args, **kwargs)
        self.gates.insert(0, gate)

    def to_qasm(self, version: int = 2) -> str:
        """
        Outputs a QASM of the current circuit.

        :param version: The QASM version to be used, default 2
        :return: The initial wiring, the resulting wiring and the full QASM circuit
        """
        qasm = super().to_qasm(version)
        initial_perm = "// Initial wiring: " + str(self.row_perm)
        end_perm = "// Resulting wiring: " + str(self.col_perm)
        return "\n".join([initial_perm, end_perm, qasm])

    @staticmethod
    def from_circuit(circuit: Circuit) -> "CNOT_tracker":
        """
        Takes a circuit and outputs a CNOT_tracker of the circuit.

        :param circuit: A circuit to be converted
        :return: A CNOT_tracker of the circuit given
        """
        new_circuit = CNOT_tracker(circuit.qubits, name=circuit.name)
        new_circuit.gates = circuit.gates
        new_circuit.update_matrix()
        return new_circuit

    def update_matrix(self):
        """Rebuilds the parity matrix from the gates in the circuit."""
        self.matrix = Mat2.id(self.n_qubits)
        for gate in self.gates:
            if hasattr(gate, "name") and gate.name == "CNOT":
                self.matrix.row_add(gate.control, gate.target)
            else:
                print(
                    "Warning: CNOT tracker can only be used for circuits with only CNOT gates!"
                )

    @staticmethod
    def from_qasm_file(fname: str) -> "CNOT_tracker":
        """
        Takes the filename of a QASM and outputs a CNOT_tracker of it.

        :param fname: The filename of the QASM
        :return: A CNOT_tracker of the QASM given
        """
        circuit = Circuit.from_qasm_file(fname)
        return CNOT_tracker.from_circuit(circuit)

class Parity:
    """
    A set of qubits XORed together.
    """
    parity: List[bool]

    def __init__(self, par: Union[str, int, Iterable[Any]], n_qubits: Optional[int] = None):
        """
        Creates and returns a set of qubits XORed together.

        :param par: The input parity pattern
        :param n_qubits: An optional array of the number of qubits, default None
        """
        if isinstance(par, int):
            self.parity = [bool(par & (1 << i)) for i in reversed(range(par.bit_length()))]
        elif isinstance(par, str):
            self.parity = [p != "0" for p in par]
        else:
            self.parity = [bool(x) for x in par]
        if n_qubits is not None and len(self.parity) < n_qubits:
            self.parity += [False] * (n_qubits - len(self.parity))

    def count(self) -> int:
        """Returns the number of qubits interacting in the parity."""
        return sum(self.parity)

    def n_qubits(self) -> int:
        """Returns the total number of qubits."""
        return len(self.parity)
    
    def to_mat2_row(self) -> List[Z2]:
        """Returns an array of 1s and 0s based off the internal parity array."""
        return [1 if b else 0 for b in self.parity]

    def __str__(self):
        """Returns a string of 1s and 0s based off the internal parity array."""
        return "".join(["1" if p else "0" for p in self.parity])

    def __repr__(self):
        """Returns a string representation of the parity."""
        return "Parity(" + str(self) + ")"

    def __len__(self):
        """Returns the length of the parity."""
        return len(self.parity)

    def __iter__(self):
        """Makes the parity iterable."""
        return iter(self.parity)

    def __getitem__(self, i):
        """Returns the item at index i in the parity."""
        return self.parity[i]

    def __set_item__(self, i, v):
        """Sets the item at index i to the value given by v in the parity."""
        self.parity[i] = v

    def __eq__(self, other):
        """Returns true if other is equal to the parity."""
        if isinstance(other, Parity):
            return self.parity == other.parity
        if isinstance(other, str):
            return str(self) == other
        return False

    def __hash__(self):
        """Returns a hash of the parity."""
        return hash(str(self))
