"""
Parity Quantum Computing GmbH
Rennweg 1 Top 314
6020 Innsbruck, Austria


Copyright (c) 2020-2022.
All rights reserved.

Tools to export ParityOS circuits to Cirq.
"""
from math import pi
from typing import Any, Mapping, Union

from parityos.base.circuit import CircuitElement
from parityos.base.gates import Qubit, Gate, RMixin, CNOT, H, X, Y, Z, Rx, Ry, Rz, Rzz, CCNOT
from parityos.base.exceptions import ParityOSImportError

try:
    import cirq
except ImportError:
    raise ParityOSImportError("The Cirq exporter requires the installation of Cirq.")


GATE_MAP = {
    Gate: NotImplemented,  # Added to improve the type checking in PyCharm
    # All keys must be subclasses of the Gate class.
    CNOT: cirq.CNOT,
    X: cirq.X,
    Y: cirq.Y,
    Z: cirq.Z,
    H: cirq.H,
    Rx: cirq.XPowGate,  # We use _PowGate for consistency with Rzz: ZZPowGate
    Ry: cirq.YPowGate,  # We use _PowGate for consistency with Rzz: ZZPowGate
    Rz: cirq.ZPowGate,  # We use _PowGate for consistency with Rzz: ZZPowGate
    Rzz: cirq.ZZPowGate,  # cirq.Rzz does not exist.
    CCNOT: cirq.CCNOT
}

CirqCircuitElement = Union[cirq.GateOperation, "CirqCircuitElement"]


class CirqExporter:
    """
    Tool to convert ParityOS circuits to Cirq circuits.

    Instantiate the CirqExporter with a qubit map and a parameter map.
    Then use the `to_cirq` method to convert ParityOS circuits to Cirq circuits.

    EXAMPLE:
        qubit_map = {Qubit(i): cirq.NamedQubit(str(i)) for i in range(10)}
        parameter_map = {'theta': sympy.Symbol('theta'), 'gamma': sympy.Symbol('gamma')}
        cirq_exporter = CirqExporter(qubit_map, parameter_map)
        cirq_circuit = cirq_exporter.to_cirq(parityos_circuit)
    """

    def __init__(
        self,
        parameter_map: Mapping[str, Any] = None,
        qubit_map: Mapping[Qubit, cirq.NamedQubit] = None,
    ):
        """
        Converts the circuit to a cirq circuit.

        :param parameter_map: a mapping of the form {parameter_name: parameter_value}, where the
            parameter_name is a string that is used as a parameter_name in the ParityOS circuit,
            and parameter_value is a number like object (int, float, numpy float or a Sympy symbol
            are all valid). Optional.  If not given, then an empty dictionary is used instead.

        :param qubit_map: a mapping of the form {ParityOS_qubit: cirq_qubit}, where cirq_qubit is a
            Cirq NamedQubit instance. Optional. If not given, then ParityOS Qubits are automatically
            converted into cirq.NamedQubit instances with the same label.

        """
        self.parameter_map = {} if parameter_map is None else parameter_map
        self.qubit_map = qubit_map

    def to_cirq(self, circuit: CircuitElement) -> cirq.Circuit:
        """
        Converts the circuit to a Cirq circuit.

        :param circuit: a ParityOS circuit of quantum gates.
        :return: a Cirq circuit.
        """

        def _elements_to_cirq_gates(element: CircuitElement) -> CirqCircuitElement:
            """Recursive helper method for the to_cirq method."""
            return (
                self.gate_to_cirq(element)
                if isinstance(element, Gate)
                else [_elements_to_cirq_gates(item) for item in element]
            )

        return cirq.Circuit(_elements_to_cirq_gates(circuit))

    def gate_to_cirq(self, gate: Gate) -> cirq.GateOperation:
        """
        Converts a gate to a Cirq gate operation.

        :param gate: a ParityOS gate instance.
        :return: a Cirq gate operation.
        """
        cirq_operation = GATE_MAP[type(gate)]
        if isinstance(gate, RMixin):
            angle = (
                gate.angle
                if gate.parameter_name is None
                else gate.angle * self.parameter_map[gate.parameter_name]
            )
            cirq_operator = cirq_operation(exponent=angle / pi, global_shift=-0.5)
            # We use ZPowGate instead or rz, so we have to adjust the exponent by a factor pi
            # and the global phase with a global shift of -1/2.
            # See https://quantumai.google/reference/python/cirq/ZPowGate
        else:
            cirq_operator = cirq_operation

        parityos_qubits = (qubit for qubit in gate.make_args() if isinstance(qubit, Qubit))
        if self.qubit_map:
            cirq_qubits = (self.qubit_map[qubit] for qubit in parityos_qubits)
        else:
            cirq_qubits = (cirq.NamedQubit(str(qubit.label)) for qubit in parityos_qubits)

        return cirq_operator(*cirq_qubits)
