# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
"""
This module provides solvers for the unitary Schrodinger equation.
"""

__all__ = ['sesolve']

import os
import types
from functools import partial
import numpy as np
import scipy.integrate
from scipy.linalg import norm as la_norm
import qutip.settings as qset
from qutip.qobj import Qobj
from qutip.rhs_generate import rhs_generate
from qutip.solver import Result, Options, config, _solver_safety_check
from qutip.rhs_generate import _td_format_check, _td_wrap_array_str
from qutip.interpolate import Cubic_Spline
from qutip.superoperator import vec2mat
from qutip.settings import debug
from qutip.cy.spmatfuncs import (cy_expect_psi, cy_ode_rhs,
                                 cy_ode_psi_func_td,
                                 cy_ode_psi_func_td_with_state,
                                 spmvpy_csr)
from qutip.cy.codegen import Codegen
from qutip.cy.utilities import _cython_build_cleanup

from qutip.ui.progressbar import (BaseProgressBar, TextProgressBar)
from qutip.cy.openmp.utilities import check_use_openmp, openmp_components

if qset.has_openmp:
    from qutip.cy.openmp.parfuncs import cy_ode_rhs_openmp

if debug:
    import inspect


def sesolve(H, psi0, tlist, e_ops=[], args={}, options=None,
            progress_bar=None,
            _safe_mode=True):
    """
    Schrodinger equation evolution of a state vector or unitary matrix
    for a given Hamiltonian.

    Evolve the state vector (`psi0`) using a given
    Hamiltonian (`H`), by integrating the set of ordinary differential
    equations that define the system. Alternatively evolve a unitary matrix in
    solving the Schrodinger operator equation.

    The output is either the state vector or unitary matrix at arbitrary points
    in time (`tlist`), or the expectation values of the supplied operators
    (`e_ops`). If e_ops is a callback function, it is invoked for each
    time in `tlist` with time and the state as arguments, and the function
    does not use any return values. e_ops cannot be used in conjunction
    with solving the Schrodinger operator equation

    Parameters
    ----------

    H : :class:`qutip.qobj`
        system Hamiltonian, or a callback function for time-dependent
        Hamiltonians.

    psi0 : :class:`qutip.qobj`
        initial state vector (ket)
        or initial unitary operator `psi0 = U`

    tlist : *list* / *array*
        list of times for :math:`t`.

    e_ops : list of :class:`qutip.qobj` / callback function single
        single operator or list of operators for which to evaluate
        expectation values.
        Must be empty list operator evolution

    args : *dictionary*
        dictionary of parameters for time-dependent Hamiltonians

    options : :class:`qutip.Qdeoptions`
        with options for the ODE solver.

    progress_bar : BaseProgressBar
        Optional instance of BaseProgressBar, or a subclass thereof, for
        showing the progress of the simulation.

    Returns
    -------

    output: :class:`qutip.solver`

        An instance of the class :class:`qutip.solver`, which contains either
        an *array* of expectation values for the times specified by `tlist`, or
        an *array* or state vectors corresponding to the
        times in `tlist` [if `e_ops` is an empty list], or
        nothing if a callback function was given inplace of operators for
        which to calculate the expectation values.

    """
    # check initial state: must be a state vector


    if _safe_mode:
        if not isinstance(psi0, Qobj):
            raise TypeError("psi0 must be Qobj")
        if psi0.isket:
            pass
        elif psi0.isunitary:
            if not e_ops == []:
                raise TypeError("Must have e_ops = [] when initial condition"
                                " psi0 is a unitary operator.")
        else:
            raise TypeError("The unitary solver requires psi0 to be"
                            " a ket as initial state"
                            " or a unitary as initial operator.")
        _solver_safety_check(H, psi0, c_ops=[], e_ops=e_ops, args=args)


    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

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

    # convert array based time-dependence to string format
    H, _, args = _td_wrap_array_str(H, [], args, tlist)
    # check for type (if any) of time-dependent inputs
    n_const, n_func, n_str = _td_format_check(H, [])

    if options is None:
        options = Options()

    if (not options.rhs_reuse) or (not config.tdfunc):
        # reset config time-dependence flags to default values
        config.reset()

    #check if should use OPENMP
    check_use_openmp(options)

    if n_func > 0:
        res = _sesolve_list_func_td(H, psi0, tlist, e_ops, args, options,
                                    progress_bar)

    elif n_str > 0:
        res = _sesolve_list_str_td(H, psi0, tlist, e_ops, args, options,
                                   progress_bar)

    elif isinstance(H, (types.FunctionType,
                        types.BuiltinFunctionType,
                        partial)):
        res = _sesolve_func_td(H, psi0, tlist, e_ops, args, options,
                               progress_bar)

    elif isinstance(H, Qobj):
        res = _sesolve_const(H, psi0, tlist, e_ops, args, options,
                             progress_bar)
    else:
        raise TypeError("Invalid Hamiltonian specification")

    if e_ops_dict:
        res.expect = {e: res.expect[n]
                      for n, e in enumerate(e_ops_dict.keys())}

    return res


