# Copyright 2023 Good Chemistry Company.
#
# 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.

"""This module provides a common point for mapping operators and
 statevectors from Fermionic to Qubit encodings via any of:
 - Jordan-Wigner
 - Bravyi-Kitaev (Fenwick Tree implementation)
 - symmetry-conserving Bravyi-Kitaev (2-qubit reduction via Z2 taper)
"""

from math import ceil
import numpy as np
from collections.abc import Iterable

from openfermion import FermionOperator as ofFermionOperator

from tangelo.toolboxes.operators import FermionOperator, QubitOperator
from tangelo.toolboxes.qubit_mappings import jordan_wigner, bravyi_kitaev, symmetry_conserving_bravyi_kitaev, jkmn
from tangelo.toolboxes.qubit_mappings.hcb import hard_core_boson_operator, boson_to_qubit_mapping

available_mappings = {"JW", "BK", "SCBK", "JKMN", "HCB"}


def get_qubit_number(mapping, n_spinorbitals):
    """Get the number of qubits for a specified number of spin orbitals, after
    passing through the qubit mapping. Symmetry-conserving Bravyi-Kitaev reduces
    the qubit number by 2.

    Args:
        mapping (string): qubit-mapping, JW, BK or SCBK.
        n_spinorbitals (int): number of spin-orbitals for molecule.

    Returns:
        int: number of qubits.
    """
    if mapping.upper() == "SCBK":
        return n_spinorbitals - 2
    elif mapping.upper() == "HCB":
        return ceil(n_spinorbitals / 2)
    else:
        return n_spinorbitals


def get_fermion_operator(operator):
    """Cast operator to FermionOperator datatype. Input is of SymbolicOperator
    type, but term words must be valid input for FermionOperator, as for example
    for InteractionOperator.

    Args:
        operator (SymbolicOperator): input operator to be cast.

    Returns:
        FermionOperator: Self-explanatory.
    """
    if not isinstance(operator, Iterable):
        raise TypeError("Input must be iterable suitable for casting to FermionOperator type.")
    fermion_operator = FermionOperator()
    for term in operator:
        try:
            fermion_operator += FermionOperator(term, operator[term])
        except Exception:
            raise TypeError("Operator terms are not formatted as valid input for FermionOperator type.")
    return fermion_operator


def fermion_to_qubit_mapping(fermion_operator, mapping, n_spinorbitals=None, n_electrons=None, up_then_down=False, spin=0):
    """Perform mapping of fermionic operator to qubit operator. This function is
    mostly a wrapper around standard openfermion code, with some important
    distinctions. We strictly enforce the specification of n_qubits for
    Bravyi-Kitaev type transformations, and n_electrons for scBK. In the absence
    of this information, these methods can return unexpected results if the
    input operator does not specifically address the highest-orbital/qubit
    index.

    Args:
        fermion_operator (FermionOperator): operator to translate to qubit
            representation.
        mapping (string): options are -- "JW" (Jordan Wigner), "BK"
            (Bravyi Kitaev), "scBK" (symmetry-conserving Bravyi Kitaev).
        n_spinorbitals (int): number of spin-orbitals for problem. Not required
            for Jordan-Wigner.
        n_electrons (int): number of occupied electron modes in problem.
            Required for symmetry conserving Bravyi-Kitaev only.
        up_then_down (bool): flag to change basis ordering, putting all spin up
            then all spin down.
        spin (int): The spin number (n_alpha - n_beta)

    Returns:
        QubitOperator: input operator, encoded in the qubit space.
    """
    # some methods may pass another operator class type. If this is the case, cast to FermionOperator where possible
    if not isinstance(fermion_operator, ofFermionOperator):
        fermion_operator = get_fermion_operator(fermion_operator)

    if mapping.upper() not in available_mappings:
        raise ValueError(f"Invalid mapping selection. Select from: {available_mappings}")

    if mapping.upper in {"BK", "SCBK", "JKMN"} and n_spinorbitals is None:
        raise ValueError(f"{mapping.upper()} requires n_spinorbitals to be set.")

    if up_then_down:
        if n_spinorbitals is None:
            raise ValueError("The number of spin-orbitals is required to execute basis re-ordering.")
        fermion_operator = make_up_then_down(fermion_operator, n_spinorbitals=n_spinorbitals)

    if mapping.upper() == "JW":
        qubit_operator = jordan_wigner(fermion_operator)
    elif mapping.upper() == "BK":
        qubit_operator = bravyi_kitaev(fermion_operator, n_qubits=n_spinorbitals)
    elif mapping.upper() == "SCBK":
        if n_electrons is None:
            raise ValueError("Symmetry-conserving Bravyi Kitaev requires specification of number of electrons.")

        qubit_operator = symmetry_conserving_bravyi_kitaev(fermion_operator=fermion_operator,
                                                           n_spinorbitals=n_spinorbitals,
                                                           n_electrons=n_electrons,
                                                           up_then_down=up_then_down,
                                                           spin=spin)
    elif mapping.upper() == "JKMN":
        qubit_operator = jkmn(fermion_operator, n_qubits=n_spinorbitals)
    elif mapping.upper() == "HCB":
        boson_operator = hard_core_boson_operator(fermion_operator)
        qubit_operator = boson_to_qubit_mapping(boson_operator)

    converted_qubit_op = QubitOperator()
    converted_qubit_op.terms = qubit_operator.terms.copy()

    return converted_qubit_op


def make_up_then_down(fermion_operator, n_spinorbitals):
    """Re-order the orbital indexing to force all spin up, followed by all spin
        down. The default ordering alternates between spin up and down.

    Args:
        fermion_operator (FermionOperator): input operator.
        n_spinorbitals (int): number of spin-orbitals in register.

    Returns:
        FermionOperator: operator with all spin up followed by all spin down.
    """
    if not isinstance(fermion_operator, ofFermionOperator):
        raise TypeError("Invalid operator input. Must be FermionOperator.")
    if n_spinorbitals % 2 != 0:
        raise ValueError("Invalid number of spin-orbitals. Expecting even number.")
    term_modes = max([factor[0] for term in fermion_operator.terms for factor in term]) + 1
    if term_modes > n_spinorbitals:
        raise ValueError("Invalid number of modes (n_spinorbitals) for input operator. Terms in operator exceed this value.")
    new_operator = FermionOperator()

    # define the new mode register
    remapped = np.linspace(0, n_spinorbitals - 1, n_spinorbitals, dtype=int)//2
    remapped[1::2] += int(np.ceil(n_spinorbitals / 2.))

    for term, coef in fermion_operator.terms.items():

        new_term = tuple([(int(remapped[ti[0]]), ti[1]) for ti in term])
        new_operator += FermionOperator(new_term, coef)

    return new_operator
