# Copyright 2021 qclib project.

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


"""
Linear-depth Multicontrolled Special Unitary
"""

from typing import Union, List

import numpy as np
from qiskit import QuantumCircuit, QuantumRegister
from qiskit.circuit.library import RZGate, RYGate
from qiskit.circuit.library import UnitaryGate
from qiskit.synthesis import OneQubitEulerDecomposer
from qiskit.circuit import Gate, Qubit

from qclib.gates.mcx import LinearMcx, McxVchainDirty
from qclib.gates.util import check_su2, apply_ctrl_state, isclose

# pylint: disable=protected-access


class Ldmcsu(Gate):
    """
    Linear depth Multi-Controlled Gate for Special Unitary
    ------------------------------------------------

    Multicontrolled gate decomposition with linear cost.
    `unitary` must be a SU(2) matrix.
    """

    def __init__(self, unitary, num_controls, ctrl_state: str = None):

        check_su2(unitary)
        self.unitary = unitary
        self.controls = QuantumRegister(num_controls)
        self.target = QuantumRegister(1)
        self.num_controls = num_controls + 1
        self.ctrl_state = ctrl_state

        super().__init__("ldmcsu", self.num_controls, [], "ldmcsu")

    def _define(self):

        self.definition = QuantumCircuit(self.controls, self.target)

        is_main_diag_real = isclose(self.unitary[0, 0].imag, 0.0) and isclose(
            self.unitary[1, 1].imag, 0.0
        )
        is_secondary_diag_real = isclose(self.unitary[0, 1].imag, 0.0) and isclose(
            self.unitary[1, 0].imag, 0.0
        )

        if not is_main_diag_real and not is_secondary_diag_real:
            # U = V D V^-1, where the entries of the diagonal D are the eigenvalues
            # `eig_vals` of U and the column vectors of V are the eigenvectors
            # `eig_vecs` of U. These columns are orthonormal and the main diagonal
            # of V is real-valued.
            eig_vals, eig_vecs = np.linalg.eig(self.unitary)

            x_vecs, z_vecs = self._get_x_z(eig_vecs)

            self.half_linear_depth_mcv(
                x_vecs,
                z_vecs,
                self.controls,
                self.target,
                self.ctrl_state,
                inverse=True,
            )
            self.linear_depth_mcv(
                np.diag(eig_vals),
                self.controls,
                self.target,
                self.ctrl_state,
                general_su2_optimization=True,
            )
            self.half_linear_depth_mcv(
                x_vecs, z_vecs, self.controls, self.target, self.ctrl_state
            )

        else:

            if not is_secondary_diag_real:
                self.definition.h(self.target)

            self.linear_depth_mcv(self.unitary, self.controls, self.target, self.ctrl_state)

            if not is_secondary_diag_real:
                self.definition.h(self.target)

    @staticmethod
    def _get_x_z(su2):
        is_secondary_diag_real = isclose(su2[0, 1].imag, 0.0) and isclose(
            su2[1, 0].imag, 0.0
        )

        if is_secondary_diag_real:
            x_value = su2[0, 1]
            z_value = su2[1, 1]
        else:
            x_value = -su2[0, 1].real
            z_value = su2[1, 1] - su2[0, 1].imag * 1.0j

        return x_value, z_value

    @staticmethod
    def _compute_gate_a(x_value, z_value):
        if x_value == 0:
            alpha = (z_value + 0j) ** (1 / 4)
            beta = 0.0
        else:
            alpha_r = np.sqrt((np.sqrt((z_value.real + 1.0) / 2.0) + 1.0) / 2.0)
            alpha_i = z_value.imag / (
                    2.0 * np.sqrt((z_value.real + 1.0) *
                                  (np.sqrt((z_value.real + 1.0) / 2.0) + 1.0))
            )
            alpha = alpha_r + 1.0j * alpha_i
            beta = x_value / (
                    2.0 * np.sqrt((z_value.real + 1.0) *
                                  (np.sqrt((z_value.real + 1.0) / 2.0) + 1.0))
            )
        s_op = np.array([[alpha, -np.conj(beta)], [beta, np.conj(alpha)]])
        return s_op

    def linear_depth_mcv(
        self,
        su2_unitary,
        controls: Union[QuantumRegister, List[Qubit]],
        target: Qubit,
        ctrl_state: str = None,
        general_su2_optimization=False,
    ):
        """
        Theorem 1 - https://arxiv.org/pdf/2302.06377.pdf
        """
        # S gate definition

        x_value, z_value = self._get_x_z(su2_unitary)

        op_a = Ldmcsu._compute_gate_a(x_value, z_value)
        gate_a = UnitaryGate(op_a)

        num_ctrl = len(controls)
        k_1 = int(np.ceil(num_ctrl / 2.0))
        k_2 = int(np.floor(num_ctrl / 2.0))

        ctrl_state_k_1 = None
        ctrl_state_k_2 = None

        if ctrl_state is not None:
            ctrl_state_k_1 = ctrl_state[::-1][:k_1][::-1]
            ctrl_state_k_2 = ctrl_state[::-1][k_1:][::-1]

        if not general_su2_optimization:
            mcx_1 = McxVchainDirty(k_1, ctrl_state=ctrl_state_k_1).definition
            self.definition.append(
                mcx_1, controls[:k_1] + controls[k_1 : 2 * k_1 - 2] + [target]
            )
        self.definition.append(gate_a, [target])

        mcx_2 = McxVchainDirty(
            k_2, ctrl_state=ctrl_state_k_2, action_only=general_su2_optimization
        ).definition
        self.definition.append(
            mcx_2.inverse(), controls[k_1:] + controls[k_1 - k_2 + 2 : k_1] + [target]
        )
        self.definition.append(gate_a.inverse(), [target])

        mcx_3 = McxVchainDirty(k_1, ctrl_state=ctrl_state_k_1).definition
        self.definition.append(
            mcx_3, controls[:k_1] + controls[k_1 : 2 * k_1 - 2] + [target]
        )
        self.definition.append(gate_a, [target])

        mcx_4 = McxVchainDirty(k_2, ctrl_state=ctrl_state_k_2).definition
        self.definition.append(
            mcx_4, controls[k_1:] + controls[k_1 - k_2 + 2 : k_1] + [target]
        )
        self.definition.append(gate_a.inverse(), [target])

    def half_linear_depth_mcv(
        self,
        x_value,
        z_value,
        controls: Union[QuantumRegister, List[Qubit]],
        target: Qubit,
        ctrl_state: str = None,
        inverse: bool = False,
    ):
        """
        Theorem 4 - https://arxiv.org/pdf/2302.06377.pdf
        """

        alpha_r = np.sqrt((z_value.real + 1.0) / 2.0)
        alpha_i = z_value.imag / np.sqrt(2 * (z_value.real + 1.0))
        alpha = alpha_r + 1.0j * alpha_i

        beta = x_value / np.sqrt(2 * (z_value.real + 1.0))

        s_op = np.array([[alpha, -np.conj(beta)], [beta, np.conj(alpha)]])

        # S gate definition
        s_gate = UnitaryGate(s_op)

        # Hadamard equivalent definition
        h_gate = UnitaryGate(np.array([[-1, 1], [1, 1]]) * 1 / np.sqrt(2))

        num_ctrl = len(controls)
        k_1 = int(np.ceil(num_ctrl / 2.0))
        k_2 = int(np.floor(num_ctrl / 2.0))

        ctrl_state_k_1 = None
        ctrl_state_k_2 = None

        if ctrl_state is not None:
            ctrl_state_k_1 = ctrl_state[::-1][:k_1][::-1]
            ctrl_state_k_2 = ctrl_state[::-1][k_1:][::-1]

        if inverse:
            self.definition.h(target)

            self.definition.append(s_gate, [target])
            mcx_2 = McxVchainDirty(
                k_2, ctrl_state=ctrl_state_k_2, action_only=True
            ).definition
            self.definition.append(
                mcx_2, controls[k_1:] + controls[k_1 - k_2 + 2 : k_1] + [target]
            )

            self.definition.append(s_gate.inverse(), [target])

            self.definition.append(h_gate, [target])

        else:
            mcx_1 = McxVchainDirty(k_1, ctrl_state=ctrl_state_k_1).definition
            self.definition.append(
                mcx_1, controls[:k_1] + controls[k_1 : 2 * k_1 - 2] + [target]
            )
            self.definition.append(h_gate, [target])

            self.definition.append(s_gate, [target])

            mcx_2 = McxVchainDirty(k_2, ctrl_state=ctrl_state_k_2).definition
            self.definition.append(
                mcx_2, controls[k_1:] + controls[k_1 - k_2 + 2 : k_1] + [target]
            )
            self.definition.append(s_gate.inverse(), [target])

            self.definition.h(target)

    @staticmethod
    def ldmcsu(
        circuit,
        unitary,
        controls: Union[QuantumRegister, List[Qubit]],
        target: Qubit,
        ctrl_state: str = None,
    ):
        """
        Apply multi-controlled SU(2)
        https://arxiv.org/abs/2302.06377
        """
        circuit.append(
            Ldmcsu(unitary, len(controls), ctrl_state=ctrl_state), [*controls, target]
        )


