# This code is part of Qiskit.
#
# (C) Copyright IBM 2019.
#
# 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.
"""
Noise model inserter module
The goal of this module is to add QuantumError gates (Kraus gates) to a circuit
based on a given noise model.
"""
import qiskit.compiler


def insert_noise(circuits, noise_model, transpile=False):
    """Return a noisy version of a QuantumCircuit.

    Args:
        circuits (QuantumCircuit or list[QuantumCircuit]): Input noise-free circuits.
        noise_model (NoiseModel):  The noise model containing the errors to add
        transpile (Boolean): Should the circuit be transpiled into the noise model basis gates

    Returns:
        QuantumCircuit: The new circuit with the Kraus noise instructions inserted.

    Additional Information:
        The noisy circuit return by this function will consist of the
        original circuit with ``Kraus`` instructions inserted after all
        instructions referenced in the ``noise_model``. The resulting circuit
        cannot be ran on a quantum computer but can be executed on the
        :class:`~qiskit.providers.aer.QasmSimulator`.
    """
    is_circuits_list = isinstance(circuits, (list, tuple))
    circuits = circuits if is_circuits_list else [circuits]
    result_circuits = []
    nonlocal_errors = noise_model._nonlocal_quantum_errors
    local_errors = noise_model._local_quantum_errors
    default_errors = noise_model._default_quantum_errors
    for circuit in circuits:
        if transpile:
            transpiled_circuit = qiskit.compiler.transpile(circuit,
                                                           basis_gates=noise_model.basis_gates)
        else:
            transpiled_circuit = circuit
        qubit_indices = {bit: index for index, bit in enumerate(transpiled_circuit.qubits)}
        result_circuit = transpiled_circuit.copy(name=transpiled_circuit.name + '_with_noise')
        result_circuit.data = []
        for inst, qargs, cargs in transpiled_circuit.data:
            result_circuit.data.append((inst, qargs, cargs))
            qubits = tuple(qubit_indices[q] for q in qargs)
            # Priority for error model used:
            # nonlocal error > local error > default error
            if inst.name in nonlocal_errors and qubits in nonlocal_errors[inst.name]:
                for noise_qubits, error in nonlocal_errors[inst.name][qubits].items():
                    result_circuit.append(error.to_instruction(), noise_qubits)
            elif inst.name in local_errors and qubits in local_errors[inst.name]:
                error = local_errors[inst.name][qubits]
                result_circuit.append(error.to_instruction(), qargs)
            elif inst.name in default_errors.keys():
                error = default_errors[inst.name]
                result_circuit.append(error.to_instruction(), qargs)
        result_circuits.append(result_circuit)
    return result_circuits if is_circuits_list else result_circuits[0]
