################################################################################
# © Copyright 2022 Zapata Computing Inc.
################################################################################
import itertools
import random
from typing import Iterable, List, Optional

import numpy as np

from orquestra.quantum.circuits import Circuit, X, Y, Z
from orquestra.quantum.measurements import ExpectationValues
from orquestra.quantum.openfermion import (
    FermionOperator,
    InteractionOperator,
    InteractionRDM,
    PolynomialTensor,
    QubitOperator,
    count_qubits,
)
from orquestra.quantum.openfermion import expectation as openfermion_expectation
from orquestra.quantum.openfermion import (
    get_interaction_operator,
    get_sparse_operator,
    normal_ordered,
)
from orquestra.quantum.openfermion.hamiltonians.special_operators import number_operator
from orquestra.quantum.openfermion.linalg import jw_get_ground_state_at_particle_number
from orquestra.quantum.openfermion.transforms import (
    freeze_orbitals,
    get_fermion_operator,
)
from orquestra.quantum.utils import ValueEstimate, bin2dec, dec2bin
from orquestra.quantum.wavefunction import Wavefunction


def get_qubitop_from_matrix(operator: List[List]) -> QubitOperator:
    r"""Expands a 2^n by 2^n matrix into n-qubit Pauli basis. The runtime of
    this function is O(2^2n).

    Args:
        operator: a list of lists (rows) representing a 2^n by 2^n
            matrix.

    Returns:
        A QubitOperator instance corresponding to the expansion of
        the input operator as a sum of Pauli strings:

        O = 2^-n \sum_P tr(O*P) P
    """

    nrows = len(operator)
    ncols = len(operator[0])

    # Check if the input operator is square
    if nrows != ncols:
        raise Exception("The input operator is not square")

    # Check if the dimensions are powers of 2
    if not (((nrows & (nrows - 1)) == 0) and nrows > 0):
        raise Exception("The number of rows is not a power of 2")
    if not (((ncols & (ncols - 1)) == 0) and ncols > 0):
        raise Exception("The number of cols is not a power of 2")

    n = int(np.log2(nrows))  # number of qubits

    def decode(bit_string):  # Helper function for converting any 2n-bit
        # string to a label vector representing a Pauli
        # string of length n

        if len(bit_string) != 2 * n:
            raise Exception("LH_expand:decode: input bit string length not 2n")

        output_label = list(np.zeros(n))
        for i in range(0, n):
            output_label[i] = bin2dec(bit_string[2 * i : 2 * i + 2])

        return output_label

    def trace_product(label_vec):  # Helper function for computing tr(OP)
        # where O is the input operator and P is a
        # Pauli string operator

        def f(j):  # Function which computes the index of the nonzero
            # element in P for a given column j

            j_str = dec2bin(j, n)
            for index in range(0, n):
                if label_vec[index] in [1, 2]:  # flip if X or Y
                    j_str[index] = int(not j_str[index])
            return bin2dec(j_str)

        def nz(j):  # Function which computes the value of the nonzero
            # element in P on the column j

            val_nz = 1.0
            j_str = dec2bin(j, n)
            for index in range(0, n):
                if label_vec[index] == 2:
                    if j_str[index] == 0:
                        val_nz = val_nz * (1j)
                    if j_str[index] == 1:
                        val_nz = val_nz * (-1j)
                if label_vec[index] == 3:
                    if j_str[index] == 1:
                        val_nz = val_nz * (-1)
            return val_nz

        # Compute the trace
        tr = 0.0
        for j in range(0, 2 ** n):  # loop over the columns
            tr = tr + operator[j][f(j)] * nz(j)

        return tr / 2 ** n

    # Expand the operator in Pauli basis
    coeffs = list(np.zeros(4 ** n))
    labels = list(np.zeros(4 ** n))
    for i in range(0, 4 ** n):  # loop over all 2n-bit strings
        current_string = dec2bin(i, 2 * n)  # see util.py
        current_label = decode(current_string)
        coeffs[i] = trace_product(current_label)
        labels[i] = current_label

    return get_qubitop_from_coeffs_and_labels(coeffs, labels)