# -----------------------------------------------------------------------------
# A time-dependent unitary wavefunction equation on the list-function format
#
def _sesolve_list_func_td(H_list, psi0, tlist, e_ops, args, opt,
                          progress_bar):
    """
    Internal function for solving the master equation. See mesolve for usage.
    """

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

    #
    # check initial state or oper
    #
    if psi0.isket:
        initial_vector = psi0.full().ravel()
        oper_evo = False
    elif psi0.isunitary:
        initial_vector = psi0.full().ravel('F')
        oper_evo = True
    else:
        raise TypeError("The unitary solver requires psi0 to be"
                        " a ket as initial state"
                        " or a unitary as initial operator.")

    #
    # construct liouvillian in list-function format
    #
    L_list = []
    if not opt.rhs_with_state:
        constant_func = lambda x, y: 1.0
    else:
        constant_func = lambda x, y, z: 1.0

    # add all hamitonian terms to the lagrangian list
    for h_spec in H_list:

        if isinstance(h_spec, Qobj):
            h = h_spec
            h_coeff = constant_func

        elif isinstance(h_spec, list):
            h = h_spec[0]
            h_coeff = h_spec[1]

        else:
            raise TypeError("Incorrect specification of time-dependent " +
                            "Hamiltonian (expected callback function)")

        L = -1j * h
        L_list.append([L.data, h_coeff])

    L_list_and_args = [L_list, args]

    #
    # setup integrator
    #
    if oper_evo:
        initial_vector = psi0.full().ravel('F')
        if opt.rhs_with_state:
            r = scipy.integrate.ode(oper_list_td_with_state)
        else:
            r = scipy.integrate.ode(oper_list_td)
    else:
        initial_vector = psi0.full().ravel()
        if opt.rhs_with_state:
            r = scipy.integrate.ode(psi_list_td_with_state)
        else:
            r = scipy.integrate.ode(psi_list_td)
    r.set_integrator('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)
    r.set_initial_value(initial_vector, tlist[0])
    r.set_f_params(L_list_and_args)

    #
    # call generic ODE code
    #
    return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar,
                              dims=psi0.dims)


#
# evaluate dpsi(t)/dt according to the master equation using the
# [Qobj, function] style time dependence API
#
def psi_list_td(t, psi, L_List_and_args):

    L_List = L_List_and_args[0]
    args = L_List_and_args[1]

    L = L_List[0][0]
    tdfunc = L_List[0][1]
    out = np.zeros(psi.shape[0],dtype=complex)
    spmvpy_csr(L.data, L.indices, L.indptr, psi, tdfunc(t, args), out)
    for n in range(1, len(L_List)):
        #
        # args[n][0] = the sparse data for a Qobj in operator form
        # args[n][1] = function callback giving the coefficient
        #
        L = L_List[n][0]
        tdfunc = L_List[n][1]
        spmvpy_csr(L.data, L.indices, L.indptr, psi, tdfunc(t, args), out)

    return out