class LdMcSpecialUnitary(Gate):
    """
    Linear-depth Multicontrolled Special Unitary
    --------------------------------------------

    Implements the gate decompostion of any gate in SU(2) with linear depth (Ld)
    presented in Lemma 7.9 in Barenco et al., 1995 (arXiv:quant-ph/9503016)
    with optimizations from Theorem 5 of Iten et al., 2016 (arXiv:1501.06911)
    """

    def __init__(self, unitary, num_controls, ctrl_state=None):

        if not check_su2(unitary):
            raise ValueError("Operator must be in SU(2)")

        self.unitary = np.array(unitary, dtype=complex)

        if num_controls > 0:
            self.control_qubits = QuantumRegister(num_controls)
        else:
            self.control_qubits = []

        self.target_qubit = QuantumRegister(1)
        self.num_qubits = num_controls + 1
        self.ctrl_state = ctrl_state

        if self.ctrl_state is None:
            self.ctrl_state = "1" * num_controls

        super().__init__("ldmc_su2", self.num_qubits, [], "LdMcSu2")

    @staticmethod
    def get_abc_operators(beta, gamma, delta):
        """
        Creates A,B and C matrices such that
        ABC = I
        """
        # A
        a_rz = RZGate(beta).to_matrix()
        a_ry = RYGate(gamma / 2).to_matrix()
        a_matrix = a_rz.dot(a_ry)

        # B
        b_ry = RYGate(-gamma / 2).to_matrix()
        b_rz = RZGate(-(delta + beta) / 2).to_matrix()
        b_matrix = b_ry.dot(b_rz)

        # C
        c_matrix = RZGate((delta - beta) / 2).to_matrix()

        a_gate = UnitaryGate(a_matrix, label="A")
        b_gate = UnitaryGate(b_matrix, label="B")
        c_gate = UnitaryGate(c_matrix, label="C")

        return a_gate, b_gate, c_gate

    def _define(self):
        self.definition = QuantumCircuit(self.control_qubits, self.target_qubit)

        if len(self.control_qubits) > 0:
            self._apply_ctrl_state()

            theta, phi, lamb, _ = OneQubitEulerDecomposer._params_zyz(self.unitary)

            a_gate, b_gate, c_gate = LdMcSpecialUnitary.get_abc_operators(
                phi, theta, lamb
            )

            self._apply_abc(a_gate, b_gate, c_gate)

            self._apply_ctrl_state()
        else:
            self.unitary(self.unitary, self.target_qubit)

    def _apply_abc(self, a_gate: UnitaryGate, b_gate: UnitaryGate, c_gate: UnitaryGate):
        """
        Applies ABC matrices to the quantum circuit according to theorem 5
        of Iten et al. 2016 (arXiv:1501.06911).
        Parameters
        ----------
            a_gate, b_gate and c_gate expceted to be special unitary gates
        """

        if len(self.control_qubits) < 3:
            self.definition.append(c_gate, [self.target_qubit])
            self.definition.mcx(self.control_qubits, self.target_qubit)
            self.definition.append(b_gate, [self.target_qubit])
            self.definition.mcx(self.control_qubits, self.target_qubit)
            self.definition.append(a_gate, [self.target_qubit])
        else:
            ancilla = self.control_qubits[-1]
            action_only = True

            if len(self.control_qubits) < 6:
                action_only = False

            # decompose A, B and C to use their optimized controlled versions
            theta_a, phi_a, lam_a, _ = OneQubitEulerDecomposer._params_zyz(
                a_gate.to_matrix()
            )
            theta_b, phi_b, lam_b, _ = OneQubitEulerDecomposer._params_zyz(
                b_gate.to_matrix()
            )
            theta_c, phi_c, lam_c, _ = OneQubitEulerDecomposer._params_zyz(
                c_gate.to_matrix()
            )
            a_a, b_a, c_a = LdMcSpecialUnitary.get_abc_operators(phi_a, theta_a, lam_a)
            a_b, b_b, c_b = LdMcSpecialUnitary.get_abc_operators(phi_b, theta_b, lam_b)
            a_c, b_c, c_c = LdMcSpecialUnitary.get_abc_operators(phi_c, theta_c, lam_c)

            # definition of left mcx, which will also be inverted as the right mcx
            mcx_gate = LinearMcx(
                len(self.control_qubits[:-1]), action_only=action_only
            ).definition

            # decomposed controlled C
            self.definition.unitary(c_c, self.target_qubit)
            self.definition.cx(ancilla, self.target_qubit)
            self.definition.unitary(b_c, self.target_qubit)
            self.definition.cx(ancilla, self.target_qubit)
            self.definition.unitary(a_c, self.target_qubit)

            self.definition.append(
                mcx_gate, self.control_qubits[:-1] + [self.target_qubit] + [ancilla]
            )

            # decomposed controlled B
            self.definition.unitary(c_b, self.target_qubit)
            self.definition.cx(ancilla, self.target_qubit)
            self.definition.unitary(b_b, self.target_qubit)
            self.definition.cx(ancilla, self.target_qubit)
            self.definition.unitary(a_b, self.target_qubit)

            self.definition.append(
                mcx_gate.inverse(),
                self.control_qubits[:-1] + [self.target_qubit] + [ancilla],
            )

            # decomposed A
            self.definition.unitary(c_a, self.target_qubit)
            self.definition.cx(ancilla, self.target_qubit)
            self.definition.unitary(b_a, self.target_qubit)
            self.definition.cx(ancilla, self.target_qubit)
            self.definition.unitary(a_a, self.target_qubit)

    @staticmethod
    def ldmcsu(circuit, unitary, controls, target, ctrl_state=None):
        """
        Linear-depth Multicontrolled Special Unitary
        --------------------------------------------

        Implements the gate decompostion of any gate in SU(2) with linear depth (Ld)
        presented in Lemma 7.9 in Barenco et al., 1995 (arXiv:quant-ph/9503016)
        with optimizations from Theorem 5 of Iten et al., 2016 (arXiv:1501.06911)
        """
        circuit.append(
            LdMcSpecialUnitary(unitary, len(controls), ctrl_state), [*controls, target]
        )


LdMcSpecialUnitary._apply_ctrl_state = apply_ctrl_state