def get_qubitop_from_coeffs_and_labels(
    coeffs: List[float], labels: List[List[int]]
) -> QubitOperator:
    """Generates a QubitOperator based on a coefficient vector and
    a label matrix.

    Args:
        coeffs: a list of floats representing the coefficients
            for the terms in the Hamiltonian
        labels: a list of lists (a matrix) where each list
            is a vector of integers representing the Pauli
            string. See pauliutil.py for details.

    Example:

        The Hamiltonian H = 0.1 X1 X2 - 0.4 Y1 Y2 Z3 Z4 can be
        initiated by calling

        H = QubitOperator([0.1, -0.4],  # coefficients
            [[1 1 0 0],  # label matrix
            [2 2 3 3]])
    """

    output = QubitOperator()
    for i in range(0, len(labels)):
        string_term = ""
        for ind, elem in enumerate(labels[i]):
            pauli_symbol = ""
            if elem == 1:
                pauli_symbol = "X" + str(ind) + " "
            if elem == 2:
                pauli_symbol = "Y" + str(ind) + " "
            if elem == 3:
                pauli_symbol = "Z" + str(ind) + " "
            string_term += pauli_symbol

        output += coeffs[i] * QubitOperator(string_term)

    return output


def generate_random_qubitop(
    nqubits: int,
    nterms: int,
    nlocality: int,
    max_coeff: float,
    fixed_coeff: bool = False,
) -> QubitOperator:
    """Generates a Hamiltonian with term coefficients uniformly distributed
    in [-max_coeff, max_coeff].

    Args:
        nqubits - number of qubits
        nterms    - number of terms in the Hamiltonian
        nlocality - locality of the Hamiltonian
        max_coeff - bound for generating the term coefficients
        fixed_coeff (bool) - If true, all the terms are assign the
            max_coeff as coefficient.

    Returns:
        A QubitOperator with the appropriate coefficient vector
        and label matrix.
    """
    # generate random coefficient vector
    if fixed_coeff:
        coeffs = [max_coeff] * nterms
    else:
        coeffs = list(np.zeros(nterms))
        for j in range(0, nterms):
            coeffs[j] = random.uniform(-max_coeff, max_coeff)

    # generate random label vector
    labels = list(np.zeros(nterms, dtype=int))
    label_set = set()
    j = 0
    while j < nterms:
        inds_nontrivial = sorted(random.sample(range(0, nqubits), nlocality))
        label = list(np.zeros(nqubits, dtype=int))
        for ind in inds_nontrivial:
            label[ind] = random.randint(1, 3)
        if str(label) not in label_set:
            labels[j] = label
            j += 1
            label_set.add(str(label))
    return get_qubitop_from_coeffs_and_labels(coeffs, labels)


def evaluate_qubit_operator(
    qubit_operator: QubitOperator, expectation_values: ExpectationValues
) -> ValueEstimate:
    """Evaluate the expectation value of a qubit operator using expectation values for
    the terms.

    Args:
        qubit_operator: the operator
        expectation_values: the expectation values

    Returns:
        value_estimate: stores the value of the expectation and its precision
    """

    # Sum the contributions from all terms
    total = 0

    # Add all non-trivial terms
    term_index = 0
    for term in qubit_operator.terms:
        total += np.real(
            qubit_operator.terms[term] * expectation_values.values[term_index]
        )
        term_index += 1

    value_estimate = ValueEstimate(total)
    return value_estimate