def psi_list_td_with_state(t, psi, L_List_and_args):

    L_List = L_List_and_args[0]
    args = L_List_and_args[1]

    L = L_List[0][0]
    tdfunc = L_List[0][1]
    out = np.zeros(psi.shape[0],dtype=complex)
    spmvpy_csr(L.data, L.indices, L.indptr, psi, tdfunc(t, psi, args), out)
    for n in range(1, len(L_List)):
        #
        # args[n][0] = the sparse data for a Qobj in operator form
        # args[n][1] = function callback giving the coefficient
        #
        L = L_List[n][0]
        tdfunc = L_List[n][1]
        spmvpy_csr(L.data, L.indices, L.indptr, psi, tdfunc(t, psi, args), out)

    return out

#
# evaluate dU(t)/dt according to the master equation using the
# [Qobj, function] style time dependence API
#
def oper_list_td(t, y, L_List_and_args):
    L_List, args = L_List_and_args
    # L_List[n][0] = operator
    # L_List[n][1] = function callback giving the coefficient
    L = L_List[0][0] * L_List[0][1](t, args)
    for n in range(1, len(L_List)):
        L = L + L_List[n][0] * L_List[n][1](t, args)

    return _ode_oper_func(t, y, L)

def oper_list_td_with_state(t, y, L_List_and_args):
    L_List, args = L_List_and_args
    # L_List[n][0] = operator
    # L_List[n][1] = function callback giving the coefficient
    L = L_List[0][0] * L_List[0][1](t, y, args)
    for n in range(1, len(L_List)):
        L = L + L_List[n][0] * L_List[n][1](t, y, args)

    return _ode_oper_func(t, y, L)

# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution) using
# a constant Hamiltonian.
#
def _sesolve_const(H, psi0, tlist, e_ops, args, opt, progress_bar):
    """
    Evolve the wave function using an ODE solver
    """
    if debug:
        print(inspect.stack()[0][3])

    #
    # setup integrator.
    #

    if psi0.isket:
        initial_vector = psi0.full().ravel()
        oper_evo = False
    elif psi0.isunitary:
        initial_vector = psi0.full().ravel('F')
        oper_evo = True
    else:
        raise TypeError("The unitary solver requires psi0 to be"
                        " a ket as initial state"
                        " or a unitary as initial operator.")

    L = -1.0j * H
    if oper_evo:
        r = scipy.integrate.ode(_ode_oper_func)
        r.set_f_params(L.data)
    else:
        if opt.use_openmp and L.data.nnz >= qset.openmp_thresh:
            r = scipy.integrate.ode(cy_ode_rhs_openmp)
            r.set_f_params(L.data.data, L.data.indices, L.data.indptr,
                            opt.openmp_threads)
        else:
            r = scipy.integrate.ode(cy_ode_rhs)
            r.set_f_params(L.data.data, L.data.indices, L.data.indptr)
    r.set_integrator('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)

    r.set_initial_value(initial_vector, tlist[0])

    #
    # call generic ODE code
    #
    return _generic_ode_solve(r, psi0, tlist, e_ops, opt,
                              progress_bar, dims=psi0.dims)


#
# evaluate dpsi(t)/dt [not used. using cython function is being used instead]
#
def _ode_psi_func(t, psi, H):
    return H * psi

# TODO cythonize this?
def _ode_oper_func(t, y, data):
    ym = vec2mat(y)
    return (data*ym).ravel('F')

# -----------------------------------------------------------------------------
# A time-dependent disipative master equation on the list-string format for
# cython compilation
#
def _sesolve_list_str_td(H_list, psi0, tlist, e_ops, args, opt,
                         progress_bar):
    """
    Internal function for solving the master equation. See mesolve for usage.
    """

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

    if psi0.isket:
        oper_evo = False
    elif psi0.isunitary:
        oper_evo = True
    else:
        raise TypeError("The unitary solver requires psi0 to be"
                        " a ket as initial state"
                        " or a unitary as initial operator.")

    #
    # construct dynamics generator
    #
    Ldata = []
    Linds = []
    Lptrs = []
    Lcoeff = []
    Lobj = []

    # loop over all hamiltonian terms, convert to superoperator form and
    # add the data of sparse matrix representation to h_coeff
    for h_spec in H_list:

        if isinstance(h_spec, Qobj):
            h = h_spec
            h_coeff = "1.0"

        elif isinstance(h_spec, list):
            h = h_spec[0]
            h_coeff = h_spec[1]

        else:
            raise TypeError("Incorrect specification of time-dependent " +
                            "Hamiltonian (expected string format)")

        L = -1j * h
        Ldata.append(L.data.data)
        Linds.append(L.data.indices)
        Lptrs.append(L.data.indptr)
        if isinstance(h_coeff, Cubic_Spline):
            Lobj.append(h_coeff.coeffs)
        Lcoeff.append(h_coeff)

    # the total number of Hamiltonian terms
    n_L_terms = len(Ldata)

    # Check which components should use OPENMP
    omp_components = None
    if qset.has_openmp:
        if opt.use_openmp:
            omp_components = openmp_components(Lptrs)

    #
    # setup ode args string: we expand the list Ldata, Linds and Lptrs into
    # and explicit list of parameters
    #
    string_list = []
    for k in range(n_L_terms):
        string_list.append("Ldata[%d], Linds[%d], Lptrs[%d]" % (k, k, k))
    # Add object terms to end of ode args string
    for k in range(len(Lobj)):
        string_list.append("Lobj[%d]" % k)

    for name, value in args.items():
        if isinstance(value, np.ndarray):
            string_list.append(name)
        else:
            string_list.append(str(value))
    parameter_string = ",".join(string_list)

    #
    # generate and compile new cython code if necessary
    #
    if not opt.rhs_reuse or config.tdfunc is None:
        if opt.rhs_filename is None:
            config.tdname = "rhs" + str(os.getpid()) + str(config.cgen_num)
        else:
            config.tdname = opt.rhs_filename
        cgen = Codegen(h_terms=n_L_terms, h_tdterms=Lcoeff, args=args,
                       config=config, use_openmp=opt.use_openmp,
                       omp_components=omp_components,
                       omp_threads=opt.openmp_threads)
        cgen.generate(config.tdname + ".pyx")

        code = compile('from ' + config.tdname + ' import cy_td_ode_rhs',
                       '<string>', 'exec')
        exec(code, globals())
        config.tdfunc = cy_td_ode_rhs

    #
    # setup integrator
    #
    if oper_evo:
        initial_vector = psi0.full().ravel('F')
        r = scipy.integrate.ode(_td_ode_rhs_oper)
        code = compile('r.set_f_params([' + parameter_string + '])',
                       '<string>', 'exec')
    else:
        initial_vector = psi0.full().ravel()
        r = scipy.integrate.ode(config.tdfunc)
        code = compile('r.set_f_params(' + parameter_string + ')',
                       '<string>', 'exec')

    r.set_integrator('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)
    r.set_initial_value(initial_vector, tlist[0])

    exec(code, locals(), args)


    # Remove RHS cython file if necessary
    if not opt.rhs_reuse and config.tdname:
        _cython_build_cleanup(config.tdname)

    #
    # call generic ODE code
    #
    return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar,
                              dims=psi0.dims)

# TODO cythonize this?
def _td_ode_rhs_oper(t, y, arglist):
    N = int(np.sqrt(len(y)))
    out = np.zeros(N, dtype=complex)
    y2 = np.zeros(len(y), dtype=complex)
    for i in range(N):
        out = cy_td_ode_rhs(t, y[i*N:(i+1)*N], *arglist)
        y2[i*N:(i+1)*N] = out
    return y2

# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution), for
# time dependent Hamiltonians
def _sesolve_list_td(H_func, psi0, tlist, e_ops, args, opt, progress_bar):
    """
    Evolve the wave function using an ODE solver with time-dependent
    Hamiltonian.
    Note: Deprecated method
    """

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

    if psi0.isket:
        pass
    elif psi0.isunitary:
        raise TypeError("The unitary operator evolution is not supported"
                        " in the list td method.")
    else:
        raise TypeError("The unitary solver requires psi0 to be"
                        " a ket as initial state"
                        " or a unitary as initial operator.")

    #
    # configure time-dependent terms and setup ODE solver
    #
    if len(H_func) != 2:
        raise TypeError('Time-dependent Hamiltonian list must have two terms.')
    if (not isinstance(H_func[0], (list, np.ndarray))) or \
       (len(H_func[0]) <= 1):
        raise TypeError('Time-dependent Hamiltonians must be a list with two '
                        + 'or more terms')
    if (not isinstance(H_func[1], (list, np.ndarray))) or \
       (len(H_func[1]) != (len(H_func[0]) - 1)):
        raise TypeError('Time-dependent coefficients must be list with ' +
                        'length N-1 where N is the number of ' +
                        'Hamiltonian terms.')
    tflag = 1
    if opt.rhs_reuse and config.tdfunc is None:
        print("No previous time-dependent RHS found.")
        print("Generating one for you...")
        rhs_generate(H_func, args)
    lenh = len(H_func[0])
    if opt.tidy:
        H_func[0] = [(H_func[0][k]).tidyup() for k in range(lenh)]
    # create data arrays for time-dependent RHS function
    Hdata = [-1.0j * H_func[0][k].data.data for k in range(lenh)]
    Hinds = [H_func[0][k].data.indices for k in range(lenh)]
    Hptrs = [H_func[0][k].data.indptr for k in range(lenh)]
    # setup ode args string
    string = ""
    for k in range(lenh):
        string += ("Hdata[" + str(k) + "], Hinds[" + str(k) +
                   "], Hptrs[" + str(k) + "],")

    if args:
        td_consts = args.items()
        for elem in td_consts:
            string += str(elem[1])
            if elem != td_consts[-1]:
                string += (",")

    # run code generator
    if not opt.rhs_reuse or config.tdfunc is None:
        if opt.rhs_filename is None:
            config.tdname = "rhs" + str(os.getpid()) + str(config.cgen_num)
        else:
            config.tdname = opt.rhs_filename
        cgen = Codegen(h_terms=n_L_terms, h_tdterms=Lcoeff, args=args,
                       config=config)
        cgen.generate(config.tdname + ".pyx")

        code = compile('from ' + config.tdname + ' import cy_td_ode_rhs',
                       '<string>', 'exec')
        exec(code, globals())
        config.tdfunc = cy_td_ode_rhs
    #
    # setup integrator
    #
    initial_vector = psi0.full().ravel()
    r = scipy.integrate.ode(config.tdfunc)
    r.set_integrator('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)
    r.set_initial_value(initial_vector, tlist[0])
    code = compile('r.set_f_params(' + string + ')', '<string>', 'exec')
    exec(code)

    #
    # call generic ODE code
    #
    return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar
                            , dims=psi0.dims)


# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution), for
# time dependent hamiltonians.
#
def _sesolve_func_td(H_func, psi0, tlist, e_ops, args, opt, progress_bar):
    """!
    Evolve the wave function using an ODE solver with time-dependent
    Hamiltonian.
    """

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

    #
    # check initial state or oper
    #
    if psi0.isket:
        initial_vector = psi0.full().ravel()
        oper_evo = False
    elif psi0.isunitary:
        initial_vector = psi0.full().ravel('F')
        oper_evo = True
    else:
        raise TypeError("The unitary solver requires psi0 to be"
                        " a ket as initial state"
                        " or a unitary as initial operator.")

    #
    # setup integrator
    #
    new_args = None

    if type(args) is dict:
        new_args = {}
        for key in args:
            if isinstance(args[key], Qobj):
                new_args[key] = args[key].data
            else:
                new_args[key] = args[key]

    elif type(args) is list or type(args) is tuple:
        new_args = []
        for arg in args:
            if isinstance(arg, Qobj):
                new_args.append(arg.data)
            else:
                new_args.append(arg)

        if type(args) is tuple:
            new_args = tuple(new_args)
    else:
        if isinstance(args, Qobj):
            new_args = args.data
        else:
            new_args = args

    if oper_evo:
        initial_vector = psi0.full().ravel('F')
        if not opt.rhs_with_state:
            r = scipy.integrate.ode(_ode_oper_func_td)
        else:
            r = scipy.integrate.ode(_ode_oper_func_td_with_state)

    else:
        initial_vector = psi0.full().ravel()
        if not opt.rhs_with_state:
            r = scipy.integrate.ode(cy_ode_psi_func_td)
        else:
            r = scipy.integrate.ode(cy_ode_psi_func_td_with_state)

    r.set_integrator('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)
    r.set_initial_value(initial_vector, tlist[0])
    r.set_f_params(H_func, new_args)

    #
    # call generic ODE code
    #
    return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar,
                              dims=psi0.dims)


#
# evaluate dpsi(t)/dt for time-dependent hamiltonian
#
# TODO remove?: Looks like these are no longer in use
def _ode_psi_func_td(t, psi, H_func, args):
    H = H_func(t, args)
    return -1j * (H * psi)

def _ode_psi_func_td_with_state(t, psi, H_func, args):
    H = H_func(t, psi, args)
    return -1j * (H * psi)
#
# evaluate dU(t)/dt according to the master equation using the
#
# TODO cythonize these?
def _ode_oper_func_td(t, y, H_func, args):
    H = H_func(t, args)
    return -1j * _ode_oper_func(t, y, H.data)

def _ode_oper_func_td_with_state(t, y, H_func, args):
    H = H_func(t, y, args)
    return -1j * _ode_oper_func(t, y, H.data)

# -----------------------------------------------------------------------------
# Solve an ODE which solver parameters already setup (r). Calculate the
# required expectation values or invoke callback function at each time step.
#
def _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar, dims=None):
    """
    Internal function for solving ODEs.
    """

    #
    # prepare output array
    #
    n_tsteps = len(tlist)
    output = Result()
    output.solver = "sesolve"
    output.times = tlist

    if psi0.isunitary:
        oper_evo = True
        oper_n = dims[0][0]
        norm_dim_factor = np.sqrt(oper_n)
    else:
        oper_evo = False
        norm_dim_factor = 1.0

    if opt.store_states:
        output.states = []

    if isinstance(e_ops, types.FunctionType):
        n_expt_op = 0
        expt_callback = True

    elif isinstance(e_ops, list):
        n_expt_op = len(e_ops)
        expt_callback = False

        if n_expt_op == 0:
            # fallback on storing states
            output.states = []
            opt.store_states = True
        else:
            output.expect = []
            output.num_expect = n_expt_op
            for op in e_ops:
                if op.isherm:
                    output.expect.append(np.zeros(n_tsteps))
                else:
                    output.expect.append(np.zeros(n_tsteps, dtype=complex))
    else:
        raise TypeError("Expectation parameter must be a list or a function")

    def get_curr_state_data():
        if oper_evo:
            return vec2mat(r.y)
        else:
            return r.y

    #
    # start evolution
    #
    progress_bar.start(n_tsteps)

    dt = np.diff(tlist)
    for t_idx, t in enumerate(tlist):
        progress_bar.update(t_idx)

        if not r.successful():
            raise Exception("ODE integration error: Try to increase "
                            "the allowed number of substeps by increasing "
                            "the nsteps parameter in the Options class.")

        # get the current state / oper data if needed
        cdata = None
        if opt.store_states or opt.normalize_output or n_expt_op > 0:
            cdata = get_curr_state_data()

        if opt.normalize_output:
            # cdata *= _get_norm_factor(cdata, oper_evo)
            cdata *= norm_dim_factor / la_norm(cdata)
            if oper_evo:
                r.set_initial_value(cdata.ravel('F'), r.t)
            else:
                r.set_initial_value(cdata, r.t)

        if opt.store_states:
            output.states.append(Qobj(cdata, dims=dims))

        if expt_callback:
            # use callback method
            e_ops(t, Qobj(cdata, dims=dims))

        for m in range(n_expt_op):
            output.expect[m][t_idx] = cy_expect_psi(e_ops[m].data,
                                                    cdata, e_ops[m].isherm)

        if t_idx < n_tsteps - 1:
            r.integrate(r.t + dt[t_idx])

    progress_bar.finished()

    if not opt.rhs_reuse and config.tdname is not None:
        try:
            os.remove(config.tdname + ".pyx")
        except:
            pass

    if opt.store_final_state:
        cdata = get_curr_state_data()
        if opt.normalize_output:
            cdata *= norm_dim_factor / la_norm(cdata)
        output.final_state = Qobj(cdata, dims=dims)

    return output
