# This file is part of QuTiP.
#
#    QuTiP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    QuTiP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QuTiP.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2011 and later, Paul D. Nation & Robert J. Johansson
#
###########################################################################

import sys
import os
import time
import copy
import numpy as np
from types import FunctionType
from multiprocessing import Pool, cpu_count
from numpy.random import RandomState, random_integers

from scipy import arange, array, cumsum, mean, ndarray, setdiff1d, sort, zeros
from scipy.integrate import ode
from scipy.linalg import norm
from scipy.linalg.blas import get_blas_funcs
dznrm2 = get_blas_funcs("znrm2", dtype=np.float64)

from qutip.qobj import *
from qutip.expect import *
from qutip.states import ket2dm
from qutip.parfor import parfor
from qutip.odeoptions import Odeoptions
from qutip.odeconfig import odeconfig
from qutip.cyQ.spmatfuncs import cy_ode_rhs, cy_expect, spmv
from qutip.cyQ.codegen import Codegen
from qutip.odedata import Odedata
from qutip.odechecks import _ode_checks
import qutip.settings
from qutip.settings import debug
import threading
threading._DummyThread._Thread__stop = lambda x: 42

from qutip.gui.progressbar import BaseProgressBar

if debug:
    import inspect

#
# Internal, global variables for storing references to dynamically loaded
# cython functions
#
_cy_col_spmv_func = None
_cy_col_expect_func = None
_cy_col_spmv_call_func = None
_cy_col_expect_call_func = None
_cy_rhs_func = None


def mcsolve(H, psi0, tlist, c_ops, e_ops, ntraj=None,
            args={}, options=Odeoptions(), progress_bar=BaseProgressBar()):
    """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 Odeoptions 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_op_list=[C0,[C1,C1_coeff]]

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

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

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

        c_op_list=[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 : qobj
        System Hamiltonian.
    psi0 : qobj
        Initial state vector
    tlist : array_like
        Times at which results are recorded.
    ntraj : int
        Number of trajectories to run.
    c_ops : array_like
        single collapse operator or ``list`` or ``array`` of collapse
        operators.
    e_ops : array_like
        single operator or ``list`` or ``array`` of operators for calculating
        expectation values.
    args : dict
        Arguments for time-dependent Hamiltonian and collapse operator terms.
    options : Odeoptions
        Instance of ODE solver options.

    Returns
    -------
    results : Odedata
        Object storing all results from simulation.

    """

    if debug:
        print(inspect.stack()[0][3])

    if ntraj is None:
        ntraj = options.ntraj

    if psi0.type != 'ket':
        raise Exception("Initial state must be a state vector.")

    if isinstance(c_ops, Qobj):
        c_ops = [c_ops]

    if isinstance(e_ops, Qobj):
        e_ops = [e_ops]

    if isinstance(e_ops, dict):
        e_ops_dict = e_ops
        e_ops = [e for e in e_ops.values()]
    else:
        e_ops_dict = None

    odeconfig.options = options
    odeconfig.progress_bar = progress_bar

    # set num_cpus to the value given in qutip.settings if none in Odeoptions
    if not odeconfig.options.num_cpus:
        odeconfig.options.num_cpus = qutip.settings.num_cpus
    # set initial value data
    if options.tidy:
        odeconfig.psi0 = psi0.tidyup(options.atol).full().ravel()
    else:
        odeconfig.psi0 = psi0.full()
    odeconfig.psi0_dims = psi0.dims
    odeconfig.psi0_shape = psi0.shape
    # set general items
    odeconfig.tlist = tlist
    if isinstance(ntraj, (list, ndarray)):
        odeconfig.ntraj = sort(ntraj)[-1]
    else:
        odeconfig.ntraj = ntraj
    # set norm finding constants
    odeconfig.norm_tol = options.norm_tol
    odeconfig.norm_steps = options.norm_steps
    #----

    #----------------------------------------------
    # SETUP ODE DATA IF NONE EXISTS OR NOT REUSING
    #----------------------------------------------
    if (not options.rhs_reuse) or (not odeconfig.tdfunc):
        # reset odeconfig collapse and time-dependence flags to default values
        odeconfig.soft_reset()

        # check for type of time-dependence (if any)
        time_type, h_stuff, c_stuff = _ode_checks(H, c_ops, 'mc')
        h_terms = len(h_stuff[0]) + len(h_stuff[1]) + len(h_stuff[2])
        c_terms = len(c_stuff[0]) + len(c_stuff[1]) + len(c_stuff[2])
        # set time_type for use in multiprocessing
        odeconfig.tflag = time_type

        #-Check for PyObjC on Mac platforms
        if sys.platform == 'darwin' and odeconfig.options.gui:
            try:
                import Foundation
            except:
                odeconfig.options.gui = False

        # check if running in iPython and using Cython compiling (then no GUI
        # to work around error)
        if odeconfig.options.gui and odeconfig.tflag in array([1, 10, 11]):
            try:
                __IPYTHON__
            except:
                pass
            else:
                odeconfig.options.gui = False
        if qutip.settings.qutip_gui == "NONE":
            odeconfig.options.gui = False

        # check for collapse operators
        if c_terms > 0:
            odeconfig.cflag = 1
        else:
            odeconfig.cflag = 0

        # Configure data
        _mc_data_config(H, psi0, h_stuff, c_ops, c_stuff, args, e_ops,
                        options, odeconfig)

        # compile and load cython functions if necessary
        _mc_func_load(odeconfig)

    else:
        # setup args for new parameters when rhs_reuse=True and tdfunc is given
        # string based
        if odeconfig.tflag in array([1, 10, 11]):
            if any(args):
                odeconfig.c_args = []
                arg_items = args.items()
                for k in range(len(args)):
                    odeconfig.c_args.append(arg_items[k][1])
        # function based
        elif odeconfig.tflag in array([2, 3, 20, 22]):
            odeconfig.h_func_args = args

    # load monte-carlo class
    mc = _MC_class(odeconfig)

    # RUN THE SIMULATION
    mc.run()

    # remove RHS cython file if necessary
    if not options.rhs_reuse and odeconfig.tdname:
        try:
            os.remove(odeconfig.tdname + ".pyx")
        except Exc