#   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.
"""Module to perform Trotter-Suzuki decompositions to output as circuits."""

from openfermion.ops import QubitOperator
import numpy
import copy
import collections

"""
Description:
    Functions for estimating the exponential of an operator
    composed of Pauli matrices, by Trotterization. Outputs QASM
    format to a python stream.

    Change-of-basis gates:
    H           : Z to X basis (and back)
    Rx(pi/2)    : Z to Y basis .....
    Rx(-pi/2)   : Y to Z basis .....
"""


def _third_order_trotter_helper(op_list):
    """Iteratively find 3rd-order Trotter ordering of a QubitOperator.

    This Trotter ordering is done according to the scheme:
    e^(A+B) = e^(7/24 * A) e^(2/3 * B) e^(3/4 * A) * e^(-2/3 * B)
        * e^(-1/24 * A) e^(B)

    Note:
        See N. Hatano and M. Suzuki Lect. Notes Phys 679, 37-68 (2005)

    Args:
        op_list (list of QubitOperator's): the terms in the
                hamiltonian as a list of QubitOperators
    Returns:
        list of QubitOperators giving the trotterized hamiltonian
    """
    # Create list to hold return value
    ret_val = collections.deque([7.0 / 24.0 * op_list[-2],
                                2.0 / 3.0 * op_list[-1],
                                3.0 / 4.0 * op_list[-2],
                                -2.0 / 3.0 * op_list[-1],
                                -1.0 / 24.0 * op_list[-2],
                                1.0 * op_list[-1]])

    for i in reversed(range(len(op_list)-2)):
        temp_ret_val = copy.deepcopy(ret_val)
        ret_val = collections.deque()
        ret_val.appendleft(7.0 / 24.0 * op_list[i])
        ret_val.extend([2.0 / 3.0 * qubop for qubop in temp_ret_val])
        ret_val.append(3.0 / 4.0 * op_list[i])
        ret_val.extend([-2.0 / 3.0 * qubop for qubop in temp_ret_val])
        ret_val.append(-1.0 / 24.0 * op_list[i])
        ret_val.extend([1.0 * qubop for qubop in temp_ret_val])
    return ret_val


def trotter_operator_grouping(hamiltonian,
                              trotter_number=1,
                              trotter_order=1,
                              term_ordering=None,
                              k_exp=1.0):
    """Trotter-decomposes operators into groups without exponentiating.

    Operators are still Hermitian at the end of this method but have been
        multiplied by k_exp.

    Note:
        The default term_ordering is simply the ordered keys of
        the QubitOperators.terms dict.

    Args:
        hamiltonian (QubitOperator): full hamiltonian
        trotter_number (int): optional number of trotter steps -
            default is 1
        trotter_order (int): optional order of trotterization as
            an integer from 1-3 - default is 1
        term_ordering (list of (tuples of tuples)): optional list
            of QubitOperator terms dictionary keys that specifies
            order of terms when trotterizing
        k_exp (float): optional exponential factor
            to all terms when trotterizing

    Yields:
        QubitOperator generator

    Raises:
        ValueError if order > 3 or order <= 0,
        TypeError for incorrect types
    """
    # Check for default arguments and type errors
    if (trotter_order > 3) or (trotter_order <= 0):
        raise ValueError("Invalid trotter order: " + str(trotter_order))
    if not isinstance(hamiltonian, QubitOperator):
        raise TypeError("Hamiltonian must be a QubitOperator.")
    if len(hamiltonian.terms) == 0:
        raise TypeError("Hamiltonian must be a non-empty QubitOperator.")
    if term_ordering is None:
        # To have consistent default behavior, ordering = sorted keys.
        term_ordering = sorted(list(hamiltonian.terms.keys()))
    if len(term_ordering) == 0:
        raise TypeError("term_ordering must None or non-empty list.")

    # Enforce float
    k_exp = float(k_exp)

    ret_val = []
    # First order trotter
    if trotter_order == 1:
        for step in range(trotter_number):
            for op in term_ordering:
                yield QubitOperator(op,
                                    hamiltonian.terms[op]
                                    * float(k_exp)
                                    / float(trotter_number)
                                    )
    # Second order trotter
    elif trotter_order == 2:
        if len(term_ordering) < 2:
            raise ValueError("Not enough terms in the Hamiltonian to do " +
                             "second order trotterization")
        for step in range(trotter_number):
            for op in term_ordering[:-1]:
                yield QubitOperator(op,
                                    (hamiltonian.terms[op]
                                     * k_exp
                                     / (2.0 * float(trotter_number))))

            yield QubitOperator(term_ordering[-1],
                                hamiltonian.terms[term_ordering[-1]]
                                * float(k_exp)
                                / float(trotter_number))

            for op in reversed(term_ordering[:-1]):
                yield QubitOperator(op, hamiltonian.terms[op]
                                    * float(k_exp)
                                    / (2.0 * float(trotter_number)))

    # Third order trotter
    elif trotter_order == 3:
        if len(term_ordering) < 2:
            raise ValueError("Not enough terms in the Hamiltonian to do " +
                             "third order trotterization")
        # Make sure original hamiltonian is not modified
        ham = hamiltonian / float(trotter_number)
        ham *= float(k_exp)
        ham_temp = []
        for term in term_ordering:
            ham_temp.append(QubitOperator(term, ham.terms[term]))
        for step in range(trotter_number):
            for returned_op in _third_order_trotter_helper(ham_temp):
                yield returned_op


