# Source code for qutip.qip.device.optpulseprocessor

    from collections.abc import Iterable
    import warnings
    import numbers

    import numpy as np

    from qutip.qobj import Qobj
    import qutip.control.pulseoptim as cpo
    from qutip.operators import identity
    from qutip.tensor import tensor
    from qutip.mesolve import mesolve
    from qutip.qip.circuit import QubitCircuit
    from qutip.qip.device.processor import Processor
    from qutip.qip.operations.gates import gate_sequence_product

    __all__ = ['OptPulseProcessor']

    class OptPulseProcessor(Processor):
        """
        A processor, which takes the Hamiltonian available
        as dynamic generators, calls the
        `qutip.control.optimize_pulse_unitary` function
        to find an optimized pulse sequence for the desired quantum circuit.
        The processor can simulate the evolution under the given
        control pulses using :func:`qutip.mesolve`.
        (For attributes documentation, please
        refer to the parent class :class:`.Processor`)

        Parameters
        ----------
        N: int
            The number of component systems.

        drift: `:class:`qutip.Qobj`
            The drift Hamiltonian. The size must match the whole quantum system.

        t1: list or float
            Characterize the decoherence of amplitude damping for
            each qubit. A list of size `N` or a float for all qubits.

        t2: list of float
            Characterize the decoherence of dephasing for
            each qubit. A list of size `N` or a float for all qubits.

        dims: list
            The dimension of each component system.
            Default value is a
            qubit system of ``dim=[2,2,2,...,2]``
        """
        def __init__(self, N, drift=None, t1=None, t2=None, dims=None):
            super(OptPulseProcessor, self).__init__(
                N, t1=t1, t2=t2, dims=dims)
            if drift is not None:
                self.add_drift(drift, list(range(N)))
            self.spline_kind = "step_func"

        def load_circuit(self, qc, min_fid_err=np.inf, merge_gates=True,
                         setting_args=None, verbose=False, **kwargs):
            """
            Find the pulses realizing a given
            :class:`.QubitCircuit` using
            `qutip.control.optimize_pulse_unitary`. Further parameter for
            for `qutip.control.optimize_pulse_unitary` needs to be given as
            keyword arguments. By default, it first merge all the gates
            into one unitary and then find the control pulses for it.
            It can be turned off and one can set different parameters
            for different gates. See examples for details.

            Examples
            --------

            Same parameter for all the gates:

            .. code-block:: python

                qc = QubitCircuit(N=1)
                qc.add_gate("SNOT", 0)
                num_tslots = 10
                evo_time = 10
                processor = OptPulseProcessor(N=1, drift=sigmaz(),
                                              ctrls=[sigmax()])
                # num_tslots and evo_time are two keyword arguments
                tlist, coeffs = processor.load_circuit(
                    qc, num_tslots=num_tslots, evo_time=evo_time)

            Different parameters for different gates:

            .. code-block:: python

                qc = QubitCircuit(N=2)
                qc.add_gate("SNOT", 0)
                qc.add_gate("SWAP", targets=[0, 1])
                qc.add_gate('CNOT', controls=1, targets=[0])

                processor = OptPulseProcessor(N=2, drift=tensor([sigmaz()]*2))
                processor.add_control(sigmax(), cyclic_permutation=True)
                processor.add_control(sigmay(), cyclic_permutation=True)
                processor.add_control(tensor([sigmay(), sigmay()]))
                setting_args = {"SNOT": {"num_tslots": 10, "evo_time": 1},
                                "SWAP": {"num_tslots": 30, "evo_time": 3},
                                "CNOT": {"num_tslots": 30, "evo_time": 3}}
                tlist, coeffs = processor.load_circuit(qc,
                                                       setting_args=setting_args,
                                                       merge_gates=False)

            Parameters
            ----------
            qc: :class:`.QubitCircuit` or list of Qobj
                The quantum circuit to be translated.

            min_fid_err: float, optional
                The minimal fidelity tolerance, if the fidelity error of any
                gate decomposition is higher, a warning will be given.
                Default is infinite.

            merge_gates: boolean, optimal
                If True, merge all gate/Qobj into one Qobj and then
                find the optimal pulses for this unitary matrix. If False,
                find the optimal pulses for each gate/Qobj.

            setting_args: dict, optional
                Only considered if merge_gates is False.
                It is a dictionary containing keyword arguments
                for different gates.

                E.g.: ::

                    setting_args = {"SNOT": {"num_tslots": 10, "evo_time": 1},
                                    "SWAP": {"num_tslots": 30, "evo_time": 3},
                                    "CNOT": {"num_tslots": 30, "evo_time": 3}}

            verbose: boolean, optional
                If true, the information for each decomposed gate
                will be shown. Default is False.

            **kwargs
                keyword arguments for `qutip.control.optimize_pulse_unitary`

            Returns
            -------
            tlist: array_like
                A NumPy array specifies the time of each coefficient

            coeffs: array_like
                A 2d NumPy array of the shape (len(ctrls), len(tlist)-1). Each
                row corresponds to the control pulse sequence for
                one Hamiltonian.

            Notes
            -----
            ``len(tlist) - 1 = coeffs.shape[1]`` since tlist gives the beginning
            and the end of the pulses.
            """
            if setting_args is None:
                setting_args = {}
            if isinstance(qc, QubitCircuit):
                props = qc.propagators()
                gates = [g.name for g in qc.gates]
            elif isinstance(qc, Iterable):
                props = qc
                gates = None  # using list of Qobj, no gates name
            else:
                raise ValueError(
                    "qc should be a "
                    "QubitCircuit or a list of Qobj")
            if merge_gates:  # merge all gates/Qobj into one Qobj
                props = [gate_sequence_product(props)]
                gates = None

            time_record = []  # a list for all the gates
            coeff_record = []
            last_time = 0.  # used in concatenation of tlist
            for prop_ind, U_targ in enumerate(props):
                U_0 = identity(U_targ.dims[0])

                # If qc is a QubitCircuit and setting_args is not empty,
                # we update the kwargs for each gate.
                # keyword arguments in setting_arg have priority
                if gates is not None and setting_args:
                    kwargs.update(setting_args[gates[prop_ind]])

                full_drift_ham = self.drift.get_ideal_qobjevo(self.dims).cte
                full_ctrls_hams = [pulse.get_ideal_qobj(self.dims)
                                   for pulse in self.pulses]
                result = cpo.optimize_pulse_unitary(
                    full_drift_ham, full_ctrls_hams, U_0, U_targ, **kwargs)

                if result.fid_err > min_fid_err:
                    warnings.warn(
                        "The fidelity error of gate {} is higher "
                        "than required limit. Use verbose=True to see"
                        "the more detailed information.".format(prop_ind))

                time_record.append(result.time[1:] + last_time)
                last_time += result.time[-1]
                coeff_record.append(result.final_amps.T)

                if verbose:
                    print("********** Gate {} **********".format(prop_ind))
                    print("Final fidelity error {}".format(result.fid_err))
                    print("Final gradient normal {}".format(
                                                    result.grad_norm_final))
                    print("Terminated due to {}".format(result.termination_reason))
                    print("Number of iterations {}".format(result.num_iter))

            tlist = np.hstack([[0.]] + time_record)
            for i in range(len(self.pulses)):
                self.pulses[i].tlist = tlist
            coeffs = np.vstack([np.hstack(coeff_record)])
            self.coeffs = coeffs

            return tlist, coeffs