def evaluate_qubit_operator_list(
    qubit_operator_list: List[QubitOperator], expectation_values: ExpectationValues
) -> ValueEstimate:
    """Evaluate the expectation value of a qubit operator list using expectation values
    for the terms. The expectation values should be in the order given by the qubit
    operator list, and the value returned is the sum of all terms in the qubit operator
    list.

    Args:
        qubit_operator_list: the operator list
        expectation_values: the expectation values

    Returns:
        value_estimate: stores the value of the expectation and its precision
    """

    # Sum the contributions from all terms
    total = 0

    # Add all non-trivial terms
    term_index = 0
    for qubit_operator in qubit_operator_list:
        for term in qubit_operator.terms:
            total += np.real(
                qubit_operator.terms[term] * expectation_values.values[term_index]
            )
            term_index += 1

    value_estimate = ValueEstimate(total)
    return value_estimate


def reverse_qubit_order(qubit_operator: QubitOperator, n_qubits: Optional[int] = None):
    """Reverse the order of qubit indices in a qubit operator.

    Args:
        qubit_operator (openfermion.QubitOperator): the operator
        n_qubits (int): total number of qubits. Needs to be provided when
            the size of the system of interest is greater than the size of qubit
            operator (optional)

    Returns:
        reversed_op (openfermion.ops.QubitOperator)
    """

    reversed_op = QubitOperator()

    if n_qubits is None:
        n_qubits = count_qubits(qubit_operator)
    if n_qubits < count_qubits(qubit_operator):
        raise ValueError("Invalid number of qubits specified.")

    for term in qubit_operator.terms:
        new_term = []
        for factor in term:
            new_factor = list(factor)
            new_factor[0] = n_qubits - 1 - new_factor[0]
            new_term.append(tuple(new_factor))
        reversed_op += QubitOperator(tuple(new_term), qubit_operator.terms[term])
    return reversed_op


def get_expectation_value(
    qubit_op: QubitOperator, wavefunction: Wavefunction, reverse_operator: bool = False
) -> complex:
    """Get the expectation value of a qubit operator with respect to a wavefunction.
    Args:
        qubit_op: the operator
        wavefunction: the wavefunction
        reverse_operator: whether to reverse order of qubit operator
            before computing expectation value. This should be True if the convention
            of the basis states used for the wavefunction is the opposite of the one in
            the qubit operator. This is the case when the wavefunction uses
            Rigetti convention (https://arxiv.org/abs/1711.02086) of ordering qubits.
    Returns:
        the expectation value
    """
    n_qubits = wavefunction.amplitudes.shape[0].bit_length() - 1

    # Convert the qubit operator to a sparse matrix. Note that the qubit indices
    # must be reversed because OpenFermion and our Wavefunction use
    # different conventions for how to order the computational basis states!
    if reverse_operator:
        qubit_op = reverse_qubit_order(qubit_op, n_qubits=n_qubits)
    sparse_op = get_sparse_operator(qubit_op, n_qubits=n_qubits)

    # Computer the expectation value
    exp_val = openfermion_expectation(sparse_op, wavefunction.amplitudes)
    return exp_val


def change_operator_type(operator, operatorType):
    """Take an operator and attempt to cast it to an operator of a different type

    Args:
        operator: The operator
        operatorType: The type of the operator that the original operator is
            cast to
    Returns:
        An operator with type operatorType
    """
    new_operator = operatorType()
    for op in operator.terms:
        new_operator += operatorType(tuple(op), operator.terms[op])

    return new_operator


def get_fermion_number_operator(n_qubits, n_particles=None):
    """Return a FermionOperator representing the number operator for n qubits.
    If `n_particles` is specified, it can be used for creating constraint on the number
    of particles.

    Args:
        n_qubits (int): number of qubits in the system
        n_particles (int): number of particles in the system.
            If specified, it is subtracted from the number
            operator such as expectation value is zero.
    Returns:
         (openfermion.ops.FermionOperator): the number operator
    """
    operator = number_operator(n_qubits)
    if n_particles is not None:
        operator += FermionOperator("", -1.0 * float(n_particles))
    return get_interaction_operator(operator)


