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

"""
    Divide-and-conquer state preparation
    https://arxiv.org/abs/2108.10182
"""

import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister

from qclib.state_preparation.util.state_tree_preparation import Amplitude, state_decomposition

from qclib.state_preparation.util.angle_tree_preparation import \
    create_angles_tree

from qclib.state_preparation.util.tree_register import add_register
from qclib.state_preparation.util.tree_walk import bottom_up

def initialize(state):
    """
        https://arxiv.org/abs/2108.10182
    """
    n_qubits = int(np.log2(len(state)))
    data = [Amplitude(i, a) for i, a in enumerate(state)]

    state_tree = state_decomposition(n_qubits, data)
    angle_tree = create_angles_tree(state_tree)

    circuit = QuantumCircuit()
    add_register(circuit, angle_tree, n_qubits-1)

    bottom_up(angle_tree, circuit, n_qubits)

    return circuit






class BinTree:
    """
        State preparation using DCSP https://www.nature.com/articles/s41598-021-85474-1.
        This class implements the original algorithm as defined in the paper.
        It is kept here for didactic reasons.
        The ``initialize`` function should preferably be used.
    """
    size = None
    values = None

    def __init__(self, values):
        self.size = len(values)
        self.values = values


    @staticmethod
    def parent(key):
        """ returns parent node index"""
        return int((key-0.5)/2)

    @staticmethod
    def left(key):
        """ returns left child index"""
        return int(2 * key + 1)

    @staticmethod
    def right(key):
        """ returns right child index"""
        return int(2 * key + 2)

    @staticmethod
    def root():
        """ returns root index """
        return 0

    def __getitem__(self, key):
        return self.values[key]


class Encoding:
    """
    Encode information into qubits
    """
    qcircuit = None
    quantum_data = None
    classical_data = None
    num_qubits = None
    tree = None
    output_qubits = []

    def __init__(self, input_vector, encode_type='amplitude_encoding'):
        if encode_type == 'amplitude_encoding':
            self.amplitude_encoding(input_vector)
        if encode_type == 'qubit_encoding':
            self.qubit_encoding(input_vector)
        if encode_type == 'dc_amplitude_encoding':
            self.dc_amplitude_encoding(input_vector)
        if encode_type == 'basis_encoding':
            self.basis_encoding(input_vector)

    def basis_encoding(self, input_vector, n_classical=1):
        """
        encoding a binary string x in a basis state |x>
        """
        self.num_qubits = int(len(input_vector))
        self.quantum_data = QuantumRegister(self.num_qubits)
        self.classical_data = ClassicalRegister(n_classical)
        self.qcircuit = QuantumCircuit(self.quantum_data, self.classical_data)
        for k, _ in enumerate(input_vector):
            if input_vector[k] == 1:
                self.qcircuit.x(self.quantum_data[k])


    def qubit_encoding(self, input_vector, n_classical=1):
        """
        encoding a binary string x as
        """
        input_pattern = QuantumRegister(len(input_vector))
        classical_register = ClassicalRegister(n_classical)
        self.qcircuit = QuantumCircuit(input_pattern, classical_register)
        for k, _ in enumerate(input_vector):
            self.qcircuit.ry(input_vector[k], input_pattern[k])

    @staticmethod
    def _recursive_compute_beta(input_vector, betas):
        if len(input_vector) > 1:
            new_x = []
            beta = []
            for k in range(0, len(input_vector), 2):
                norm = np.sqrt(input_vector[k] ** 2 + input_vector[k + 1] ** 2)
                new_x.append(norm)
                if norm == 0:
                    beta.append(0)
                else:
                    if input_vector[k] < 0:
                        ## testing
                        beta.append(2 * np.pi - 2 * np.arcsin(input_vector[k + 1] / norm))
                    else:
                        beta.append(2 * np.arcsin(input_vector[k + 1] / norm))
            Encoding._recursive_compute_beta(new_x, betas)
            betas.append(beta)

    @staticmethod
    def _index(k, circuit, control_qubits, numberof_controls):
        binary_index = '{:0{}b}'.format(k, numberof_controls)
        for j, qbit in enumerate(control_qubits):
            if binary_index[j] == '1':
                circuit.x(qbit)


    def amplitude_encoding(self, input_vector):
        """
        load real vector x to the amplitude of a quantum state
        """
        self.num_qubits = int(np.log2(len(input_vector)))
        self.quantum_data = QuantumRegister(self.num_qubits)
        self.qcircuit = QuantumCircuit(self.quantum_data)
        newx = np.copy(input_vector)
        betas = []
        Encoding._recursive_compute_beta(newx, betas)
        self._generate_circuit(betas, self.qcircuit, self.quantum_data)

    def dc_amplitude_encoding(self, input_vector):
        """
        Divide and conquer amplitude encoding
        """
        self.num_qubits = int(len(input_vector))-1
        self.quantum_data = QuantumRegister(self.num_qubits)
        self.qcircuit = QuantumCircuit(self.quantum_data)
        newx = np.copy(input_vector)
        betas = []
        Encoding._recursive_compute_beta(newx, betas)
        self._dc_generate_circuit(betas, self.qcircuit, self.quantum_data)

    def _dc_generate_circuit(self, betas, qcircuit, quantum_input):

        k = 0
        linear_angles = []
        for angles in betas:
            linear_angles = linear_angles + angles
            for angle in angles:
                qcircuit.ry(angle, quantum_input[k])
                k += 1

        self.tree = BinTree(quantum_input)
        my_tree = self.tree

        last = my_tree.size - 1
        actual = my_tree.parent(last)
        level = my_tree.parent(last)
        while actual >= 0:
            left_index = my_tree.left(actual)
            right_index = my_tree.right(actual)
            while right_index <= last:

                qcircuit.cswap(my_tree[actual], my_tree[left_index], my_tree[right_index])

                left_index = my_tree.left(left_index)
                right_index = my_tree.left(right_index)
            actual -= 1
            if level != my_tree.parent(actual):
                level -= 1

        # set output qubits
        next_index = 0
        while next_index < my_tree.size:
            self.output_qubits.append(next_index)
            next_index = my_tree.left(next_index)

    def _generate_circuit(self, betas, qcircuit, quantum_input):
        numberof_controls = 0  # number of controls
        control_bits = []
        for angles in betas:
            if numberof_controls == 0:
                qcircuit.ry(angles[0], quantum_input[self.num_qubits-1])
                numberof_controls += 1
                control_bits.append(quantum_input[self.num_qubits-1])
            else:
                for k, angle in enumerate(reversed(angles)):
                    Encoding._index(k, qcircuit, control_bits, numberof_controls)

                    qcircuit.mcry(angle,
                                  control_bits,
                                  quantum_input[self.num_qubits - 1 - numberof_controls],
                                  None,
                                  mode='noancilla')

                    Encoding._index(k, qcircuit, control_bits, numberof_controls)
                control_bits.append(quantum_input[self.num_qubits - 1 - numberof_controls])
                numberof_controls += 1
