# Source code for qutip.mcsolve

    __all__ = ['mcsolve']

    import os
    import numpy as np
    from numpy.random import RandomState, randint
    import scipy.sparse as sp
    from scipy.integrate import ode
    from scipy.integrate._ode import zvode

    from types import FunctionType, BuiltinFunctionType
    from functools import partial
    from qutip.fastsparse import csr2fast
    from qutip.qobj import Qobj
    from qutip.qobjevo import QobjEvo
    from qutip.parallel import parfor, parallel_map, serial_map
    from qutip.cy.mcsolve import CyMcOde, CyMcOdeDiag
    from qutip.cy.spconvert import dense1D_to_fastcsr_ket
    from qutip.sesolve import sesolve
    from qutip.solver import (Options, Result, ExpectOps,
                              solver_safe, SolverSystem)
    from qutip.settings import debug
    from qutip.ui.progressbar import TextProgressBar, BaseProgressBar
    import qutip.settings

    if debug:
        import inspect

    #
    # Internal, global variables for storing references to dynamically loaded
    # cython functions

    # Todo: use real warning
    def warn(text):
        print(text)

    class qutip_zvode(zvode):
        def step(self, *args):
            itask = self.call_args[2]
            self.rwork[0] = args[4]
            self.call_args[2] = 5
            r = self.run(*args)
            self.call_args[2] = itask
            return r

    def mcsolve(H, psi0, tlist, c_ops=[], e_ops=[], ntraj=0,
                args={}, options=None, progress_bar=True,
                map_func=parallel_map, map_kwargs={}, _safe_mode=True):
        r"""Monte Carlo evolution of a state vector :math:`|\psi \rangle` for a
        given Hamiltonian and sets of collapse operators, and possibly, operators
        for calculating expectation values. Options for the underlying ODE solver
        are given by the Options class.

        mcsolve supports time-dependent Hamiltonians and collapse operators using
        either Python functions of strings to represent time-dependent
        coefficients. Note that, the system Hamiltonian MUST have at least one
        constant term.

        As an example of a time-dependent problem, consider a Hamiltonian with two
        terms ``H0`` and ``H1``, where ``H1`` is time-dependent with coefficient
        ``sin(w*t)``, and collapse operators ``C0`` and ``C1``, where ``C1`` is
        time-dependent with coeffcient ``exp(-a*t)``.  Here, w and a are constant
        arguments with values ``W`` and ``A``.

        Using the Python function time-dependent format requires two Python
        functions, one for each collapse coefficient. Therefore, this problem could
        be expressed as::

            def H1_coeff(t,args):
                return sin(args['w']*t)

            def C1_coeff(t,args):
                return exp(-args['a']*t)

            H = [H0, [H1, H1_coeff]]

            c_ops = [C0, [C1, C1_coeff]]

            args={'a': A, 'w': W}

        or in String (Cython) format we could write::

            H = [H0, [H1, 'sin(w*t)']]

            c_ops = [C0, [C1, 'exp(-a*t)']]

            args={'a': A, 'w': W}

        Constant terms are preferably placed first in the Hamiltonian and collapse
        operator lists.

        Parameters
        ----------
        H : :class:`qutip.Qobj`, ``list``
            System Hamiltonian.

        psi0 : :class:`qutip.Qobj`
            Initial state vector

        tlist : array_like
            Times at which results are recorded.

        ntraj : int
            Number of trajectories to run.

        c_ops : :class:`qutip.Qobj`, ``list``
            single collapse operator or a ``list`` of collapse operators.

        e_ops : :class:`qutip.Qobj`, ``list``
            single operator as Qobj or ``list`` or equivalent of Qobj operators
            for calculating expectation values.

        args : dict
            Arguments for time-dependent Hamiltonian and collapse operator terms.

        options : Options
            Instance of ODE solver options.

        progress_bar: BaseProgressBar
            Optional instance of BaseProgressBar, or a subclass thereof, for
            showing the progress of the simulation. Set to None to disable the
            progress bar.

        map_func: function
            A map function for managing the calls to the single-trajactory solver.

        map_kwargs: dictionary
            Optional keyword arguments to the map_func function.

        Returns
        -------
        results : :class:`qutip.solver.Result`
            Object storing all results from the simulation.

        .. note::

            It is possible to reuse the random number seeds from a previous run
            of the mcsolver by passing the output Result object seeds via the
            Options class, i.e. Options(seeds=prev_result.seeds).
        """
        if isinstance(c_ops, (Qobj, QobjEvo)):
            c_ops = [c_ops]

        if options is None:
            options = Options()

        if options.rhs_reuse and not isinstance(H, SolverSystem):
            # TODO: deprecate when going to class based solver.
            if "mcsolve" in solver_safe:
                # print(" ")
                H = solver_safe["mcsolve"]
            else:
                pass
                # raise Exception("Could not find the Hamiltonian to reuse.")

        if not ntraj:
            ntraj = options.ntraj

        if len(c_ops) == 0 and not options.rhs_reuse:
            warn("No c_ops, using sesolve")
            return sesolve(H, psi0, tlist, e_ops=e_ops, args=args,
                           options=options, progress_bar=progress_bar,
                           _safe_mode=_safe_mode)

        try:
            num_traj = int(ntraj)
        except TypeError:
            num_traj = max(ntraj)

        # set the physics
        if not psi0.isket:
            raise Exception("Initial state must be a state vector.")

        # load monte carlo class
        mc = _MC(options)

        if isinstance(H, SolverSystem):
            mc.ss = H
        else:
            mc.make_system(H, c_ops, tlist, args, options)

        mc.reset(tlist[0], psi0)

        mc.set_e_ops(e_ops)

        if options.seeds is not None:
            mc.seed(num_traj, options.seeds)

        if _safe_mode:
            mc.run_test()

        # Run the simulation
        mc.run(num_traj=num_traj, tlist=tlist,
               progress_bar=progress_bar,
               map_func=map_func, map_kwargs=map_kwargs)

        return mc.get_result(ntraj)

    # -----------------------------------------------------------------------------
    # MONTE CARLO CLASS
    # -----------------------------------------------------------------------------
    class _MC():
        """
        Private class for solving Monte Carlo evolution from mcsolve
        """
        def __init__(self, options=None):
            if options is None:
                options = Options()
            self.options = options
            self.ss = None
            self.tlist = None
            self.e_ops = None
            self.ran = False
            self.psi0 = None
            self.seeds = []
            self.t = 0.
            self.num_traj = 0
            self.args_col = None

            self._psi_out = []
            self._expect_out = []
            self._collapse = []
            self._ss_out = []

        def reset(self, t=0., psi0=None):
            if psi0 is not None:
                self.psi0 = psi0
            if self.psi0 is not None:
                self.initial_vector = self.psi0.full().ravel("F")
                if self.ss is not None and self.ss.type == "Diagonal":
                    self.initial_vector = np.dot(self.ss.Ud, self.initial_vector)

            self.t = t
            self.ran = False
            self._psi_out = []
            self._expect_out = []
            self._collapse = []
            self._ss_out = []

        def seed(self, ntraj, seeds=[]):
            # setup seeds array
            np.random.seed()
            try:
                seed = int(seeds)
                np.random.seed(seed)
                seeds = []
            except TypeError:
                pass

            if len(seeds) < ntraj:
                self.seeds = seeds + list(randint(0, 2**32,
                                                  size=ntraj-len(seeds),
                                                  dtype=np.uint32))
            else:
                self.seeds = seeds[:ntraj]

        def make_system(self, H, c_ops, tlist=None, args={}, options=None):
            if options is None:
                options = self.options
            else:
                self.options = options
            var = _collapse_args(args)

            ss = SolverSystem()
            ss.td_c_ops = []
            ss.td_n_ops = []
            ss.args = args
            ss.col_args = var
            for c in c_ops:
                cevo = QobjEvo(c, args, tlist=tlist)
                cdc = cevo._cdc()
                cevo.compile()
                cdc.compile()
                ss.td_c_ops.append(cevo)
                ss.td_n_ops.append(cdc)

            try:
                H_td = QobjEvo(H, args, tlist=tlist)
                H_td *= -1j
                for c in ss.td_n_ops:
                    H_td += -0.5 * c
                if options.rhs_with_state:
                    H_td._check_old_with_state()
                H_td.compile()
                ss.H_td = H_td
                ss.makefunc = _qobjevo_set
                ss.set_args = _qobjevo_args
                ss.type = "QobjEvo"

            except Exception:
                ss.h_func = H
                ss.Hc_td = -0.5 * sum(ss.td_n_ops)
                ss.Hc_td.compile()
                ss.with_state = options.rhs_with_state
                ss.makefunc = _func_set
                ss.set_args = _func_args
                ss.type = "callback"

            solver_safe["mcsolve"] = ss
            self.ss = ss
            self.reset()

        def set_e_ops(self, e_ops=[]):
            if e_ops:
                self.e_ops = ExpectOps(e_ops)
            else:
                self.e_ops = ExpectOps([])

            self.e_ops.check_dims(self.ss.td_c_ops[0].cte.dims)

            ss = self.ss
            if ss is not None and ss.type == "Diagonal" and not self.e_ops.isfunc:
                e_ops = [
                    Qobj(ss.Ud @ e.full() @ ss.U, dims=e.dims)
                    for e in self.e_ops.e_ops
                ]
                self.e_ops = ExpectOps(e_ops)

            if not self.e_ops:
                self.options.store_states = True

        def run_test(self):
            try:
                for c_op in self.ss.td_c_ops:
                    c_op.mul_vec(0, self.psi0)
            except Exception as e:
                raise Exception("c_ops are not consistant with psi0") from e

            if self.ss.type == "QobjEvo":
                try:
                    self.ss.H_td.mul_vec(0., self.psi0)
                except Exception as e:
                    raise Exception("Error calculating H") from e
            else:
                try:
                    rhs, ode_args = self.ss.makefunc(self.ss)
                    rhs(0, self.psi0.full().ravel(), ode_args)
                except Exception as e:
                    raise Exception("Error calculating H") from e

        def run(self, num_traj=0, psi0=None, tlist=None,
                args={}, e_ops=None, options=None,
                progress_bar=True,
                map_func=parallel_map, map_kwargs={}):
            # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            # 4 situation for run:
            # - first run
            # - change parameters
            # - add  trajectories
            #       (self.add_traj)      Not Implemented
            # - continue from the last time and states
            #       (self.continue_runs) Not Implemented
            # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            options = options if options is not None else self.options

            if self.ran and tlist[0] == self.t:
                # psi0 is ignored since we restart from a
                # different states for each trajectories
                self.continue_runs(num_traj, tlist, args, e_ops, options,
                                   progress_bar, map_func, map_kwargs)
                return

            if args and args != self.ss.args:
                self.ss.set_args(self.ss, args)
                self.reset()

            if e_ops and e_ops != self.e_ops:
                self.set_e_ops(e_ops)
                self.reset()

            if psi0 is not None and psi0 != self.psi0:
                self.psi0 = psi0
                self.reset()

            tlist = np.array(tlist)
            if tlist is not None and np.all(tlist != self.tlist):
                self.tlist = tlist
                self.reset()

            if self.ran:
                if options.store_states and self._psi_out[0].shape[0] == 1:
                    self.reset()
                else:
                    # if not reset here, add trajectories
                    self.add_traj(num_traj, progress_bar, map_func, map_kwargs)
                    return

            if not num_traj:
                num_traj = options.ntraj

            if options.num_cpus == 1 or num_traj == 1:
                map_func = serial_map

            if len(self.seeds) != num_traj:
                self.seed(num_traj, self.seeds)

            if not progress_bar:
                progress_bar = BaseProgressBar()
            elif progress_bar is True:
                progress_bar = TextProgressBar()

            # set arguments for input to monte carlo
            map_kwargs_ = {'progress_bar': progress_bar,
                           'num_cpus': options.num_cpus}
            map_kwargs_.update(map_kwargs)
            map_kwargs = map_kwargs_

            if self.e_ops is None:
                self.set_e_ops()

            if self.ss.type == "Diagonal":
                results = map_func(self._single_traj_diag, list(range(num_traj)),
                                   **map_kwargs)
            else:
                results = map_func(self._single_traj, list(range(num_traj)),
                                   **map_kwargs)

            self.t = self.tlist[-1]
            self.num_traj = num_traj
            self.ran = True

            for result in results:
                state_out, ss_out, expect, collapse = result
                self._psi_out.append(state_out)
                self._ss_out.append(ss_out)
                self._expect_out.append(expect)
                self._collapse.append(collapse)
            self._psi_out = np.stack(self._psi_out)
            self._ss_out = np.stack(self._ss_out)

        def add_traj(self, num_traj,
                     progress_bar=True,
                     map_func=parallel_map, map_kwargs={}):
            raise NotImplementedError

        def continue_runs(self, num_traj, tlist, args={}, e_ops=[], options=None,
                          progress_bar=True,
                          map_func=parallel_map, map_kwargs={}):
            raise NotImplementedError

        # --------------------------------------------------------------------------
        # results functions
        # --------------------------------------------------------------------------
        @property
        def states(self):
            dims = self.psi0.dims[0]
            len_ = self._psi_out.shape[2]
            if self._psi_out.shape[1] == 1:
                dm_t = np.zeros((len_, len_), dtype=complex)
                for i in range(self.num_traj):
                    vec = self._psi_out[i, 0]
                    dm_t += np.outer(vec, vec.conj())
                return Qobj(dm_t/self.num_traj, dims=[dims, dims])
            else:
                states = np.empty((len(self.tlist)), dtype=object)
                for j in range(len(self.tlist)):
                    dm_t = np.zeros((len_, len_), dtype=complex)
                    for i in range(self.num_traj):
                        vec = self._psi_out[i, j]
                        dm_t += np.outer(vec, vec.conj())
                    states[j] = Qobj(dm_t/self.num_traj, dims=[dims, dims])
                return states

        @property
        def final_state(self):
            dims = self.psi0.dims[0]
            len_ = self._psi_out.shape[2]
            dm_t = np.zeros((len_, len_), dtype=complex)
            for i in range(self.num_traj):
                vec = self._psi_out[i, -1]
                dm_t += np.outer(vec, vec.conj())
            return Qobj(dm_t/self.num_traj, dims=[dims, dims])

        @property
        def runs_final_states(self):
            dims = self.psi0.dims[0]
            psis = np.empty((self.num_traj), dtype=object)
            for i in range(self.num_traj):
                psis[i] = Qobj(dense1D_to_fastcsr_ket(self._psi_out[i, -1]),
                               dims=dims, fast='mc')
            return psis

        @property
        def expect(self):
            return self.expect_traj_avg()

        @property
        def runs_expect(self):
            return [expt.finish() for expt in self._expect_out]

        def expect_traj_avg(self, ntraj=0):
            if not ntraj:
                ntraj = len(self._expect_out)
            expect = np.stack([expt.raw_out for expt in self._expect_out[:ntraj]])
            expect = np.mean(expect, axis=0)

            result = []
            for ii in range(self.e_ops.e_num):
                if self.e_ops.e_ops_isherm[ii]:
                    result.append(np.real(expect[ii, :]))
                else:
                    result.append(expect[ii, :])

            if self.e_ops.e_ops_dict:
                result = {e: result[n]
                          for n, e in enumerate(self.e_ops.e_ops_dict.keys())}
            return result

        @property
        def steady_state(self):
            if self._ss_out is not None:
                dims = self.psi0.dims[0]
                len_ = self.psi0.shape[0]
                return Qobj(np.mean(self._ss_out, axis=0),
                            dims=[dims, dims], shape=(len_, len_))
            # TO-DO rebuild steady_state from _psi_out if needed
            # elif self._psi_out is not None:
            #     return sum(self.state_average) / self.num_traj
            else:
                return None

        @property
        def runs_states(self):
            dims = self.psi0.dims
            psis = np.empty((self.num_traj, len(self.tlist)), dtype=object)
            for i in range(self.num_traj):
                for j in range(len(self.tlist)):
                    psis[i, j] = Qobj(dense1D_to_fastcsr_ket(self._psi_out[i, j]),
                                      dims=dims, fast='mc')
            return psis

        @property
        def collapse(self):
            return self._collapse

        @property
        def collapse_times(self):
            out = []
            for col_ in self._collapse:
                col = list(zip(*col_))
                col = ([] if len(col) == 0 else col[0])
                out.append(np.array(col))
            return out
            return [np.array(list(zip(*col_))[0]) for col_ in self._collapse]

        @property
        def collapse_which(self):
            out = []
            for col_ in self._collapse:
                col = list(zip(*col_))
                col = ([] if len(col) == 0 else col[1])
                out.append(np.array(col))
            return out
            return [np.array(list(zip(*col_))[1]) for col_ in self._collapse]

        def get_result(self, ntraj=[]):
            # Store results in the Result object
            if not ntraj:
                ntraj = [self.num_traj]
            elif not isinstance(ntraj, list):
                ntraj = [ntraj]

            output = Result()
            output.solver = 'mcsolve'
            output.seeds = self.seeds

            options = self.options
            output.options = options

            if options.steady_state_average:
                output.states = self.steady_state
            elif options.average_states and options.store_states:
                output.states = self.states
            elif options.store_states:
                output.states = self.runs_states

            if options.store_final_state:
                if options.average_states:
                    output.final_state = self.final_state
                else:
                    output.final_state = self.runs_final_states

            if options.average_expect:
                output.expect = [self.expect_traj_avg(n) for n in ntraj]
                if len(output.expect) == 1:
                    output.expect = output.expect[0]
            else:
                output.expect = self.runs_expect

            # simulation parameters
            output.times = self.tlist
            output.num_expect = self.e_ops.e_num
            output.num_collapse = len(self.ss.td_c_ops)
            output.ntraj = self.num_traj
            output.col_times = self.collapse_times
            output.col_which = self.collapse_which

            return output

        # --------------------------------------------------------------------------
        # single-trajectory for monte carlo
        # --------------------------------------------------------------------------
        def _single_traj(self, nt):
            """
            Monte Carlo algorithm returning state-vector or expectation values
            at times tlist for a single trajectory.
            """
            # SEED AND RNG AND GENERATE
            prng = RandomState(self.seeds[nt])
            opt = self.options

            # set initial conditions
            ss = self.ss
            tlist = self.tlist
            e_ops = self.e_ops.copy()
            opt = self.options
            rhs, ode_args = self.ss.makefunc(ss)
            ODE = self._build_integration_func(rhs, ode_args, opt)
            ODE.set_initial_value(self.initial_vector, tlist[0])
            e_ops.init(tlist)

            cymc = CyMcOde(ss, opt)
            states_out, ss_out, collapses = cymc.run_ode(ODE, tlist, e_ops, prng)

            # Run at end of mc_alg function
            # -----------------------------
            if opt.steady_state_average:
                ss_out /= float(len(tlist))

            return (states_out, ss_out, e_ops, collapses)

        def _build_integration_func(self, rhs, ode_args, opt):
            """
            Create the integration function while fixing the parameters
            """
            ODE = ode(rhs)
            if ode_args:
                ODE.set_f_params(ode_args)
            # initialize ODE solver for RHS
            ODE.set_integrator('zvode', method="adams")
            ODE._integrator = qutip_zvode(
                method=opt.method, order=opt.order, atol=opt.atol,
                rtol=opt.rtol, nsteps=opt.nsteps, first_step=opt.first_step,
                min_step=opt.min_step, max_step=opt.max_step)
            return ODE

        # --------------------------------------------------------------------------
        # In development diagonalize the Hamiltonian before solving
        # Same seeds give same evolution
        # 3~5 time faster
        # constant system only.
        # --------------------------------------------------------------------------
        def make_diag_system(self, H, c_ops):
            ss = SolverSystem()
            ss.td_c_ops = []
            ss.td_n_ops = []

            H_ = H.copy()
            H_ *= -1j
            for c in c_ops:
                H_ += -0.5 * c.dag() * c

            w, v = np.linalg.eig(H_.full())
            arg = np.argsort(np.abs(w))
            eig = w[arg]
            U = v.T[arg].T
            Ud = U.T.conj()

            for c in c_ops:
                c_diag = Qobj(Ud @ c.full() @ U, dims=c.dims)
                cevo = QobjEvo(c_diag)
                cdc = cevo._cdc()
                cevo.compile()
                cdc.compile()
                ss.td_c_ops.append(cevo)
                ss.td_n_ops.append(cdc)

            ss.H_diag = eig
            ss.Ud = Ud
            ss.U = U
            ss.args = {}
            ss.type = "Diagonal"
            solver_safe["mcsolve"] = ss

            if self.e_ops and not self.e_ops.isfunc:
                e_ops = [
                    Qobj(Ud @ e.full() @ U, dims=e.dims)
                    for e in self.e_ops.e_ops
                ]
                self.e_ops = ExpectOps(e_ops)
            self.ss = ss
            self.reset()

        def _single_traj_diag(self, nt):
            """
            Monte Carlo algorithm returning state-vector or expectation values
            at times tlist for a single trajectory.
            """
            # SEED AND RNG AND GENERATE
            prng = RandomState(self.seeds[nt])
            opt = self.options

            ss = self.ss
            tlist = self.tlist
            e_ops = self.e_ops.copy()
            opt = self.options
            e_ops.init(tlist)

            cymc = CyMcOdeDiag(ss, opt)
            states_out, ss_out, collapses =\
                cymc.run_ode(self.initial_vector, tlist, e_ops, prng)

            if opt.steady_state_average:
                ss_out = ss.U @ ss_out @ ss.Ud
            states_out = np.inner(ss.U, states_out).T
            if opt.steady_state_average:
                ss_out /= float(len(tlist))
            return (states_out, ss_out, e_ops, collapses)

    # -----------------------------------------------------------------------------
    # CODES FOR PYTHON FUNCTION BASED TIME-DEPENDENT RHS
    # -----------------------------------------------------------------------------
    def _qobjevo_set(ss, psi0=None, args={}, opt=None):
        if args:
            self.set_args(args)
        rhs = ss.H_td.compiled_qobjevo.mul_vec
        return rhs, ()

    def _qobjevo_args(ss, args):
        var = _collapse_args(args)
        ss.col_args = var
        ss.args = args
        ss.H_td.solver_set_args(args, psi0, e_ops)
        for c in ss.td_c_ops:
            c.solver_set_args(args, psi0, e_ops)
        for c in ss.td_n_ops:
            c.solver_set_args(args, psi0, e_ops)

    def _func_set(HS, psi0=None, args={}, opt=None):
        if args:
            self.set_args(args)
        else:
            args = ss.args
        if ss.with_state:
            rhs = _funcrhs
        else:
            rhs = _funcrhs_with_state
        return rhs, (ss.h_func, ss.Hc_td, args)

    def _func_args(ss, args):
        var = _collapse_args(args)
        ss.col_args = var
        ss.args = args
        for c in ss.td_c_ops:
            c.solver_set_args(args, psi0, e_ops)
        for c in ss.td_n_ops:
            c.solver_set_args(args, psi0, e_ops)
        return rhs, (ss.h_func, ss.Hc_td, args)

    # RHS of ODE for python function Hamiltonian
    def _funcrhs(t, psi, h_func, Hc_td, args):
        h_func_data = -1.0j * h_func(t, args).data
        h_func_term = h_func_data * psi
        return h_func_term + Hc_td.mul_vec(t, psi)

    def _funcrhs_with_state(t, psi, h_func, Hc_td, args):
        h_func_data = - 1.0j * h_func(t, psi, args).data
        h_func_term = h_func_data * psi
        return h_func_term + Hc_td.mul_vec(t, psi)

    def _mc_dm_avg(psi_list):
        """
        Private function that averages density matrices in parallel
        over all trajectories for a single time using parfor.
        """
        ln = len(psi_list)
        dims = psi_list[0].dims
        shape = psi_list[0].shape
        out_data = sum([psi.data for psi in psi_list]) / ln
        return Qobj(out_data, dims=dims, shape=shape, fast='mc-dm')

    def _collapse_args(args):
        for key in args:
            if key == "collapse":
                if not isinstance(args[key], list):
                    args[key] = []
                return key
        return ""