def get_diagonal_component(operator):
    if isinstance(operator, InteractionOperator):
        return _get_diagonal_component_interaction_operator(operator)
    elif isinstance(operator, PolynomialTensor):
        return _get_diagonal_component_polynomial_tensor(operator)
    else:
        raise TypeError(
            f"Getting diagonal component not supported for {0}".format(type(operator))
        )


def _get_diagonal_component_polynomial_tensor(polynomial_tensor):
    """Get the component of an interaction operator that is
    diagonal in the computational basis under Jordan-Wigner
    transformation (i.e., the terms that can be expressed
    as products of number operators).
    Args:
        interaction_operator (openfermion.ops.InteractionOperator): the operator

    Returns:
        tuple: two openfermion.ops.InteractionOperator objects. The first is the
            diagonal component, and the second is the remainder.
    """
    n_modes = count_qubits(polynomial_tensor)
    remainder_tensors = {}
    diagonal_tensors = {}

    diagonal_tensors[()] = polynomial_tensor.constant
    for key in polynomial_tensor.n_body_tensors:
        if key == ():
            continue
        remainder_tensors[key] = np.zeros((n_modes,) * len(key), complex)
        diagonal_tensors[key] = np.zeros((n_modes,) * len(key), complex)

        for indices in itertools.product(range(n_modes), repeat=len(key)):
            creation_counts = {}
            annihilation_counts = {}

            for meta_index, index in enumerate(indices):
                if key[meta_index] == 0:
                    if annihilation_counts.get(index) is None:
                        annihilation_counts[index] = 1
                    else:
                        annihilation_counts[index] += 1
                elif key[meta_index] == 1:
                    if creation_counts.get(index) is None:
                        creation_counts[index] = 1
                    else:
                        creation_counts[index] += 1

            term_is_diagonal = True
            for index in creation_counts:
                if creation_counts[index] != annihilation_counts.get(index):
                    term_is_diagonal = False
                    break
            if term_is_diagonal:
                for index in annihilation_counts:
                    if annihilation_counts[index] != creation_counts.get(index):
                        term_is_diagonal = False
                        break
            if term_is_diagonal:
                diagonal_tensors[key][indices] = polynomial_tensor.n_body_tensors[key][
                    indices
                ]
            else:
                remainder_tensors[key][indices] = polynomial_tensor.n_body_tensors[key][
                    indices
                ]

    return PolynomialTensor(diagonal_tensors), PolynomialTensor(remainder_tensors)


def _get_diagonal_component_interaction_operator(interaction_operator):
    """Get the component of an interaction operator that is
    diagonal in the computational basis under Jordan-Wigner
    transformation (i.e., the terms that can be expressed
    as products of number operators).
    Args:
        interaction_operator (openfermion.ops.InteractionOperator): the operator

    Returns:
        tuple: two openfermion.ops.InteractionOperator objects. The first is the
            diagonal component, and the second is the remainder.
    """

    one_body_tensor = np.zeros(
        interaction_operator.one_body_tensor.shape, dtype=complex
    )
    two_body_tensor = np.zeros(
        interaction_operator.two_body_tensor.shape, dtype=complex
    )
    diagonal_op = InteractionOperator(
        interaction_operator.constant, one_body_tensor, two_body_tensor
    )

    one_body_tensor = np.copy(interaction_operator.one_body_tensor).astype(complex)
    two_body_tensor = np.copy(interaction_operator.two_body_tensor).astype(complex)
    remainder_op = InteractionOperator(0.0, one_body_tensor, two_body_tensor)

    n_spin_orbitals = interaction_operator.two_body_tensor.shape[0]

    for p in range(n_spin_orbitals):
        for q in range(n_spin_orbitals):
            diagonal_op.two_body_tensor[
                p, q, p, q
            ] = interaction_operator.two_body_tensor[p, q, p, q]
            diagonal_op.two_body_tensor[
                p, q, q, p
            ] = interaction_operator.two_body_tensor[p, q, q, p]
            remainder_op.two_body_tensor[p, q, p, q] = 0.0
            remainder_op.two_body_tensor[p, q, q, p] = 0.0

    for p in range(n_spin_orbitals):
        diagonal_op.one_body_tensor[p, p] = interaction_operator.one_body_tensor[p, p]
        remainder_op.one_body_tensor[p, p] = 0.0

    return diagonal_op, remainder_op