def pauli_exp_to_qasm(qubit_operator_list,
                      evolution_time=1.0):
    """Exponentiate a list of QubitOperators to a QASM string generator.

    Exponentiates a list of QubitOperators, and yields string generators in
        QASM format using the formula:  exp(-1.0j * evolution_time * op).

    Args:
        qubit_operator_list (list of QubitOperators): list of single Pauli-term
            QubitOperators to be exponentiated
        evolution_time (float): evolution time of the operators in
            the list

    Yields:
        string
    """
    for qubit_operator in qubit_operator_list:
        # ret_val = ""
        ret_list = []

        for term in qubit_operator.terms:
            term_coeff = qubit_operator.terms[term]

            # List of operators and list of qubit ids
            ops = []
            qids = []
            string_basis_1 = []  # Basis rotations 1
            string_basis_2 = []  # Basis rotations 2

            for p in term:  # p = single pauli term
                qid = p[0]  # Qubit index
                pop = p[1]  # Pauli op

                qids.append(qid)  # Qubit index
                ops.append(pop)  # Pauli operator

                if pop == 'X':
                    string_basis_1.append("H {}".format(qid))    # Hadamard
                    string_basis_2.append("H {}".format(qid))    # Hadamard
                elif pop == 'Y':
                    string_basis_1.append(
                        "Rx 1.5707963267948966 {}".format(qid))
                    string_basis_2.append(
                        "Rx -1.5707963267948966 {}".format(qid))

            # Prep for CNOTs
            cnot_pairs = numpy.vstack((qids[:-1], qids[1:]))
            cnots1 = []
            cnots2 = []
            for i in range(cnot_pairs.shape[1]):
                pair = cnot_pairs[:, i]
                cnots1.append("CNOT {} {}".format(pair[0], pair[1]))
            for i in numpy.arange(cnot_pairs.shape[1])[::-1]:
                pair = cnot_pairs[:, i]
                cnots2.append("CNOT {} {}".format(pair[0], pair[1]))

            # Exponentiating each Pauli string requires five parts

            # 1. Perform basis rotations
            ret_list = ret_list + string_basis_1

            # 2. First set CNOTs
            ret_list = ret_list + cnots1

            # 3. Rotation (Note kexp & Ntrot)
            ret_list = ret_list + ["Rz {} {}".format(
                        term_coeff * evolution_time, qids[-1])]

            # 4. Second set of CNOTs
            ret_list = ret_list + cnots2

            # 5. Rotate back to Z basis
            ret_list = ret_list + string_basis_2

            for gate in ret_list:
                yield gate


def trotterize_exp_qubop_to_qasm(hamiltonian,
                                 trotter_number=1,
                                 trotter_order=1,
                                 term_ordering=None,
                                 k_exp=1.0):
    """Trotterize a Qubit hamiltonian and write it to QASM format.

    Assumes input hamiltonian is still hermitian and -1.0j has not yet been
    applied. Therefore, signs of coefficients should reflect this. Returns
    a generator which generates a QASM file.

    Args:
        hamiltonian (QubitOperator): hamiltonian
        trotter_number (int): optional number of trotter steps (slices) for
            trotterization as an integer - default = 1
        trotter_order: optional order of trotterization as an integer -
            default = 1
        term_ordering (list of (tuples of tuples)): list of tuples
            (QubitOperator terms dictionary keys) that specifies
            order of terms when trotterizing
        k_exp (float): optional exponential factor to all
            terms when trotterizing

        Yields:
            string generator

    """

    for trotterized_op in trotter_operator_grouping(hamiltonian,
                                                    trotter_number,
                                                    trotter_order,
                                                    term_ordering,
                                                    k_exp):
        for exponentiated_qasm_string in pauli_exp_to_qasm([trotterized_op]):
            yield exponentiated_qasm_string
