#  Copyright 2021 The QUARK Authors. All Rights Reserved.
#
#  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.

import logging
from typing import TypedDict

import numpy as np
from more_itertools import locate
from qiskit_optimization import QuadraticProgram
from qiskit_optimization.converters import QuadraticProgramToQubo

from quark.modules.applications.Mapping import Mapping
from quark.modules.Core import Core
from quark.utils import start_time_measurement, end_time_measurement


class Ising(Mapping):
    """
    Ising formulation of the auto-carrier loading (ACL) problem.
    """

    def __init__(self):
        """
        Constructor method.
        """
        super().__init__()
        self.submodule_options = ["QAOA", "QiskitQAOA"]
        self.global_variables = []
        logging.warning("Currently, all scenarios are too large to be solved with an Ising model.")
        logging.warning("Consider using another mapping until the modelling is refined.")

    @staticmethod
    def get_requirements() -> list[dict]:
        """
        Return requirements of this module.

        :return: list of dict with requirements of this module
        """
        return [
            {"name": "numpy", "version": "1.26.4"},
            {"name": "more-itertools", "version": "10.5.0"},
            {"name": "qiskit-optimization", "version": "0.6.1"},
        ]

    def get_parameter_options(self) -> dict:
        """
        Returns empty dict as this mapping has no configurable settings.

        :return: Empty dict
        """
        return {}

    class Config(TypedDict):
        """
        Empty config as this solver has no configurable settings.
        """
        pass

    def map_pulp_to_qiskit(self, problem: dict) -> QuadraticProgram:
        """
        Maps the problem dict to a quadratic program.

        :param problem: Problem formulation in dict form
        :return: Quadratic program in qiskit-optimization format
        """
        # Details at:
        # https://coin-or.github.io/pulp/guides/how_to_export_models.html
        # https://qiskit.org/documentation/stable/0.26/tutorials/optimization/2_converters_for_quadratic_programs.html

        qp = QuadraticProgram()

        # Variables
        for variable_dict in problem["variables"]:
            if variable_dict["cat"] == "Integer":
                lb = variable_dict["lowBound"]
                ub = variable_dict["upBound"]
                name = variable_dict["name"]
                # If the integer variable is actually a binary variable
                if lb == 0 and ub == 1:
                    qp.binary_var(name)
                # If the integer variable is non-binary
                else:
                    qp.integer_var(lowerbound=lb, upperbound=ub, name=name)

        # Objective function
        obj_arguments = {arg["name"]: arg["value"] for arg in problem["objective"]["coefficients"]}
        # Maximize
        if problem["parameters"]["sense"] == -1:
            qp.maximize(linear=obj_arguments)
        # Minimize
        else:
            qp.minimize(linear=obj_arguments)

        # Constraints
        for constraint in problem["constraints"]:
            const_arguments = {arg["name"]: arg["value"] for arg in constraint["coefficients"]}
            sense = constraint["sense"]
            const_sense = "LE" if sense == -1 else "GE" if sense == 1 else "E"
            qp.linear_constraint(
                linear=const_arguments,
                sense=const_sense,
                rhs=-1 * constraint["constant"],
                name=constraint["name"]
            )

        return qp

    def map(self, problem: dict, config: Config) -> tuple[dict, float]:
        """
        Use Ising mapping of qiskit-optimize.

        :param problem: Dict containing the problem parameters
        :param config: Config with the parameters specified in Config class
        :return: Dict with the Ising, time it took to map it
        """
        start = start_time_measurement()

        # Map Linear problem from dictionary (generated by pulp) to quadratic program
        qp = self.map_pulp_to_qiskit(problem)
        logging.info(qp.export_as_lp_string())

        # convert quadratic problem to qubo to ising
        conv = QuadraticProgramToQubo()
        qubo = conv.convert(qp)

        variables = [variable.name for variable in qubo.variables]
        qubit_op, _ = qubo.to_ising()

        self.global_variables = variables

        # reverse generate J and t out of qubit PauliSumOperator from qiskit
        t_matrix = np.zeros(qubit_op.num_qubits, dtype=complex)
        j_matrix = np.zeros((qubit_op.num_qubits, qubit_op.num_qubits), dtype=complex)

        for pauli_op in qubit_op:
            pauli_str, coeff = pauli_op.to_list()[0]
            index_pos_list = list(locate(pauli_str, lambda a: a == 'Z'))

            if len(index_pos_list) == 1:
                t_matrix[index_pos_list[0]] = coeff
            elif len(index_pos_list) == 2:
                j_matrix[index_pos_list[0]][index_pos_list[1]] = coeff

        return {"J": j_matrix, "t": t_matrix}, end_time_measurement(start)

    def reverse_map(self, solution: dict) -> tuple[dict, float]:
        """
        Maps the solution back to the representation needed by the ACL class for validation/evaluation.

        :param solution: Dict with a bit_string containing the solution
        :return: Solution mapped accordingly, time it took to map it
        """
        start = start_time_measurement()

        if np.any(solution == "-1"):
            solution = self._convert_ising_to_qubo(solution)

        result = {"status": [0]}
        variables = {}
        objective_value = 0

        for bit in solution:
            if solution[bit] > 0:
                if "x" in self.global_variables[bit]:
                    variables[self.global_variables[bit]] = solution[bit]
                    result["status"] = 'Optimal'
                    objective_value += solution[bit]

        result["variables"] = variables
        result["obj_value"] = objective_value

        return result, end_time_measurement(start)

    @staticmethod
    def _convert_ising_to_qubo(solution: any) -> any:
        solution = np.array(solution)
        with np.nditer(solution, op_flags=['readwrite']) as it:
            for x in it:
                if x == -1:
                    x[...] = 0
        return solution

    def get_default_submodule(self, option: str) -> Core:
        """
        Returns the default submodule based on the provided option.

        :param option: Option specifying the submodule
        :return: Instance of the corresponding submodule
        :raises NotImplementedError: If the option is not recognized
        """
        if option == "QAOA":
            from quark.modules.solvers.QAOA import QAOA  # pylint: disable=C0415
            return QAOA()
        elif option == "QiskitQAOA":
            from quark.modules.solvers.QiskitQAOA import QiskitQAOA  # pylint: disable=C0415
            return QiskitQAOA()
        else:
            raise NotImplementedError(f"Solver Option {option} not implemented")