def get_polynomial_tensor(fermion_operator, n_qubits=None):
    r"""Convert a fermionic operator to a Polynomial Tensor.

    Args:
        fermion_operator (openferion.ops.FermionOperator): The operator.
        n_qubits (int): The number of qubits to be included in the
            PolynomialTensor. Must be at least equal to the number of qubits
            that are acted on by fermion_operator. If None, then the number of
            qubits is inferred from fermion_operator.

    Returns:
        openfermion.ops.PolynomialTensor: The tensor representation of the
            operator.
    """
    if not isinstance(fermion_operator, FermionOperator):
        raise TypeError("Input must be a FermionOperator.")

    if n_qubits is None:
        n_qubits = count_qubits(fermion_operator)
    if n_qubits < count_qubits(fermion_operator):
        raise ValueError("Invalid number of qubits specified.")

    # Normal order the terms and initialize.
    fermion_operator = normal_ordered(fermion_operator)
    tensor_dict = {}

    # Loop through terms and assign to matrix.
    for term in fermion_operator.terms:
        coefficient = fermion_operator.terms[term]

        # Handle constant shift.
        if len(term) == 0:
            tensor_dict[()] = coefficient

        else:
            key = tuple([operator[1] for operator in term])
            if tensor_dict.get(key) is None:
                tensor_dict[key] = np.zeros((n_qubits,) * len(key), complex)

            indices = tuple([operator[0] for operator in term])
            tensor_dict[key][indices] = coefficient

    return PolynomialTensor(tensor_dict)


def create_circuits_from_qubit_operator(qubit_operator: QubitOperator) -> List[Circuit]:
    """Creates a list of circuit objects from the Pauli terms of a QubitOperator
    Args:
        qubit_operator: operator for which the Pauli terms are converted into circuits

    Return:
        circuit_set: a list of Pauli string gate circuits
    """

    # Get the Pauli terms, ignoring coefficients
    pauli_terms = list(qubit_operator.terms.keys())
    term_gate_map = {"X": X, "Y": Y, "Z": Z}
    circuit_set = []

    # Loop over Pauli terms and populate circuit set list
    for term in pauli_terms:

        circuit = Circuit()

        # Loop over Pauli factors in Pauli term and construct Pauli term circuit
        for pauli in term:  # loop over pauli operators in an n qubit pauli term
            pauli_index = pauli[0]
            pauli_factor = pauli[1]
            circuit += term_gate_map[pauli_factor](pauli_index)

        circuit_set += [circuit]

    return circuit_set


