# This code is part of Qiskit.
#
# (C) Copyright IBM 2019, 2021.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

""" Measurement error mitigation """

import copy

from qiskit import compiler
from ..exceptions import QiskitError, MissingOptionalLibraryError

# pylint: disable=invalid-name


def get_measured_qubits(transpiled_circuits):
    """
    Retrieve the measured qubits from transpiled circuits.

    Args:
        transpiled_circuits ([QuantumCircuit]): a list of transpiled circuits

    Returns:
        list[int]: the used and sorted qubit index
        dict: key is qubit index str connected by '_',
              value is the experiment index. {str: list[int]}
    Raises:
        QiskitError: invalid qubit mapping
    """
    qubit_index = None
    qubit_mappings = {}
    for idx, qc in enumerate(transpiled_circuits):
        measured_qubits = []
        for inst, qargs, _ in qc.data:
            if inst.name != "measure":
                continue
            for qreg in qc.qregs:
                if qargs[0] in qreg:
                    index = qreg[:].index(qargs[0])
                    measured_qubits.append(index)
                    break
        measured_qubits_str = "_".join([str(x) for x in measured_qubits])
        if measured_qubits_str not in qubit_mappings:
            qubit_mappings[measured_qubits_str] = []
        qubit_mappings[measured_qubits_str].append(idx)
        if qubit_index is None:
            qubit_index = measured_qubits
        elif set(qubit_index) != set(measured_qubits):
            raise QiskitError(
                "The used qubit index are different. ({}) vs ({}).\nCurrently, "
                "we only support all circuits using the same set of qubits "
                "regardless qubit order.".format(qubit_index, measured_qubits)
            )

    return sorted(qubit_index), qubit_mappings


def get_measured_qubits_from_qobj(qobj):
    """
    Retrieve the measured qubits from transpiled circuits.

    Args:
        qobj (QasmObj): qobj

    Returns:
        list[int]: the used and sorted qubit index
        dict: key is qubit index str connected by '_',
              value is the experiment index. {str: list[int]}
     Raises:
        QiskitError: invalid qubit mapping
    """

    qubit_index = None
    qubit_mappings = {}

    for idx, exp in enumerate(qobj.experiments):
        measured_qubits = []
        for instr in exp.instructions:
            if instr.name != "measure":
                continue
            measured_qubits.append(instr.qubits[0])
        measured_qubits_str = "_".join([str(x) for x in measured_qubits])
        if measured_qubits_str not in qubit_mappings:
            qubit_mappings[measured_qubits_str] = []
        qubit_mappings[measured_qubits_str].append(idx)
        if qubit_index is None:
            qubit_index = measured_qubits
        else:
            if set(qubit_index) != set(measured_qubits):
                raise QiskitError(
                    "The used qubit index are different. ({}) vs ({}).\nCurrently, "
                    "we only support all circuits using the same set of qubits "
                    "regardless qubit order.".format(qubit_index, measured_qubits)
                )

    return sorted(qubit_index), qubit_mappings


def build_measurement_error_mitigation_circuits(
    qubit_list, fitter_cls, backend, backend_config=None, compile_config=None
):
    """Build measurement error mitigation circuits
    Args:
        qubit_list (list[int]): list of ordered qubits used in the algorithm
        fitter_cls (callable): CompleteMeasFitter or TensoredMeasFitter
        backend (BaseBackend): backend instance
        backend_config (dict, optional): configuration for backend
        compile_config (dict, optional): configuration for compilation
    Returns:
        QasmQobj: the Qobj with calibration circuits at the beginning
        list[str]: the state labels for build MeasFitter
        list[str]: the labels of the calibration circuits
    Raises:
        QiskitError: when the fitter_cls is not recognizable.
        MissingOptionalLibraryError: Qiskit-Ignis not installed
    """
    try:
        from spinqit.qiskit.ignis.mitigation.measurement import (
            complete_meas_cal,
            CompleteMeasFitter,
            TensoredMeasFitter,
        )
    except ImportError as ex:
        raise MissingOptionalLibraryError(
            libname="qiskit-ignis",
            name="build_measurement_error_mitigation_qobj",
            pip_install="pip install qiskit-ignis",
        ) from ex

    circlabel = "mcal"

    if not qubit_list:
        raise QiskitError("The measured qubit list can not be [].")

    if fitter_cls == CompleteMeasFitter:
        meas_calibs_circuits, state_labels = complete_meas_cal(
            qubit_list=range(len(qubit_list)), circlabel=circlabel
        )
    elif fitter_cls == TensoredMeasFitter:
        # TODO support different calibration
        raise QiskitError("Does not support TensoredMeasFitter yet.")
    else:
        raise QiskitError("Unknown fitter {}".format(fitter_cls))

    # the provided `qubit_list` would be used as the initial layout to
    # assure the consistent qubit mapping used in the main circuits.

    tmp_compile_config = copy.deepcopy(compile_config)
    tmp_compile_config["initial_layout"] = qubit_list
    t_meas_calibs_circuits = compiler.transpile(
        meas_calibs_circuits, backend, **backend_config, **tmp_compile_config
    )
    return t_meas_calibs_circuits, state_labels, circlabel


def build_measurement_error_mitigation_qobj(
    qubit_list, fitter_cls, backend, backend_config=None, compile_config=None, run_config=None
):
    """
    Args:
        qubit_list (list[int]): list of ordered qubits used in the algorithm
        fitter_cls (callable): CompleteMeasFitter or TensoredMeasFitter
        backend (BaseBackend): backend instance
        backend_config (dict, optional): configuration for backend
        compile_config (dict, optional): configuration for compilation
        run_config (RunConfig, optional): configuration for running a circuit

    Returns:
        QasmQobj: the Qobj with calibration circuits at the beginning
        list[str]: the state labels for build MeasFitter
        list[str]: the labels of the calibration circuits

    Raises:
        QiskitError: when the fitter_cls is not recognizable.
        MissingOptionalLibraryError: Qiskit-Ignis not installed
    """
    try:
        from spinqit.qiskit.ignis.mitigation.measurement import (
            complete_meas_cal,
            CompleteMeasFitter,
            TensoredMeasFitter,
        )
    except ImportError as ex:
        raise MissingOptionalLibraryError(
            libname="qiskit-ignis",
            name="build_measurement_error_mitigation_qobj",
            pip_install="pip install qiskit-ignis",
        ) from ex

    circlabel = "mcal"

    if not qubit_list:
        raise QiskitError("The measured qubit list can not be [].")

    if fitter_cls == CompleteMeasFitter:
        meas_calibs_circuits, state_labels = complete_meas_cal(
            qubit_list=range(len(qubit_list)), circlabel=circlabel
        )
    elif fitter_cls == TensoredMeasFitter:
        # TODO support different calibration
        raise QiskitError("Does not support TensoredMeasFitter yet.")
    else:
        raise QiskitError("Unknown fitter {}".format(fitter_cls))

    # the provided `qubit_list` would be used as the initial layout to
    # assure the consistent qubit mapping used in the main circuits.

    tmp_compile_config = copy.deepcopy(compile_config)
    tmp_compile_config["initial_layout"] = qubit_list
    t_meas_calibs_circuits = compiler.transpile(
        meas_calibs_circuits, backend, **backend_config, **tmp_compile_config
    )
    cals_qobj = compiler.assemble(t_meas_calibs_circuits, backend, **run_config.to_dict())
    if hasattr(cals_qobj.config, "parameterizations"):
        del cals_qobj.config.parameterizations
    return cals_qobj, state_labels, circlabel