def get_ground_state_rdm_from_qubit_op(
    qubit_operator: QubitOperator, n_particles: int
) -> InteractionRDM:
    """Diagonalize operator and compute the ground state 1- and 2-RDM

    Args:
        qubit_operator: The openfermion operator to diagonalize
        n_particles: number of particles in the target ground state

    Returns:
        rdm: interaction RDM of the ground state with the particle number n_particles
    """

    sparse_operator = get_sparse_operator(qubit_operator)
    e, ground_state_wf = jw_get_ground_state_at_particle_number(
        sparse_operator, n_particles
    )  # float/np.array pair
    n_qubits = count_qubits(qubit_operator)

    one_body_tensor_list = []
    for i in range(n_qubits):
        for j in range(n_qubits):
            idag_j = get_sparse_operator(
                FermionOperator(f"{i}^ {j}"), n_qubits=n_qubits
            )
            idag_j = idag_j.toarray()
            one_body_tensor_list.append(
                np.conjugate(ground_state_wf) @ idag_j @ ground_state_wf
            )

    one_body_tensor = np.array(one_body_tensor_list)
    one_body_tensor = one_body_tensor.reshape(n_qubits, n_qubits)

    two_body_tensor = np.zeros((n_qubits,) * 4, dtype=complex)
    for p in range(n_qubits):
        for q in range(0, p + 1):
            for r in range(n_qubits):
                for s in range(0, r + 1):
                    pdag_qdag_r_s = get_sparse_operator(
                        FermionOperator(f"{p}^ {q}^ {r} {s}"), n_qubits=n_qubits
                    )
                    pdag_qdag_r_s = pdag_qdag_r_s.toarray()
                    rdm_element = (
                        np.conjugate(ground_state_wf) @ pdag_qdag_r_s @ ground_state_wf
                    )
                    two_body_tensor[p, q, r, s] = rdm_element
                    two_body_tensor[q, p, r, s] = -rdm_element
                    two_body_tensor[q, p, s, r] = rdm_element
                    two_body_tensor[p, q, s, r] = -rdm_element

    return InteractionRDM(one_body_tensor, two_body_tensor)


def remove_inactive_orbitals(
    interaction_op: InteractionOperator, n_active: int = None, n_core: int = 0
) -> InteractionOperator:
    """Remove orbitals not in the active space from an interaction operator.

    Args:
        interaction_op: the operator, assumed to be ordered with alternating spin-up and
            spin-down spin orbitals.
        n_active: the number of active molecular orbitals. If None, include all orbitals
            beyond n_core. Note that the number of active spin orbitals will be twice
            the number of active molecular orbitals.
        n_core: the number of core molecular orbitals to be frozen.

    Returns:
        The interaction operator with inactive orbitals removed, and the Hartree-Fock
            energy of the core orbitals added to the constant.
    """

    # This implementation is probably not very efficient, because it converts the
    # interaction operator into a fermion operator and then back to an interaction
    # operator.

    # Convert the InteractionOperator to a FermionOperator
    fermion_op = get_fermion_operator(interaction_op)

    # Determine which occupied spin-orbitals are to be frozen
    occupied = range(2 * n_core)

    unoccupied: Iterable
    # Determine which unoccupied spin-orbitals are to be frozen
    if n_active is not None:
        unoccupied = range(
            2 * n_core + 2 * n_active, interaction_op.one_body_tensor.shape[0]
        )
    else:
        unoccupied = []

    # Freeze the spin-orbitals
    frozen_fermion_op = freeze_orbitals(fermion_op, occupied, unoccupied)

    # Convert back to an interaction operator
    frozen_interaction_op = get_interaction_operator(frozen_fermion_op)

    return frozen_interaction_op


def hf_rdm(n_alpha: int, n_beta: int, n_orbitals: int) -> InteractionRDM:
    """Construct the RDM corresponding to a Hartree-Fock state.

    Args:
        n_alpha (int): number of spin-up electrons
        n_beta (int): number of spin-down electrons
        n_orbitals (int): number of spatial orbitals (not spin orbitals)

    Returns:
        openfermion.ops.InteractionRDM: the reduced density matrix
    """
    # Determine occupancy of each spin orbital
    occ = np.zeros(2 * n_orbitals)
    occ[: (2 * n_alpha) : 2] = 1
    occ[1 : (2 * n_beta + 1) : 2] = 1

    one_body_tensor = np.diag(occ)

    two_body_tensor = np.zeros([2 * n_orbitals for i in range(4)])
    for i in range(2 * n_orbitals):
        for j in range(2 * n_orbitals):
            if i != j and occ[i] and occ[j]:
                two_body_tensor[i, j, j, i] = 1
                two_body_tensor[i, j, i, j] = -1

    return InteractionRDM(one_body_tensor, two_body_tensor)
