# -*- coding: utf-8 -*-
#
# 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.
###############################################################################
import numpy as np
import scipy.sparse as sp
from qutip.cy.stochastic import (SSESolver, SMESolver, PcSSESolver, PcSMESolver,
                                 PmSMESolver, GenericSSolver, Solvers)
from qutip.qobj import Qobj, isket, isoper, issuper
from qutip.states import ket2dm
from qutip.solver import Result
from qutip.qobjevo import QobjEvo
from qutip.superoperator import (spre, spost, mat2vec, vec2mat,
                                 liouvillian, lindblad_dissipator)
from qutip.solver import Options, _solver_safety_check
from qutip.parallel import serial_map
from qutip.ui.progressbar import TextProgressBar
from qutip.pdpsolve import main_ssepdpsolve, main_smepdpsolve

__all__ = ['ssesolve', 'photocurrent_sesolve', 'smepdpsolve',
           'smesolve', 'photocurrent_mesolve', 'ssepdpsolve',
           'stochastic_solvers', 'general_stochastic']

def stochastic_solvers():
    """Available solvers for ssesolve and smesolve
    euler-maruyama:
        A simple generalization of the Euler method for ordinary
        differential equations to stochastic differential equations.
        Only solver which could take non-commuting sc_ops. *not tested*
        -Order 0.5
        -Code: 'euler-maruyama', 'euler', 0.5

    milstein, Order 1.0 strong Taylor scheme:
        Better approximate numerical solution to stochastic
        differential equations.
        -Order strong 1.0
        -Code: 'milstein', 1.0
        Numerical Solution of Stochastic Differential Equations
        Chapter 10.3 Eq. (3.1), By Peter E. Kloeden, Eckhard Platen

    milstein-imp, Order 1.0 implicit strong Taylor scheme:
        Implicit milstein scheme for the numerical simulation of stiff
        stochastic differential equations.
        -Order strong 1.0
        -Code: 'milstein-imp'
        Numerical Solution of Stochastic Differential Equations
        Chapter 12.2 Eq. (2.9), By Peter E. Kloeden, Eckhard Platen

    predictor-corrector:
        Generalization of the trapezoidal method to stochastic
        differential equations. More stable than explicit methods.
        -Order strong 0.5, weak 1.0
        Only the stochastic part is corrected.
            (alpha = 0, eta = 1/2)
            -Code: 'pred-corr', 'predictor-corrector', 'pc-euler'
        Both the deterministic and stochastic part corrected.
            (alpha = 1/2, eta = 1/2)
            -Code: 'pc-euler-imp', 'pc-euler-2', 'pred-corr-2'
        Numerical Solution of Stochastic Differential Equations
        Chapter 15.5 Eq. (5.4), By Peter E. Kloeden, Eckhard Platen

    platen:
        Explicit scheme, create the milstein using finite difference instead of
        derivatives. Also contain some higher order terms, thus converge better
        than milstein while staying strong order 1.0.
        Do not require derivatives, therefore usable for
        :func:`qutip.stochastic.general_stochastic`
        -Order strong 1.0, weak 2.0
        -Code: 'platen', 'platen1', 'explicit1'
        The Theory of Open Quantum Systems
        Chapter 7 Eq. (7.47), H.-P Breuer, F. Petruccione

    rouchon:
        Scheme keeping the positivity of the density matrix. (smesolve only)
        -Order strong 1.0?
        -Code: 'rouchon', 'Rouchon'
        Eq. 4 of arXiv:1410.5345 with eta=1
        Efficient Quantum Filtering for Quantum Feedback Control
        Pierre Rouchon, Jason F. Ralph
        arXiv:1410.5345 [quant-ph]
        Phys. Rev. A 91, 012118, (2015)

    taylor1.5, Order 1.5 strong Taylor scheme:
        Solver with more terms of the Ito-Taylor expansion.
        Default solver for smesolve and ssesolve.
        -Order strong 1.5
        -Code: 'taylor1.5', 'taylor15', 1.5, None
        Numerical Solution of Stochastic Differential Equations
        Chapter 10.4 Eq. (4.6), By Peter E. Kloeden, Eckhard Platen

    taylor1.5-imp, Order 1.5 implicit strong Taylor scheme:
        implicit Taylor 1.5 (alpha = 1/2, beta = doesn't matter)
        -Order strong 1.5
        -Code: 'taylor1.5-imp', 'taylor15-imp'
        Numerical Solution of Stochastic Differential Equations
        Chapter 12.2 Eq. (2.18), By Peter E. Kloeden, Eckhard Platen

    explicit1.5, Explicit Order 1.5 Strong Schemes:
        Reproduce the order 1.5 strong Taylor scheme using finite difference
        instead of derivatives. Slower than taylor15 but usable by
        :func:`qutip.stochastic.general_stochastic`
        -Order strong 1.5
        -Code: 'explicit1.5', 'explicit15', 'platen15'
        Numerical Solution of Stochastic Differential Equations
        Chapter 11.2 Eq. (2.13), By Peter E. Kloeden, Eckhard Platen

    taylor2.0, Order 2 strong Taylor scheme:
        Solver with more terms of the Stratonovich expansion.
        -Order strong 2.0
        -Code: 'taylor2.0', 'taylor20', 2.0
        Numerical Solution of Stochastic Differential Equations
        Chapter 10.5 Eq. (5.2), By Peter E. Kloeden, Eckhard Platen

    ---All solvers, except taylor2.0, are usable in both smesolve and ssesolve
    and for both heterodyne and homodyne. taylor2.0 only work for 1 stochastic
    operator not dependent of time with the homodyne method.
    The :func:`qutip.stochastic.general_stochastic` only accept derivatives
    free solvers: ['euler', 'platen', 'explicit1.5'].

Available solver for photocurrent_sesolve and photocurrent_mesolve:
        Photocurrent use ordinary differential equations between
        stochastic "jump/collapse".
    euler:
        Euler method for ordinary differential equations between jumps.
        Only 1 jumps per time interval.
        Default solver
        -Order 1.0
        -Code: 'euler'
        Quantum measurement and control
        Chapter 4, Eq 4.19, 4.40, By Howard M. Wiseman, Gerard J. Milburn

    predictor–corrector:
        predictor–corrector method (PECE) for ordinary differential equations.
        Use poisson distribution to obtain the number of jump at each timestep.
        -Order 2.0
        -Code: 'pred-corr'

    """
    pass

class StochasticSolverOptions:
    """Class of options for stochastic solvers such as
    :func:`qutip.stochastic.ssesolve`, :func:`qutip.stochastic.smesolve`, etc.

    The stochastic solvers :func:`qutip.stochastic.general_stochastic`,
    :func:`qutip.stochastic.ssesolve`, :func:`qutip.stochastic.smesolve`,
    :func:`qutip.stochastic.photocurrent_sesolve` and
    :func:`qutip.stochastic.photocurrent_mesolve`
    all take the same keyword arguments as
    the constructor of these class, and internally they use these arguments to
    construct an instance of this class, so it is rarely needed to explicitly
    create an instance of this class.

    Attributes
    ----------

    H : :class:`qutip.Qobj`, time-dependent Qobj as a list*
        System Hamiltonian.

    state0 : :class:`qutip.Qobj`
        Initial state vector (ket) or density matrix.

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    c_ops : list of :class:`qutip.Qobj`, :class:`qutip.QobjEvo` or [Qobj, coeff*]
        List of deterministic collapse operators.

    sc_ops : list of :class:`qutip.Qobj`, :class:`qutip.QobjEvo` or [Qobj, coeff*]
        List of stochastic collapse operators. Each stochastic collapse
        operator will give a deterministic and stochastic contribution
        to the equation of motion according to how the d1 and d2 functions
        are defined.

    e_ops : list of :class:`qutip.Qobj`
        Single operator or list of operators for which to evaluate
        expectation values.

    m_ops : list of :class:`qutip.Qobj`
        List of operators representing the measurement operators. The expected
        format is a nested list with one measurement operator for each
        stochastic increament, for each stochastic collapse operator.

    args : dict
        Dictionary of parameters for time dependent systems.

    tol : float
        Tolerance of the solver for implicit methods.

    ntraj : int
        Number of trajectors.

    nsubsteps : int
        Number of sub steps between each time-spep given in `times`.

    dW_factors : array
        Array of length len(sc_ops), containing scaling factors for each
        measurement operator in m_ops.

    solver : string
        Name of the solver method to use for solving the stochastic
        equations. Valid values are:
        order 1/2 algorithms: 'euler-maruyama', 'pc-euler', 'pc-euler-imp'
        order 1 algorithms: 'milstein', 'platen', 'milstein-imp', 'rouchon'
        order 3/2 algorithms: 'taylor1.5', 'taylor1.5-imp', 'explicit1.5'
        order 2 algorithms: 'taylor2.0'
        call help of :func:`qutip.stochastic.stochastic_solvers`
        for a description of the solvers.
        Implicit methods can adjust tolerance via the kw 'tol'
        default is {'tol':1e-6}

    method : string ('homodyne', 'heterodyne')
        The name of the type of measurement process that give rise to the
        stochastic equation to solve.

    store_all_expect : bool (default False)
        Whether or not to store the e_ops expect values for all paths.

    store_measurement : bool (default False)
        Whether or not to store the measurement results in the
        :class:`qutip.solver.Result` instance returned by the solver.

    noise : int, array[int, 1d], array[double, 4d]
        int : seed of the noise
        array[int, 1d], length = ntraj, seeds for each trajectories
        array[double, 4d] (ntraj, len(times), nsubsteps, len(sc_ops)*[1|2])
            vector for the noise, the len of the last dimensions is doubled for
            solvers of order 1.5. The correspond to results.noise

    noiseDepth : int
        Number of terms kept of the truncated series used to create the
        noise used by taylor2.0 solver.

    normalize : bool
        (default True for (photo)ssesolve, False for (photo)smesolve)
        Whether or not to normalize the wave function during the evolution.
        Normalizing density matrices introduce numerical errors.

    options : :class:`qutip.solver.Options`
        Generic solver options. Only options.average_states and
        options.store_states are used.

    map_func: function
        A map function or managing the calls to single-trajactory solvers.

    map_kwargs: dictionary
        Optional keyword arguments to the map_func function function.

    progress_bar : :class:`qutip.ui.BaseProgressBar`
        Optional progress bar class instance.

    *
    time-dependent Qobj can be used for H, c_ops and sc_ops.
    The format for time-dependent system hamiltonian is:
    H = [Qobj0,[Qobj1,coeff1],[Qobj2,coeff2],...]
      = Qobj0 + Qobj1 * coeff1(t) + Qobj2 * coeff2(t)

    coeff function can be:
        function: coeff(t, args) -> complex
        str: "sin(1j*w*t)"
        np.array[complex, 1d] of length equal to the times array
    The argument args for the function coeff is the args keyword argument of
        the stochastic solver.
    Likewisem in str cases, the parameters ('w' in this case) are taken from
        the args keywords argument.
    *While mixing coeff type does not results in errors, it is not recommended.*

    For the collapse operators (c_ops, sc_ops):
    Each operators can only be composed of 1 Qobj.
    c_ops = [c_op1, c_op2, ...]
    where, c_opN = Qobj or [Qobj,coeff]
    The coeff format is the same as for the Hamiltonian.
    """
    def __init__(self, me, H=None, c_ops=[], sc_ops=[], state0=None,
                 e_ops=[], m_ops=None, store_all_expect=False,
                 store_measurement=False, dW_factors=None,
                 solver=None, method="homodyne", normalize=None,
                 times=None, nsubsteps=1, ntraj=1, tol=None,
                 generate_noise=None, noise=None,
                 progress_bar=None, map_func=None, map_kwargs=None,
                 args={}, options=None, noiseDepth=20):

        if options is None:
            options = Options()

        if progress_bar is None:
            progress_bar = TextProgressBar()

        # System
        # Cast to QobjEvo so the code has only one version for both the
        # constant and time-dependent case.
        self.me = me

        if H is not None:
            msg = "The Hamiltonian format is not valid. "
            try:
                self.H = QobjEvo(H, args=args, tlist=times)
            except Exception as e:
                raise ValueError(msg + str(e)) from e
            except:
                raise ValueError(msg)
        else:
            self.H = H

        if sc_ops:
            msg = ("The sc_ops format is not valid. Options are "
                   "[ Qobj / QobjEvo / [Qobj, coeff]]. ")
            try:
                self.sc_ops = [QobjEvo(op, args=args, tlist=times)
                               for op in sc_ops]
            except Exception as e:
                raise ValueError(msg + str(e)) from e
            except:
                raise ValueError(msg)
        else:
            self.sc_ops = sc_ops

        if c_ops:
            msg = ("The c_ops format is not valid. Options are "
                   "[ Qobj / QobjEvo / [Qobj, coeff]]. ")
            try:
                self.c_ops = [QobjEvo(op, args=args, tlist=times)
                              for op in c_ops]
            except Exception as e:
                raise ValueError(msg + str(e)) from e
            except:
                raise ValueError(msg)
        else:
            self.c_ops = c_ops

        self.state0 = state0
        self.rho0 = mat2vec(state0.full()).ravel()

        # Observation
        self.e_ops = e_ops
        self.m_ops = m_ops
        self.store_measurement = store_measurement
        self.store_all_expect = store_all_expect
        self.store_states = options.store_states
        self.dW_factors = dW_factors

        # Solver
        self.solver = solver
        self.method = method
        if normalize is None and me:
            self.normalize = 0
        elif normalize is None and not me:
            self.normalize = 1
        elif normalize:
            self.normalize = 1
        else:
            self.normalize = 0

        self.times = times
        self.nsubsteps = nsubsteps
        self.dt = (times[1] - times[0]) / self.nsubsteps
        self.ntraj = ntraj
        if tol is not None:
            self.tol = tol
        elif "tol" in args:
            self.tol = args["tol"]
        else:
            self.tol = 1e-7

        # Noise
        if noise is not None:
            if isinstance(noise, int):
                # noise contain a seed
                np.random.seed(noise)
                noise = np.random.randint(0, 2**32, ntraj)
            noise = np.array(noise)
            if len(noise.shape) == 1:
                if noise.shape[0] < ntraj:
                    raise ValueError("'noise' does not have enought seeds " +
                                     "len(noise) >= ntraj")
                # numpy seed must be between 0 and 2**32-1
                # 'u4': unsigned 32bit int
                self.noise = noise.astype("u4")
                self.noise_type = 0

            elif len(noise.shape) == 4:
                # taylor case not included
                dw_len = (2 if method == "heterodyne" else 1)
                dw_len_str = (" * 2" if method == "heterodyne" else "")
                msg = "Incorrect shape for 'noise': "
                if noise.shape[0] < ntraj:
                    raise ValueError(msg + "shape[0] >= ntraj")
                if noise.shape[1] < len(times):
                    raise ValueError(msg + "shape[1] >= len(times)")
                if noise.shape[2] < nsubsteps:
                    raise ValueError(msg + "shape[2] >= nsubsteps")
                if noise.shape[3] < len(self.sc_ops) * dw_len:
                    raise ValueError(msg + "shape[3] >= len(self.sc_ops)" +
                                     dw_len_str)
                self.noise_type = 1
                self.noise = noise

        else:
            self.noise = np.random.randint(0, 2**32, ntraj).astype("u4")
            self.noise_type = 0

        # Map
        self.progress_bar = progress_bar
        if self.ntraj > 1 and map_func:
            self.map_func = map_func
        else:
            self.map_func = serial_map
        self.map_kwargs = map_kwargs if map_kwargs is not None else {}

        # Other
        self.options = options
        self.args = args
        self.set_solver()
        self.p = noiseDepth

    def set_solver(self):
        if self.solver in ['euler-maruyama', 'euler', 50, 0.5]:
            self.solver_code = 50
            self.solver = 'euler-maruyama'
        elif self.solver in ['platen', 'platen1', 'explicit1', 100]:
            self.solver_code = 100
            self.solver = 'platen'
        elif self.solver in ['pred-corr', 'predictor-corrector',
                             'pc-euler', 101]:
            self.solver_code = 101
            self.solver = 'pred-corr'
        elif self.solver in ['milstein', 102, 1.0]:
            self.solver_code = 102
            self.solver = 'milstein'
        elif self.solver in ['milstein-imp', 103]:
            self.solver_code = 103
            self.solver = 'milstein-imp'
        elif self.solver in ['pred-corr-2', 'pc-euler-2', 'pc-euler-imp', 104]:
            self.solver_code = 104
            self.solver = 'pred-corr-2'
        elif self.solver in ['Rouchon', 'rouchon', 120]:
            self.solver_code = 120
            self.solver = 'rouchon'
            if not all((op.const for op in self.sc_ops)):
                raise ValueError("Rouchon only works with constant sc_ops")
        elif self.solver in ['platen15', 'explicit1.5', 'explicit15', 150]:
            self.solver_code = 150
            self.solver = 'explicit1.5'
        elif self.solver in ['taylor15', 'taylor1.5', None, 1.5, 152]:
            self.solver_code = 152
            self.solver = 'taylor1.5'
        elif self.solver in ['taylor15-imp', 'taylor1.5-imp', 153]:
            self.solver_code = 153
            self.solver = 'taylor1.5-imp'
        elif self.solver in ['taylor2.0', 'taylor20', 2.0, 202]:
            self.solver_code = 202
            self.solver = 'taylor2.0'
            if not len(self.sc_ops) == 1 or \
                    not self.sc_ops[0].const or \
                    not self.method == "homodyne":
                raise ValueError("Taylor2.0 only works with 1 constant " +
                                "sc_ops and for homodyne method")
        else:
            raise ValueError((
                    "The solver should be one of "
                    "[None, 'euler-maruyama', 'platen', 'pc-euler', "
                    "'pc-euler-imp', 'milstein', 'milstein-imp', "
                    "'rouchon', "
                    "'taylor1.5', 'taylor1.5-imp', 'explicit1.5' "
                    "'taylor2.0']"))

class StochasticSolverOptionsPhoto(StochasticSolverOptions):
    """
    Attributes
    ----------

    solver : string
        Name of the solver method to use for solving the evolution
        of the system.*
        order 1 algorithms: 'euler'
        order 2 algorithms: 'pred-corr'
        In photocurrent evolution
    """
    def set_solver(self):
        if self.solver in [None, 'euler', 1, 60]:
            self.solver_code = 60
            self.solver = 'euler'
        elif self.solver in ['pred-corr', 'predictor-corrector', 110, 2]:
            self.solver_code = 110
            self.solver = 'pred-corr'
        else:
            raise Exception("The solver should be one of " +
                            "[None, 'euler', 'predictor-corrector']")

def smesolve(H, rho0, times, c_ops=[], sc_ops=[], e_ops=[],
             _safe_mode=True, args={}, **kwargs):
    """
    Solve stochastic master equation. Dispatch to specific solvers
    depending on the value of the `solver` keyword argument.

    Parameters
    ----------

    H : :class:`qutip.Qobj`, or time dependent system.
        System Hamiltonian.
        Can depend on time, see StochasticSolverOptions help for format.

    rho0 : :class:`qutip.Qobj`
        Initial density matrix or state vector (ket).

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    c_ops : list of :class:`qutip.Qobj`, or time dependent Qobjs.
        Deterministic collapse operator which will contribute with a standard
        Lindblad type of dissipation.
        Can depend on time, see StochasticSolverOptions help for format.

    sc_ops : list of :class:`qutip.Qobj`, or time dependent Qobjs.
        List of stochastic collapse operators. Each stochastic collapse
        operator will give a deterministic and stochastic contribution
        to the eqaution of motion according to how the d1 and d2 functions
        are defined.
        Can depend on time, see StochasticSolverOptions help for format.

    e_ops : list of :class:`qutip.Qobj`
        single operator or list of operators for which to evaluate
        expectation values.

    kwargs : *dictionary*
        Optional keyword arguments. See
        :class:`qutip.stochastic.StochasticSolverOptions`.

    Returns
    -------

    output: :class:`qutip.solver.Result`

        An instance of the class :class:`qutip.solver.Result`.

    """
    if "method" in kwargs and kwargs["method"] == "photocurrent":
        print("stochastic solver with photocurrent method has been moved to "
              "it's own function: photocurrent_mesolve")
        return photocurrent_mesolve(H, rho0, times, c_ops=c_ops, sc_ops=sc_ops,
                                   e_ops=e_ops, _safe_mode=_safe_mode,
                                   args=args, **kwargs)
    if isket(rho0):
        rho0 = ket2dm(rho0)

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

    sso = StochasticSolverOptions(True, H=H, state0=rho0, times=times,
                                  c_ops=c_ops, sc_ops=sc_ops, e_ops=e_ops,
                                  args=args, **kwargs)

    if _safe_mode:
        _safety_checks(sso)

    if sso.solver_code == 120:
        return _positive_map(sso, e_ops_dict)

    sso.LH = liouvillian(sso.H, c_ops=sso.sc_ops + sso.c_ops) * sso.dt
    if sso.method == 'homodyne' or sso.method is None:
        if sso.m_ops is None:
            sso.m_ops = [op + op.dag() for op in sso.sc_ops]
        sso.sops = [spre(op) + spost(op.dag()) for op in sso.sc_ops]
        if not isinstance(sso.dW_factors, list):
            sso.dW_factors = [1] * len(sso.m_ops)
        elif len(sso.dW_factors) != len(sso.m_ops):
            raise Exception("The len of dW_factors is not the same as m_ops")

    elif sso.method == 'heterodyne':
        if sso.m_ops is None:
            m_ops = []
        sso.sops = []
        for c in sso.sc_ops:
            if sso.m_ops is None:
                m_ops += [c + c.dag(), -1j * c - c.dag()]
            sso.sops += [(spre(c) + spost(c.dag())) / np.sqrt(2),
                         (spre(c) - spost(c.dag())) * -1j / np.sqrt(2)]
        sso.m_ops = m_ops
        if not isinstance(sso.dW_factors, list):
            sso.dW_factors = [np.sqrt(2)] * len(sso.sops)
        elif len(sso.dW_factors) == len(sso.m_ops):
            pass
        elif len(sso.dW_factors) == len(sso.sc_ops):
            dW_factors = []
            for fact in sso.dW_factors:
                dW_factors += [np.sqrt(2) * fact, np.sqrt(2) * fact]
            sso.dW_factors = dW_factors
        elif len(sso.dW_factors) != len(sso.m_ops):
            raise Exception("The len of dW_factors is not the same as sc_ops")

    elif sso.method == "photocurrent":
        raise NotImplementedError("Moved to 'photocurrent_mesolve'")

    else:
        raise Exception("The method must be one of None, homodyne, heterodyne")

    sso.ce_ops = [QobjEvo(spre(op)) for op in sso.e_ops]
    sso.cm_ops = [QobjEvo(spre(op)) for op in sso.m_ops]

    sso.LH.compile()
    [op.compile() for op in sso.sops]
    [op.compile() for op in sso.cm_ops]
    [op.compile() for op in sso.ce_ops]

    if sso.solver_code in [103, 153]:
        sso.imp = 1 - sso.LH * 0.5
        sso.imp.compile()

    sso.solver_obj = SMESolver
    sso.solver_name = "smesolve_" + sso.solver

    res = _sesolve_generic(sso, sso.options, sso.progress_bar)

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

def ssesolve(H, psi0, times, sc_ops=[], e_ops=[],
             _safe_mode=True, args={}, **kwargs):
    """
    Solve stochastic schrodinger equation. Dispatch to specific solvers
    depending on the value of the `solver` keyword argument.

    Parameters
    ----------

    H : :class:`qutip.Qobj`, or time dependent system.
        System Hamiltonian.
        Can depend on time, see StochasticSolverOptions help for format.

    psi0 : :class:`qutip.Qobj`
        State vector (ket).

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    sc_ops : list of :class:`qutip.Qobj`, or time dependent Qobjs.
        List of stochastic collapse operators. Each stochastic collapse
        operator will give a deterministic and stochastic contribution
        to the eqaution of motion according to how the d1 and d2 functions
        are defined.
        Can depend on time, see StochasticSolverOptions help for format.

    e_ops : list of :class:`qutip.Qobj`
        single operator or list of operators for which to evaluate
        expectation values.

    kwargs : *dictionary*
        Optional keyword arguments. See
        :class:`qutip.stochastic.StochasticSolverOptions`.

    Returns
    -------

    output: :class:`qutip.solver.Result`

        An instance of the class :class:`qutip.solver.Result`.
    """
    if "method" in kwargs and kwargs["method"] == "photocurrent":
        print("stochastic solver with photocurrent method has been moved to "
              "it's own function: photocurrent_sesolve")
        return photocurrent_sesolve(H, psi0, times, c_ops=c_ops,
                                   e_ops=e_ops, _safe_mode=_safe_mode,
                                   args=args, **kwargs)

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

    sso = StochasticSolverOptions(False, H=H, state0=psi0, times=times,
                                  sc_ops=sc_ops, e_ops=e_ops,
                                  args=args, **kwargs)

    if _safe_mode:
        _safety_checks(sso)

    if sso.solver_code == 120:
        raise Exception("rouchon only work with smesolve")

    if sso.method == 'homodyne' or sso.method is None:
        if sso.m_ops is None:
            sso.m_ops = [op + op.dag() for op in sso.sc_ops]
        sso.sops = [[op, op + op.dag()] for op in sso.sc_ops]
        if not isinstance(sso.dW_factors, list):
            sso.dW_factors = [1] * len(sso.sops)
        elif len(sso.dW_factors) != len(sso.sops):
            raise Exception("The len of dW_factors is not the same as sc_ops")

    elif sso.method == 'heterodyne':
        if sso.m_ops is None:
            m_ops = []
        sso.sops = []
        for c in sso.sc_ops:
            if sso.m_ops is None:
                m_ops += [c + c.dag(), -1j * (c - c.dag())]
            c1 = c / np.sqrt(2)
            c2 = c * (-1j / np.sqrt(2))
            sso.sops += [[c1, c1 + c1.dag()],
                         [c2, c2 + c2.dag()]]
        sso.m_ops = m_ops
        if not isinstance(sso.dW_factors, list):
            sso.dW_factors = [np.sqrt(2)] * len(sso.sops)
        elif len(sso.dW_factors) == len(sso.sc_ops):
            dW_factors = []
            for fact in sso.dW_factors:
                dW_factors += [np.sqrt(2) * fact, np.sqrt(2) * fact]
            sso.dW_factors = dW_factors
        elif len(sso.dW_factors) != len(sso.sops):
            raise Exception("The len of dW_factors is not the same as sc_ops")

    elif sso.method == "photocurrent":
        NotImplementedError("Moved to 'photocurrent_sesolve'")

    else:
        raise Exception("The method must be one of None, homodyne, heterodyne")

    sso.LH = sso.H * (-1j*sso.dt)
    for ops in sso.sops:
        sso.LH -= ops[0]._cdc()*0.5*sso.dt

    sso.ce_ops = [QobjEvo(op) for op in sso.e_ops]
    sso.cm_ops = [QobjEvo(op) for op in sso.m_ops]

    sso.LH.compile()
    [[op.compile() for op in ops] for ops in sso.sops]
    [op.compile() for op in sso.cm_ops]
    [op.compile() for op in sso.ce_ops]

    sso.solver_obj = SSESolver
    sso.solver_name = "ssesolve_" + sso.solver

    res = _sesolve_generic(sso, sso.options, sso.progress_bar)

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

    return res

def _positive_map(sso, e_ops_dict):
    if sso.method == 'homodyne' or sso.method is None:
        sops = sso.sc_ops
        if sso.m_ops is None:
            sso.m_ops = [op + op.dag() for op in sso.sc_ops]
        if not isinstance(sso.dW_factors, list):
            sso.dW_factors = [1] * len(sops)
        elif len(sso.dW_factors) != len(sops):
            raise Exception("The len of dW_factors is not the same as sc_ops")

    elif sso.method == 'heterodyne':
        if sso.m_ops is None:
            m_ops = []
        sops = []
        for c in sso.sc_ops:
            if sso.m_ops is None:
                m_ops += [c + c.dag(), -1j * c - c.dag()]
            sops += [c / np.sqrt(2), -1j / np.sqrt(2) * c]
        sso.m_ops = m_ops
        if not isinstance(sso.dW_factors, list):
            sso.dW_factors = [np.sqrt(2)] * len(sops)
        elif len(sso.dW_factors) == len(sso.sc_ops):
            dW_factors = []
            for fact in sso.dW_factors:
                dW_factors += [np.sqrt(2) * fact, np.sqrt(2) * fact]
            sso.dW_factors = dW_factors
        elif len(sso.dW_factors) != len(sops):
            raise Exception("The len of dW_factors is not the same as sc_ops")
    else:
        raise Exception("The method must be one of homodyne or heterodyne")

    LH = 1 - (sso.H * 1j * sso.dt)
    sso.pp = spre(sso.H) * 0
    sso.sops = []
    sso.preops = []
    sso.postops = []
    sso.preops2 = []
    sso.postops2 = []

    def _prespostdag(op):
        return spre(op) * spost(op.dag())

    for op in sso.c_ops:
        LH -= op._cdc() * sso.dt * 0.5
        sso.pp += op.apply(_prespostdag)._f_norm2() * sso.dt

    for i, op in enumerate(sops):
        LH -= op._cdc() * sso.dt * 0.5
        sso.sops += [(spre(op) + spost(op.dag())) * sso.dt]
        sso.preops += [spre(op)]
        sso.postops += [spost(op.dag())]
        for op2 in sops[i:]:
            sso.preops2 += [spre(op * op2)]
            sso.postops2 += [spost(op.dag() * op2.dag())]

    sso.ce_ops = [QobjEvo(spre(op)) for op in sso.e_ops]
    sso.cm_ops = [QobjEvo(spre(op)) for op in sso.m_ops]
    sso.preLH = spre(LH)
    sso.postLH = spost(LH.dag())
    sso.preLH.compile()
    sso.postLH.compile()
    sso.pp.compile()
    [op.compile() for op in sso.sops]
    [op.compile() for op in sso.preops]
    [op.compile() for op in sso.postops]
    [op.compile() for op in sso.preops2]
    [op.compile() for op in sso.postops2]
    [op.compile() for op in sso.cm_ops]
    [op.compile() for op in sso.ce_ops]

    sso.solver_obj = PmSMESolver
    sso.solver_name = "smesolve_" + sso.solver
    res = _sesolve_generic(sso, sso.options, sso.progress_bar)

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

    return res

def photocurrent_mesolve(H, rho0, times, c_ops=[], sc_ops=[], e_ops=[],
                        _safe_mode=True, args={}, **kwargs):
    """
    Solve stochastic master equation using the photocurrent method.

    Parameters
    ----------

    H : :class:`qutip.Qobj`, or time dependent system.
        System Hamiltonian.
        Can depend on time, see StochasticSolverOptions help for format.

    rho0 : :class:`qutip.Qobj`
        Initial density matrix or state vector (ket).

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    c_ops : list of :class:`qutip.Qobj`, or time dependent Qobjs.
        Deterministic collapse operator which will contribute with a standard
        Lindblad type of dissipation.
        Can depend on time, see StochasticSolverOptions help for format.

    sc_ops : list of :class:`qutip.Qobj`, or time dependent Qobjs.
        List of stochastic collapse operators. Each stochastic collapse
        operator will give a deterministic and stochastic contribution
        to the eqaution of motion according to how the d1 and d2 functions
        are defined.
        Can depend on time, see StochasticSolverOptions help for format.

    e_ops : list of :class:`qutip.Qobj` / callback function single
        single operator or list of operators for which to evaluate
        expectation values.

    kwargs : *dictionary*
        Optional keyword arguments. See
        :class:`qutip.stochastic.StochasticSolverOptions`.

    Returns
    -------

    output: :class:`qutip.solver.Result`

        An instance of the class :class:`qutip.solver.Result`.
    """
    if isket(rho0):
        rho0 = ket2dm(rho0)

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

    sso = StochasticSolverOptionsPhoto(True, H=H, state0=rho0, times=times,
                                       c_ops=c_ops, sc_ops=sc_ops, e_ops=e_ops,
                                       args=args, **kwargs)

    if _safe_mode:
        _safety_checks(sso)

    if sso.m_ops is None:
        sso.m_ops = [op * 0 for op in sso.sc_ops]
    if not isinstance(sso.dW_factors, list):
        sso.dW_factors = [1] * len(sso.sc_ops)
    elif len(sso.dW_factors) != len(sso.sc_ops):
        raise Exception("The len of dW_factors is not the same as sc_ops")

    sso.solver_obj = PcSMESolver
    sso.solver_name = "photocurrent_mesolve"
    sso.LH = liouvillian(sso.H, c_ops=sso.c_ops) * sso.dt

    def _prespostdag(op):
        return spre(op) * spost(op.dag())

    sso.sops = [[spre(op._cdc()) + spost(op._cdc()),
                 spre(op._cdc()),
                 op.apply(_prespostdag)._f_norm2()] for op in sso.sc_ops]
    sso.ce_ops = [QobjEvo(spre(op)) for op in sso.e_ops]
    sso.cm_ops = [QobjEvo(spre(op)) for op in sso.m_ops]

    sso.LH.compile()
    [[op.compile() for op in ops] for ops in sso.sops]
    [op.compile() for op in sso.cm_ops]
    [op.compile() for op in sso.ce_ops]

    res = _sesolve_generic(sso, sso.options, sso.progress_bar)
    res.num_collapse = [np.count_nonzero(noise) for noise in res.noise]

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

    return res

def photocurrent_sesolve(H, psi0, times, sc_ops=[], e_ops=[],
                        _safe_mode=True, args={}, **kwargs):
    """
    Solve stochastic schrodinger equation using the photocurrent method.

    Parameters
    ----------

    H : :class:`qutip.Qobj`, or time dependent system.
        System Hamiltonian.
        Can depend on time, see StochasticSolverOptions help for format.

    psi0 : :class:`qutip.Qobj`
        Initial state vector (ket).

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    sc_ops : list of :class:`qutip.Qobj`, or time dependent Qobjs.
        List of stochastic collapse operators. Each stochastic collapse
        operator will give a deterministic and stochastic contribution
        to the eqaution of motion according to how the d1 and d2 functions
        are defined.
        Can depend on time, see StochasticSolverOptions help for format.

    e_ops : list of :class:`qutip.Qobj` / callback function single
        single operator or list of operators for which to evaluate
        expectation values.

    kwargs : *dictionary*
        Optional keyword arguments. See
        :class:`qutip.stochastic.StochasticSolverOptions`.

    Returns
    -------

    output: :class:`qutip.solver.Result`

        An instance of the class :class:`qutip.solver.Result`.
    """
    if isinstance(e_ops, dict):
        e_ops_dict = e_ops
        e_ops = [e for e in e_ops.values()]
    else:
        e_ops_dict = None

    sso = StochasticSolverOptionsPhoto(False, H=H, state0=psi0, times=times,
                                       sc_ops=sc_ops, e_ops=e_ops,
                                       args=args, **kwargs)

    if _safe_mode:
        _safety_checks(sso)

    if sso.m_ops is None:
        sso.m_ops = [op * 0 for op in sso.sc_ops]
    if not isinstance(sso.dW_factors, list):
        sso.dW_factors = [1] * len(sso.sc_ops)
    elif len(sso.dW_factors) != len(sso.sc_ops):
        raise Exception("The len of dW_factors is not the same as sc_ops")

    sso.solver_obj = PcSSESolver
    sso.solver_name = "photocurrent_sesolve"
    sso.sops = [[op, op._cdc()] for op in sso.sc_ops]
    sso.LH = sso.H * (-1j*sso.dt)
    for ops in sso.sops:
        sso.LH -= ops[0]._cdc()*0.5*sso.dt
    sso.ce_ops = [QobjEvo(op) for op in sso.e_ops]
    sso.cm_ops = [QobjEvo(op) for op in sso.m_ops]

    sso.LH.compile()
    [[op.compile() for op in ops] for ops in sso.sops]
    [op.compile() for op in sso.cm_ops]
    [op.compile() for op in sso.ce_ops]

    res = _sesolve_generic(sso, sso.options, sso.progress_bar)
    res.num_collapse = [np.count_nonzero(noise) for noise in res.noise]

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

    return res

def general_stochastic(state0, times, d1, d2, e_ops=[], m_ops=[],
                       _safe_mode=True, len_d2=1, args={}, **kwargs):
    """
    Solve stochastic general equation. Dispatch to specific solvers
    depending on the value of the `solver` keyword argument.

    Parameters
    ----------

    state0 : :class:`qutip.Qobj`
        Initial state vector (ket) or density matrix as a vector.

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    d1 : function, callable class
        Function representing the deterministic evolution of the system.

        def d1(time (double), state (as a np.array vector)):
            return 1d np.array

    d2 : function, callable class
        Function representing the stochastic evolution of the system.

        def d2(time (double), state (as a np.array vector)):
            return 2d np.array (N_sc_ops, len(state0))

    len_d2 : int
        Number of output vector produced by d2

    e_ops : list of :class:`qutip.Qobj`
        single operator or list of operators for which to evaluate
        expectation values.
        Must be a superoperator if the state vector is a density matrix.

    kwargs : *dictionary*
        Optional keyword arguments. See
        :class:`qutip.stochastic.StochasticSolverOptions`.

    Returns
    -------

    output: :class:`qutip.solver.Result`
        An instance of the class :class:`qutip.solver.Result`.
    """

    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 "solver" not in kwargs:
        kwargs["solver"] = 50

    sso = StochasticSolverOptions(False, H=None, state0=state0, times=times,
                                  e_ops=e_ops, args=args, **kwargs)
    if sso.solver_code not in [50, 100, 150]:
        raise ValueError("Only Euler, platen, platen15 can be " +
                         "used for the general stochastic solver.")

    sso.d1 = d1
    sso.d2 = d2
    if _safe_mode:
        # This state0_vec is computed as mat2vec(state0.full()).ravel()
        # in the sso init.
        state0_vec = sso.rho0
        l_vec = state0_vec.shape[0]
        try:
            out_d1 = d1(0., sso.rho0)
        except Exception as e:
            raise RuntimeError("Safety check: d1(0., state0_vec) failed.:\n" +
                               str(e)) from e
        except:
            raise RuntimeError("Safety check: d1(0., state0_vec) failed.")
        try:
            out_d2 = d2(0., sso.rho0)
        except Exception as e:
            raise RuntimeError("Safety check: d2(0., state0_vec) failed:\n" +
                               str(e)) from e
        except:
            raise RuntimeError("Safety check: d2(0., state0_vec) failed.")

        msg_d1 = ("d1 must return an 1d numpy array with the same number "
                  "of elements as the initial state as a vector.")
        if not isinstance(out_d1, np.ndarray):
            raise TypeError(msg_d1)
        if (out_d1.ndim != 1
                or out_d1.shape[0] != l_vec or len(out_d1.shape) != 1):
            raise ValueError(msg_d1)

        msg_d2 = ("Safety check: d2 must return a 2d numpy array "
                  "with the shape (len_d2, len(state0_vec) ).")
        if not isinstance(out_d2, np.ndarray):
            raise TypeError(msg_d2)
        if (out_d2.ndim != 2
                or out_d2.shape[1] != l_vec or out_d2.shape[0] != len_d2):
            raise ValueError(msg_d2)
        if out_d1.dtype != np.dtype('complex128') or \
           out_d2.dtype != np.dtype('complex128'):
            raise ValueError("Safety check: d1 and d2 must return " +
                             "complex numpy array.")
        msg_e_ops = ("Safety check: The shape of the e_ops "
                     "does not fit the intial state.")
        for op in sso.e_ops:
            shape_op = op.shape
            if sso.me:
                if shape_op[0]**2 != l_vec or shape_op[1]**2 != l_vec:
                    raise ValueError(msg_e_ops)
            else:
                if shape_op[0] != l_vec or shape_op[1] != l_vec:
                    raise ValueError(msg_e_ops +
                                     " Expecting e_ops as superoperators.")

    sso.m_ops = []
    sso.cm_ops = []
    if sso.store_measurement:
        if not m_ops:
            raise ValueError("General stochastic needs explicit " +
                             "m_ops to store measurement.")
        sso.m_ops = m_ops
        sso.cm_ops = [QobjEvo(op) for op in sso.m_ops]
        [op.compile() for op in sso.cm_ops]
        if sso.dW_factors is None:
            sso.dW_factors = [1.] * len(sso.m_ops)
        elif len(sso.dW_factors) == 1:
                sso.dW_factors = sso.dW_factors * len(sso.m_ops)
        elif len(sso.dW_factors) != len(sso.m_ops):
            raise ValueError("The number of dW_factors must fit" +
                             " the number of m_ops.")

    if sso.dW_factors is None:
        sso.dW_factors = [1.] * len_d2
    sso.sops = [None] * len_d2
    sso.ce_ops = [QobjEvo(op) for op in sso.e_ops]
    [op.compile() for op in sso.ce_ops]

    sso.solver_obj = GenericSSolver
    sso.solver_name = "general_stochastic_solver_" + sso.solver

    ssolver = GenericSSolver()
    # ssolver.set_data(sso)
    ssolver.set_solver(sso)

    res = _sesolve_generic(sso, sso.options, sso.progress_bar)

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

    return res

def _safety_checks(sso):
    l_vec = sso.rho0.shape[0]
    if sso.H.cte.issuper:
        if not sso.me:
            raise
        shape_op = sso.H.cte.shape
        if shape_op[0] != l_vec or shape_op[1] != l_vec:
            raise Exception("The size of the hamiltonian does "
                            "not fit the intial state")
    else:
        shape_op = sso.H.cte.shape
        if sso.me:
            if shape_op[0]**2 != l_vec or shape_op[1]**2 != l_vec:
                raise Exception("The size of the hamiltonian does "
                                "not fit the intial state")
        else:
            if shape_op[0] != l_vec or shape_op[1] != l_vec:
                raise Exception("The size of the hamiltonian does "
                                "not fit the intial state")

    for op in sso.sc_ops:
        if op.cte.issuper:
            if not sso.me:
                raise
            shape_op = op.cte.shape
            if shape_op[0] != l_vec or shape_op[1] != l_vec:
                raise Exception("The size of the sc_ops does "
                                "not fit the intial state")
        else:
            shape_op = op.cte.shape
            if sso.me:
                if shape_op[0]**2 != l_vec or shape_op[1]**2 != l_vec:
                    raise Exception("The size of the sc_ops does "
                                    "not fit the intial state")
            else:
                if shape_op[0] != l_vec or shape_op[1] != l_vec:
                    raise Exception("The size of the sc_ops does "
                                    "not fit the intial state")

    for op in sso.c_ops:
        if op.cte.issuper:
            if not sso.me:
                raise
            shape_op = op.cte.shape
            if shape_op[0] != l_vec or shape_op[1] != l_vec:
                raise Exception("The size of the c_ops does "
                                "not fit the intial state")
        else:
            shape_op = op.cte.shape
            if sso.me:
                if shape_op[0]**2 != l_vec or shape_op[1]**2 != l_vec:
                    raise Exception("The size of the c_ops does "
                                    "not fit the intial state")
            else:
                if shape_op[0] != l_vec or shape_op[1] != l_vec:
                    raise Exception("The size of the c_ops does "
                                    "not fit the intial state")

    for op in sso.e_ops:
        shape_op = op.shape
        if sso.me:
            if shape_op[0]**2 != l_vec or shape_op[1]**2 != l_vec:
                raise Exception("The size of the e_ops does "
                                "not fit the intial state")
        else:
            if shape_op[0] != l_vec or shape_op[1] != l_vec:
                raise Exception("The size of the e_ops does "
                                "not fit the intial state")

    if sso.m_ops is not None:
        for op in sso.m_ops:
            shape_op = op.shape
            if sso.me:
                if shape_op[0]**2 != l_vec or shape_op[1]**2 != l_vec:
                    raise Exception("The size of the m_ops does "
                                    "not fit the intial state")
            else:
                if shape_op[0] != l_vec or shape_op[1] != l_vec:
                    raise Exception("The size of the m_ops does "
                                    "not fit the intial state")

def _sesolve_generic(sso, options, progress_bar):
    """
    Internal function. See smesolve.
    """
    res = Result()
    res.times = sso.times
    res.expect = np.zeros((len(sso.e_ops), len(sso.times)), dtype=complex)
    res.ss = np.zeros((len(sso.e_ops), len(sso.times)), dtype=complex)
    res.measurement = []
    res.solver = sso.solver_name
    res.ntraj = sso.ntraj
    res.num_expect = len(sso.e_ops)

    nt = sso.ntraj
    task = _single_trajectory
    map_kwargs = {'progress_bar': sso.progress_bar}
    map_kwargs.update(sso.map_kwargs)
    task_args = (sso,)
    task_kwargs = {}

    results = sso.map_func(task, list(range(sso.ntraj)),
                           task_args, task_kwargs, **map_kwargs)
    noise = []
    for result in results:
        states_list, dW, m, expect = result
        res.states.append(states_list)
        noise.append(dW)
        res.measurement.append(m)
        res.expect += expect
        res.ss += expect * expect
    res.noise = np.stack(noise)

    if sso.store_all_expect:
        paths_expect = []
        for result in results:
            paths_expect.append(result[3])
        res.runs_expect = np.stack(paths_expect)

    # average density matrices (vectorized maybe)
    # ajgpitch 2019-10-25: np.any(res.states) seems to error
    # I guess there may be a potential exception if there are no states?
    # store individual trajectory states
    res.traj_states = res.states
    res.avg_states = None
    if options.average_states and options.store_states:
        avg_states_list = []
        for n in range(len(res.times)):
            tslot_states = [res.states[mm][n].data for mm in range(nt)]
            if len(tslot_states) > 0:
                state = Qobj(np.sum(tslot_states),
                             dims=res.states[0][n].dims).unit()
                avg_states_list.append(state)
        # store average states
        res.states = res.avg_states = avg_states_list

    # average
    res.expect = res.expect / nt

    # standard error
    if nt > 1:
        res.se = (res.ss - nt * (res.expect ** 2)) / (nt * (nt - 1))
    else:
        res.se = None

    # convert complex data to real if hermitian
    res.expect = [np.real(res.expect[n, :])
                  if e.isherm else res.expect[n, :]
                  for n, e in enumerate(sso.e_ops)]

    return res

def _single_trajectory(i, sso):
    # Only one step?
    ssolver = sso.solver_obj()
    #ssolver.set_data(sso)
    ssolver.set_solver(sso)
    result = ssolver.cy_sesolve_single_trajectory(i)#, sso)
    return result

# The code for ssepdpsolve have been moved to the file pdpsolve.
# The call is still in stochastic for consistance.
def ssepdpsolve(H, psi0, times, c_ops, e_ops, **kwargs):
    """
    A stochastic (piecewse deterministic process) PDP solver for wavefunction
    evolution. For most purposes, use :func:`qutip.mcsolve` instead for quantum
    trajectory simulations.

    Parameters
    ----------

    H : :class:`qutip.Qobj`
        System Hamiltonian.

    psi0 : :class:`qutip.Qobj`
        Initial state vector (ket).

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    c_ops : list of :class:`qutip.Qobj`
        Deterministic collapse operator which will contribute with a standard
        Lindblad type of dissipation.

    e_ops : list of :class:`qutip.Qobj` / callback function single
        single operator or list of operators for which to evaluate
        expectation values.

    kwargs : *dictionary*
        Optional keyword arguments. See
        :class:`qutip.stochastic.StochasticSolverOptions`.

    Returns
    -------

    output: :class:`qutip.solver.Result`

        An instance of the class :class:`qutip.solver.Result`.

    """
    return main_ssepdpsolve(H, psi0, times, c_ops, e_ops, **kwargs)

# The code for smepdpsolve have been moved to the file pdpsolve.
# The call is still in stochastic for consistance.
def smepdpsolve(H, rho0, times, c_ops, e_ops, **kwargs):
    """
    A stochastic (piecewse deterministic process) PDP solver for density matrix
    evolution.

    Parameters
    ----------

    H : :class:`qutip.Qobj`
        System Hamiltonian.

    rho0 : :class:`qutip.Qobj`
        Initial density matrix.

    times : *list* / *array*
        List of times for :math:`t`. Must be uniformly spaced.

    c_ops : list of :class:`qutip.Qobj`
        Deterministic collapse operator which will contribute with a standard
        Lindblad type of dissipation.

    sc_ops : list of :class:`qutip.Qobj`
        List of stochastic collapse operators. Each stochastic collapse
        operator will give a deterministic and stochastic contribution
        to the eqaution of motion according to how the d1 and d2 functions
        are defined.

    e_ops : list of :class:`qutip.Qobj` / callback function single
        single operator or list of operators for which to evaluate
        expectation values.

    kwargs : *dictionary*
        Optional keyword arguments. See
        :class:`qutip.stochastic.StochasticSolverOptions`.

    Returns
    -------

    output: :class:`qutip.solver.Result`

        An instance of the class :class:`qutip.solver.Result`.

    """
    return main_smepdpsolve(H, rho0, times, c_ops, e_ops, **kwargs)


# Copyright 2020 NREL

# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

import numpy as np

from ....utilities import cosd, sind, tand
from ..base_velocity_deficit import VelocityDeficit

class GaussianModel(VelocityDeficit):
    """
    This is the super-class for all Gaussian-type wake models. It includes
    implementations of functions that subclasses should use to perform
    Gaussian-related calculations (see :cite:`gmb-King2019Controls`)

    References:
        .. bibliography:: /source/zrefs.bib
            :style: unsrt
            :filter: docname in docnames
            :keyprefix: gmb-
    """

    def __init__(self, parameter_dictionary):
        """
        See super-class for initialization details.

        Args:
            parameter_dictionary (dict): Model-specific parameters.
        """
        super().__init__(parameter_dictionary)

    def correction_steps(
        self, U_local, U, V, W, x_locations, y_locations, turbine, turbine_coord
    ):
        """
        This method corrects the U-component velocities when yaw added recovery
        is enabled. For more details on how the velocities are changed, see [1].
        # TODO add reference to 1

        Args:
            U_local (np.array): U-component velocities across the flow field.
            U (np.array): U-component velocity deficits across the flow field.
            V (np.array): V-component velocity deficits across the flow field.
            W (np.array): W-component velocity deficits across the flow field.
            x_locations (np.array): Streamwise locations in wake.
            y_locations (np.array): Spanwise locations in wake.
            turbine (:py:class:`floris.simulation.turbine.Turbine`):
                Turbine object.
            turbine_coord (:py:obj:`floris.simulation.turbine_map.TurbineMap.coords`):
                Spatial coordinates of wind turbine.

        Returns:
            np.array: U-component velocity deficits across the flow field.
        """
        if self.use_yaw_added_recovery:
            U = self.yaw_added_recovery_correction(
                U_local, U, W, x_locations, y_locations, turbine, turbine_coord
            )
        return U

    def calculate_VW(
        self, V, W, coord, turbine, flow_field, x_locations, y_locations, z_locations
    ):
        """
        This method calculates the V- and W-component velocities using
        methods developed in [1].
        # TODO add reference to 1
        # TODO is this function needed? It simply calls another function

        Args:
            V (np.array): V-component velocity deficits across the flow field.
            W (np.array): W-component velocity deficits across the flow field.
            coord (:py:obj:`floris.simulation.turbine_map.TurbineMap.coords`):
                Spatial coordinates of wind turbine.
            turbine (:py:class:`floris.simulation.turbine.Turbine`):
                Turbine object.
            flow_field ([type]): [description]
            x_locations (np.array): Streamwise locations in wake.
            y_locations (np.array): Spanwise locations in wake.
            z_locations (np.array): Vertical locations in wake.

        Raises:
            ValueError: It appears that 'use_yaw_added_recovery' is set
                to True and 'calculate_VW_velocities' is set to False.
                This configuration is not valid. Please set
                'calculate_VW_velocities' to True if you wish to use
                yaw-added recovery.

        Returns:
            np.array, np.array:

                - V-component velocity deficits across the flow field.
                - W-component velocity deficits across the flow field.
        """
        if self.use_yaw_added_recovery:
            if not self.calculate_VW_velocities:
                err_msg = (
                    "It appears that 'use_yaw_added_recovery' is set "
                    + "to True and 'calculate_VW_velocities' is set to False. "
                    + "This configuration is not valid. Please set "
                    + "'calculate_VW_velocities' to True if you wish to use "
                    + "yaw-added recovery."
                )
                self.logger.error(err_msg, stack_info=True)
                raise ValueError(err_msg)
        if self.calculate_VW_velocities:
            V, W = self.calc_VW(
                coord, turbine, flow_field, x_locations, y_locations, z_locations
            )
        return V, W

    def yaw_added_recovery_correction(
        self, U_local, U, W, x_locations, y_locations, turbine, turbine_coord
    ):
        """
        This method corrects the U-component velocities when yaw added recovery
        is enabled. For more details on how the velocities are changed, see [1].
        # TODO add reference to 1

        Args:
            U_local (np.array): U-component velocities across the flow field.
            U (np.array): U-component velocity deficits across the flow field.
            W (np.array): W-component velocity deficits across the flow field.
            x_locations (np.array): Streamwise locations in wake.
            y_locations (np.array): Spanwise locations in wake.
            turbine (:py:class:`floris.simulation.turbine.Turbine`):
                Turbine object.
            turbine_coord (:py:obj:`floris.simulation.turbine_map.TurbineMap.coords`):
                Spatial coordinates of wind turbine.

        Returns:
            np.array: U-component velocity deficits across the flow field.
        """
        # compute the velocity without modification
        U1 = U_local - U

        # set dimensions
        D = turbine.rotor_diameter
        xLocs = x_locations - turbine_coord.x1
        ky = self.ka * turbine.current_turbulence_intensity + self.kb
        U2 = (np.mean(W) * xLocs) / ((ky * xLocs + D / 2))
        U_total = U1 + np.nan_to_num(U2)

        # turn it back into a deficit
        U = U_local - U_total

        # zero out anything before the turbine
        U[x_locations < turbine_coord.x1] = 0

        return U

    def calc_VW(
        self, coord, turbine, flow_field, x_locations, y_locations, z_locations
    ):
        """
        This method calculates the V- and W-component velocities using
        methods developed in [1].
        # TODO add reference to 1

        Args:
            coord (:py:obj:`floris.simulation.turbine_map.TurbineMap.coords`):
                Spatial coordinates of wind turbine.
            turbine (:py:class:`floris.simulation.turbine.Turbine`):
                Turbine object.
            flow_field ([type]): [description]
            x_locations (np.array): Streamwise locations in wake.
            y_locations (np.array): Spanwise locations in wake.
            z_locations (np.array): Vertical locations in wake.

        Returns:
            np.array, np.array:

                - V-component velocity deficits across the flow field.
                - W-component velocity deficits across the flow field.
        """
        # turbine parameters
        D = turbine.rotor_diameter
        HH = turbine.hub_height
        yaw = turbine.yaw_angle
        Ct = turbine.Ct
        TSR = turbine.tsr
        aI = turbine.aI

        # flow parameters
        Uinf = np.mean(flow_field.wind_map.grid_wind_speed)

        scale = 1.0
        vel_top = (
            Uinf
            * ((HH + D / 2) / flow_field.specified_wind_height) ** flow_field.wind_shear
        ) / Uinf
        vel_bottom = (
            Uinf
            * ((HH - D / 2) / flow_field.specified_wind_height) ** flow_field.wind_shear
        ) / Uinf
        Gamma_top = (
            scale * (np.pi / 8) * D * vel_top * Uinf * Ct * sind(yaw) * cosd(yaw)
        )
        Gamma_bottom = (
            -scale * (np.pi / 8) * D * vel_bottom * Uinf * Ct * sind(yaw) * cosd(yaw)
        )
        Gamma_wake_rotation = (
            0.25 * 2 * np.pi * D * (aI - aI ** 2) * turbine.average_velocity / TSR
        )

        # compute the spanwise and vertical velocities induced by yaw
        eps = self.eps_gain * D  # Use set value

        # decay the vortices as they move downstream - using mixing length
        lmda = D / 8
        kappa = 0.41
        lm = kappa * z_locations / (1 + kappa * z_locations / lmda)
        z = np.linspace(
            np.min(z_locations), np.max(z_locations), np.shape(flow_field.u_initial)[2]
        )
        dudz_initial = np.gradient(flow_field.u_initial, z, axis=2)
        nu = lm ** 2 * np.abs(dudz_initial[0, :, :])

        # top vortex
        yLocs = y_locations + 0.01 - (coord.x2)
        zT = z_locations + 0.01 - (HH + D / 2)
        rT = yLocs ** 2 + zT ** 2
        V1 = (
            (zT * Gamma_top)
            / (2 * np.pi * rT)
            * (1 - np.exp(-rT / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        W1 = (
            (-yLocs * Gamma_top)
            / (2 * np.pi * rT)
            * (1 - np.exp(-rT / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        # bottom vortex
        zB = z_locations + 0.01 - (HH - D / 2)
        rB = yLocs ** 2 + zB ** 2
        V2 = (
            (zB * Gamma_bottom)
            / (2 * np.pi * rB)
            * (1 - np.exp(-rB / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        W2 = (
            ((-yLocs * Gamma_bottom) / (2 * np.pi * rB))
            * (1 - np.exp(-rB / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        # top vortex - ground
        yLocs = y_locations + 0.01 - (coord.x2)
        zLocs = z_locations + 0.01 + (HH + D / 2)
        V3 = (
            (
                ((zLocs * -Gamma_top) / (2 * np.pi * (yLocs ** 2 + zLocs ** 2)))
                * (1 - np.exp(-(yLocs ** 2 + zLocs ** 2) / (eps ** 2)))
                + 0.0
            )
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        W3 = (
            ((-yLocs * -Gamma_top) / (2 * np.pi * (yLocs ** 2 + zLocs ** 2)))
            * (1 - np.exp(-(yLocs ** 2 + zLocs ** 2) / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        # bottom vortex - ground
        yLocs = y_locations + 0.01 - (coord.x2)
        zLocs = z_locations + 0.01 + (HH - D / 2)
        V4 = (
            (
                ((zLocs * -Gamma_bottom) / (2 * np.pi * (yLocs ** 2 + zLocs ** 2)))
                * (1 - np.exp(-(yLocs ** 2 + zLocs ** 2) / (eps ** 2)))
                + 0.0
            )
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        W4 = (
            ((-yLocs * -Gamma_bottom) / (2 * np.pi * (yLocs ** 2 + zLocs ** 2)))
            * (1 - np.exp(-(yLocs ** 2 + zLocs ** 2) / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        # wake rotation vortex
        zC = z_locations + 0.01 - (HH)
        rC = yLocs ** 2 + zC ** 2
        V5 = (
            (zC * Gamma_wake_rotation)
            / (2 * np.pi * rC)
            * (1 - np.exp(-rC / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        W5 = (
            (-yLocs * Gamma_wake_rotation)
            / (2 * np.pi * rC)
            * (1 - np.exp(-rC / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        # wake rotation vortex - ground effect
        yLocs = y_locations + 0.01 - coord.x2
        zLocs = z_locations + 0.01 + HH
        V6 = (
            (
                (
                    (zLocs * -Gamma_wake_rotation)
                    / (2 * np.pi * (yLocs ** 2 + zLocs ** 2))
                )
                * (1 - np.exp(-(yLocs ** 2 + zLocs ** 2) / (eps ** 2)))
                + 0.0
            )
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        W6 = (
            ((-yLocs * -Gamma_wake_rotation) / (2 * np.pi * (yLocs ** 2 + zLocs ** 2)))
            * (1 - np.exp(-(yLocs ** 2 + zLocs ** 2) / (eps ** 2)))
            * eps ** 2
            / (4 * nu * (x_locations - coord.x1) / Uinf + eps ** 2)
        )

        # print(Gamma_wake_rotation, np.mean(W5), np.mean(W6))

        # total spanwise velocity
        V = V1 + V2 + V3 + V4 + V5 + V6
        W = W1 + W2 + W3 + W4 + W5 + W6

        # no spanwise and vertical velocity upstream of the turbine
        V[x_locations < coord.x1 + 10] = 0.0
        W[x_locations < coord.x1 + 10] = 0.0
        W[W < 0] = 0

        return V, W

    @property
    def calculate_VW_velocities(self):
        """
        Flag to enable the calculation of V- and W-component velocities using
        methods developed in [1].

        **Note:** This is a virtual property used to "get" or "set" a value.

        Args:
            value (bool): Value to set.

        Returns:
            float: Value currently set.

        Raises:
            ValueError: Invalid value.
        """
        return self._calculate_VW_velocities

    @calculate_VW_velocities.setter
    def calculate_VW_velocities(self, value):
        if type(value) is not bool:
            err_msg = (
                "Value of calculate_VW_velocities must be type "
                + "float; {} given.".format(type(value))
            )
            self.logger.error(err_msg, stack_info=True)
            raise ValueError(err_msg)
        self._calculate_VW_velocities = value

    @property
    def use_yaw_added_recovery(self):
        """
        Flag to use yaw added recovery on the wake velocity using methods
        developed in [1].

        **Note:** This is a virtual property used to "get" or "set" a value.

        Args:
            value (bool): Value to set.

        Returns:
            float: Value currently set.

        Raises:
            ValueError: Invalid value.
        """
        return self._use_yaw_added_recovery

    @use_yaw_added_recovery.setter
    def use_yaw_added_recovery(self, value):
        if type(value) is not bool:
            # TODO Shouldn't this be a bool?
            err_msg = (
                "Value of use_yaw_added_recovery must be type "
                + "float; {} given.".format(type(value))
            )
            self.logger.error(err_msg, stack_info=True)
            raise ValueError(err_msg)
        self._use_yaw_added_recovery = value

    @property
    def eps_gain(self):
        """
        Tuning value for calculating the V- and W- component velocities using
        methods developed in [1].

        **Note:** This is a virtual property used to "get" or "set" a value.

        Args:
            value (bool): Value to set.

        Returns:
            float: Value currently set.

        Raises:
            ValueError: Invalid value.
        """
        return self._eps_gain

    @eps_gain.setter
    def eps_gain(self, value):
        if type(value) is not float:
            err_msg = "Value of eps_gain must be type " + "float; {} given.".format(
                type(value)
            )
            self.logger.error(err_msg, stack_info=True)
            raise ValueError(err_msg)
        self._eps_gain = value

    @staticmethod
    def mask_upstream_wake(y_locations, turbine_coord, yaw):
        """
        Calculates values to be used for masking the upstream wake relative to
        the current turbine.

        Args:
            y_locations (np.array): Spanwise locations in wake.
            turbine_coord (:py:obj:`floris.simulation.turbine_map.TurbineMap.coords`):
                Spatial coordinates of wind turbine.
            yaw (float): The turbine yaw angle.

        Returns:
            tuple: tuple containing:

                -   yR (np.array): Y locations to mask upstream wake.
                -   xR (np.array): X locations to mask upstream wake.
        """
        yR = y_locations - turbine_coord.x2
        xR = yR * tand(yaw) + turbine_coord.x1
        return xR, yR

    @staticmethod
    def initial_velocity_deficits(U_local, Ct):
        """
        Calculates the initial velocity deficits used in determining the wake
        expansion in a Gaussian wake velocity model.

        Args:
            U_local (np.array): U-component velocities across the flow field.
            Ct (float): The thrust coefficient of a turbine at the current
                operating conditions.

        Returns:
            tuple: tuple containing:

                -   uR (np.array): Initial velocity deficit used in calculation
                    of wake expansion.
                -   u0 (np.array): Initial velocity deficit used in calculation
                    of wake expansion.
        """
        uR = U_local * Ct / (2.0 * (1 - np.sqrt(1 - Ct)))
        u0 = U_local * np.sqrt(1 - Ct)
        return uR, u0

    @staticmethod
    def initial_wake_expansion(turbine, U_local, veer, uR, u0):
        """
        Calculates the initial wake widths associated with wake expansion.

        Args:
            turbine (:py:class:`floris.simulation.turbine.Turbine`):
                Turbine object.
            U_local (np.array): U-component velocities across the flow field.
            veer (float): The amount of veer across the rotor.
            uR (np.array): Initial velocity deficit used in calculation of wake
                expansion.
            u0 (np.array): Initial velocity deficit used in calculation of wake
                expansion.

        Returns:
            tuple: tuple containing:

                -   sigma_y0 (np.array): Initial wake width in the spanwise
                    direction.
                -   sigma_z0 (np.array): Initial wake width in the vertical
                    direction.
        """
        yaw = -1 * turbine.yaw_angle
        sigma_z0 = turbine.rotor_diameter * 0.5 * np.sqrt(uR / (U_local + u0))
        sigma_y0 = sigma_z0 * cosd(yaw) * cosd(veer)
        return sigma_y0, sigma_z0

    @staticmethod
    def gaussian_function(U, C, r, n, sigma):
        """
        A general form of the Gaussian function used in the Gaussian wake
        models.

        Args:
            U (np.array): U-component velocities across the flow field.
            C (np.array): Velocity deficit at the wake center normalized by the
                incoming wake velocity.
            r (float): Radial distance from the wake center.
            n (float): Exponent of radial distance from the wake center.
            sigma (np.array): Standard deviation of the wake.

        Returns:
            np.array: U (np.array): U-component velocity deficits across the
            flow field.
        """
        return U * C * np.exp(-1 * r ** n / (2 * sigma ** 2))


import numbers
from typing import Optional, Iterable, Tuple, Union

import numpy as np

from ..algorithms.matrixaverager import MatrixAverager, ErrorPropagationMethod

class Curve:
    _data: np.ndarray

    def __init__(self):
        pass

    @property
    def q(self) -> np.ndarray:
        return self._data[:, 0]

    @q.setter
    def q(self, newvalue):
        self._data[:, 0] = newvalue

    @property
    def intensity(self) -> np.ndarray:
        return self._data[:, 1]

    @intensity.setter
    def intensity(self, newvalue):
        self._data[:, 1] = newvalue

    @property
    def uncertainty(self) -> np.ndarray:
        return self._data[:, 2]

    @uncertainty.setter
    def uncertainty(self, newvalue):
        self._data[:, 2] = newvalue

    @property
    def quncertainty(self) -> np.ndarray:
        return self._data[:, 3]

    @quncertainty.setter
    def quncertainty(self, newvalue):
        self._data[:, 3] = newvalue

    @property
    def binarea(self) -> np.ndarray:
        return self._data[:, 4]

    @binarea.setter
    def binarea(self, newvalue):
        self._data[:, 4] = newvalue

    @property
    def pixel(self) -> np.ndarray:
        return self._data[:, 5]

    @pixel.setter
    def pixel(self, newvalue):
        self._data[:, 5] = newvalue

    @classmethod
    def fromFile(cls, filename: str, *args, **kwargs) -> "Curve":
        data = np.loadtxt(filename, *args, **kwargs)
        self = cls()
        self._data = np.empty((data.shape[0], 6)) * np.nan
        self._data[:, :min(data.shape[1], 6)] = data[:, :min(data.shape[1], 6)]
        return self

    @classmethod
    def fromArray(cls, array: np.ndarray) -> "Curve":
        self = cls()
        self._data = np.empty((array.shape[0], 6), array.dtype) + np.nan
        self._data[:, :array.shape[1]] = array
        return self

    def __array__(self) -> np.ndarray:
        return self._data

    asArray = __array__

    @classmethod
    def fromVectors(cls, q: np.ndarray, intensity: np.ndarray, uncertainty: Optional[np.ndarray] = None,
                    quncertainty: Optional[np.ndarray] = None, binarea: Optional[np.ndarray] = None,
                    pixel: Optional[np.ndarray] = None) -> "Curve":
        assert (q is not None) and (intensity is not None)
        if not all([
            q.ndim == 1,
            intensity.ndim == 1,
            (uncertainty is None) or (uncertainty.ndim == 1),
            (quncertainty is None) or (quncertainty.ndim == 1),
            (binarea is None) or (binarea.ndim == 1),
            (pixel is None) or (pixel.ndim == 1)
        ]):
            raise ValueError('Supplied arrays must be one-dimensional.')
        if not all([
            len(intensity) == len(q),
            (uncertainty is None) or (len(uncertainty) == len(q)),
            (quncertainty is None) or (len(quncertainty) == len(q)),
            (binarea is None) or (len(binarea) == len(q)),
            (pixel is None) or (len(pixel) == len(q))
        ]):
            raise ValueError('All vectors must have the same length.')
        self = cls()
        self._data = np.empty((len(q), 6))
        self._data[:, 0] = q
        self._data[:, 1] = intensity
        self._data[:, 2] = uncertainty if uncertainty is not None else np.nan
        self._data[:, 3] = quncertainty if quncertainty is not None else np.nan
        self._data[:, 4] = binarea if binarea is not None else np.nan
        self._data[:, 5] = pixel if pixel is not None else np.nan
        return self

    def trim(self, left: float = -np.inf, right: float = np.inf, bottom=-np.inf, top=np.inf, bypixel: bool = False):
        if bypixel:
            xtrimidx = np.logical_and(self._data[:, 5] <= right, self._data[:, 5] >= left)
        else:
            xtrimidx = np.logical_and(self._data[:, 0] <= right, self._data[:, 0] >= left)
        ytrimidx = np.logical_and(self._data[:, 1] <= top, self._data[:, 1] >= bottom)
        curve = Curve()
        curve._data = self._data[np.logical_and(xtrimidx, ytrimidx)]
        return curve

    def __len__(self) -> int:
        return self._data.shape[0]

    def sanitize(self) -> "Curve":
        return Curve.fromArray(self._data[self.isvalid(), :])

    @classmethod
    def average(cls, curves: Iterable["Curve"], ierrorpropagation: ErrorPropagationMethod,
                qerrorpropagation: ErrorPropagationMethod) -> "Curve":
        qavg = MatrixAverager(errorpropagationmethod=qerrorpropagation)
        iavg = MatrixAverager(errorpropagationmethod=ierrorpropagation)
        aavg = MatrixAverager(errorpropagationmethod=ierrorpropagation)
        pavg = MatrixAverager(errorpropagationmethod=ierrorpropagation)
        for c in curves:
            qavg.add(c.q, c.quncertainty)
            iavg.add(c.intensity, c.uncertainty)
            aavg.add(c.binarea, c.binarea)
            pavg.add(c.pixel, c.pixel)
        q, dq = qavg.get()
        i, di = iavg.get()
        a = aavg.get()[0]
        p = pavg.get()[0]
        return Curve.fromVectors(q, i, di, dq, a, p)

    def isfinite(self) -> np.ndarray:
        idx = np.logical_and(np.isfinite(self.q, np.isfinite(self.intensity)))
        for vector in [self.uncertainty, self.quncertainty, self.pixel, self.binarea]:
            if np.isfinite(vector).sum() > 0:  # if not all elements are NaN
                idx = np.logical_and(idx, np.isfinite(vector))
        return idx

    def isvalid(self) -> np.ndarray:
        idx = np.logical_and(np.isfinite(self.q), np.isfinite(self.intensity))
        idx = np.logical_and(idx, self.q > 0)
        for vector in [self.uncertainty, self.quncertainty, self.pixel]:
            if np.isfinite(vector).sum() > 0:
                idx = np.logical_and(
                    idx,
                    np.logical_and(
                        np.isfinite(vector),
                        vector >= 0
                    ))
        if np.isfinite(self.binarea).sum() > 0:
            idx = np.logical_and(
                idx,
                np.logical_and(
                    np.isfinite(self.binarea),
                    self.binarea > 0
                ))
        return idx

    def __getitem__(self, item) -> "Curve":
        if isinstance(item, np.ndarray) and (item.dtype == np.bool):
            return Curve.fromArray(self._data[item, :])

    def _checkcompatibility(self, other: "Curve", maxdifferenceratio: float = 0.005):
        incompatibility = np.abs(self.q - other.q) / np.max(np.mean((self.q, other.q)), axis=0) * 2
        if np.any(incompatibility[np.isfinite(incompatibility)] > maxdifferenceratio):  # 0.01 means 1%
            raise ValueError(
                f'The two q-scales are incompatible. Max. incompatibility: {incompatibility[np.isfinite(incompatibility)].max()}')

    def __sub__(self, other: Union["Curve", numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, Curve):
            self._checkcompatibility(other)
            return self.fromVectors(
                q=0.5 * (self.q + other.q),
                intensity=self.intensity - other.intensity,
                uncertainty=(self.uncertainty ** 2 + other.uncertainty ** 2) ** 0.5,
                quncertainty=(self.quncertainty ** 2 + other.quncertainty ** 2) ** 0.5 / 2.0,
                binarea=None,
                pixel=None
            )
        elif isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity - other,
                uncertainty=self.uncertainty,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity - other[0],
                uncertainty=(self.uncertainty ** 2 + other[1] ** 2) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented

    def __add__(self, other: Union["Curve", numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, Curve):
            self._checkcompatibility(other)
            return self.fromVectors(
                q=0.5 * (self.q + other.q),
                intensity=self.intensity + other.intensity,
                uncertainty=(self.uncertainty ** 2 + other.uncertainty ** 2) ** 0.5,
                quncertainty=(self.quncertainty ** 2 + other.quncertainty ** 2) ** 0.5 / 2.0,
                binarea=None,
                pixel=None
            )
        elif isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity + other,
                uncertainty=self.uncertainty,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity + other[0],
                uncertainty=(self.uncertainty ** 2 + other[1] ** 2) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented

    def __rsub__(self, other: Union[numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=other - self.intensity,
                uncertainty=self.uncertainty,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=other[0] - self.intensity,
                uncertainty=(self.uncertainty ** 2 + other[1] ** 2) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented

    def __radd__(self, other: Union[numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity + other,
                uncertainty=self.uncertainty,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity + other[0],
                uncertainty=(self.uncertainty ** 2 + other[1] ** 2) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented

    def __truediv__(self, other: Union["Curve", numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, Curve):
            self._checkcompatibility(other)
            return self.fromVectors(
                q=0.5 * (self.q + other.q),
                intensity=self.intensity / other.intensity,
                uncertainty=(
                                    self.uncertainty ** 2 / other.intensity ** 2 + other.uncertainty ** 2 * self.intensity ** 2 / other.intensity ** 4) ** 0.5,
                quncertainty=(self.quncertainty ** 2 + other.quncertainty ** 2) ** 0.5 / 2.0,
                binarea=None,
                pixel=None
            )
        elif isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity / other,
                uncertainty=np.abs(self.uncertainty / other),
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity / other[0],
                uncertainty=(self.uncertainty ** 2 / other[0] ** 2 + other[1] ** 2 * self.intensity ** 2 / other[
                    0] ** 4) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented

    def __mul__(self, other: Union["Curve", numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, Curve):
            self._checkcompatibility(other)
            return self.fromVectors(
                q=0.5 * (self.q + other.q),
                intensity=self.intensity * other.intensity,
                uncertainty=(
                                    self.uncertainty ** 2 * other.intensity ** 2 + other.uncertainty ** 2 * self.intensity ** 2) ** 0.5,
                quncertainty=(self.quncertainty ** 2 + other.quncertainty ** 2) ** 0.5 / 2.0,
                binarea=None,
                pixel=None
            )
        elif isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity * other,
                uncertainty=np.abs(self.uncertainty * other),
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity * other[0],
                uncertainty=(self.uncertainty ** 2 * other[0] ** 2 + self.intensity ** 2 * other[1] ** 2) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented

    def __rtruediv__(self, other: Union[numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=other / self.intensity,
                uncertainty=np.abs(other * self.uncertainty / self.intensity ** 2),
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=other[0] / self.intensity,
                uncertainty=(self.uncertainty ** 2 / self.intensity ** 4 * other[0] ** 2 + other[
                    1] ** 2 / self.intensity ** 2) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented

    def __rmul__(self, other: Union[numbers.Real, Tuple[numbers.Real, numbers.Real]]) -> "Curve":
        if isinstance(other, numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity * other,
                uncertainty=np.abs(self.uncertainty * other),
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        elif isinstance(other, tuple) and (len(other) == 2) and isinstance(other[0], numbers.Real) and isinstance(
                other[1], numbers.Real):
            return self.fromVectors(
                q=self.q,
                intensity=self.intensity * other[0],
                uncertainty=(self.uncertainty ** 2 * other[0] ** 2 + self.intensity ** 2 * other[1] ** 2) ** 0.5,
                quncertainty=self.quncertainty,
                binarea=self.binarea,
                pixel=self.pixel
            )
        else:
            return NotImplemented


"""
Functions for aeri retrievals.

"""

import numpy as np
from scipy.optimize import brentq

from act.retrievals.irt import irt_response_function, sum_function_irt

def aeri2irt(
    aeri_ds,
    wnum_name='wnum',
    rad_name='mean_rad',
    hatch_name='hatchOpen',
    tolerance=0.1,
    temp_low=150.0,
    temp_high=320.0,
    maxiter=200,
):
    """
    This function will integrate over the correct wavenumber values to produce
    the effective IRT temperature.

    As a note from the ARM IRT Instrument Handbook
    A positive bias of the sky temperature is exhibited by the downwelling IRT,
    compared to the AERI, during clear-sky conditions when the sky temperature
    is less than ~180K. The effect depends on the characteristics of the
    individual IRT and the internal reference temperature of the IRT. The
    greatest difference compared to AERI will occur when the sky is very clear,
    dry, and cold and the ambient temperature is relatively hot, maximizing the
    difference in temperature between the sky and instrument, and the
    calibration of the IRT at the lower limit of 223K was not performed
    accurately. This bias is especially apparent at high-latitude sites
    (e.g., NSA, OLI, and AWR).

    https://www.arm.gov/publications/tech_reports/handbooks/irt_handbook.pdf

    Author - Ken Kehoe

    Parameters
    ----------
    aeri_ds : Xarray Dataset Object
        The Dataset object containing AERI data.
    wnum_name : str
        The variable name for coordinate dimention of wave number Xarray Dataset.
    hatch_name : str or None
        The variable name for hatch status. If set to None will not try to set
        when hatch is not opent to NaN.
    rad_name : str
        The variable name for mean radiance in Xarray Dataset.
    tolerance : float
        The tolerance value to try and match for returned temperature.
    temp_low : float
        The initial low value to use in zbren function to invert radiances.
    temp_high : float
        The initial low value to use in zbren function to invert radiances.
    maxiter : int
        The maximum number if iterations to use with invertion process.
        Prevents runaway processes.

    Returns
    -------
    obj : Xarray Dataset Object or None
        The aeri_ds Dataset with new DataArray of temperatures added under
        variable name 'aeri_irt_equiv_temperature'.

    """
    # Get data values
    rf_wnum, rf = irt_response_function()
    wnum = aeri_ds[wnum_name].values
    mean_rad = aeri_ds[rad_name].values

    # Pull out AERI data for correct wavenumbers and apply response function --;
    index = np.where((wnum >= (rf_wnum[0] - 0.001)) & (wnum <= (rf_wnum[-1] + 0.001)))[0]
    if index.size == 0:
        raise ValueError('No wavenumbers match for aeri2irt')

    wnum = wnum[index]
    mean_rad = mean_rad[:, index]
    # If the wavenumbers in AERI data are not close enough to response function
    # match the wavenumbers and adjust.
    atol = 0.001
    if not np.all(np.isclose(wnum, rf_wnum, atol=atol)):
        index_wnum = []
        index_rf = []
        for ii in range(wnum.size):
            idx = (np.abs(wnum[ii] - rf_wnum)).argmin()
            if np.isclose(wnum[ii], rf_wnum[idx], atol=atol):
                index_wnum.append(ii)
                index_rf.append(idx)
        mean_rad = mean_rad[:, index_wnum]
        rf = rf[index_rf]

    # Apply response function to AERI radiance
    mean_rad = mean_rad * rf

    # Sum along wavenumber dimention to get a single value for each time step
    mean_rad = np.nansum(mean_rad, axis=1)

    # Loop over each time step of the AERI data and through the use of
    # solving for zero determine the AERI equivlante IRT sky temperature.
    aeri_irt_vals = np.full(mean_rad.size, np.nan, dtype=mean_rad.dtype)

    # Look for when the hatch is not in Open position and set values to NaN.
    if hatch_name is not None:
        flag_values = aeri_ds[hatch_name].attrs['flag_values']
        flag_meanings = aeri_ds[hatch_name].attrs['flag_meanings']
        if not isinstance(flag_meanings, list):
            flag_meanings = flag_meanings.split()
        flag_meanings = [att.lower() for att in flag_meanings]
        if not isinstance(flag_values, list):
            flag_values = flag_values.split()
            flag_values = [int(att) for att in flag_values]
        value = flag_values[flag_meanings.index('open')]
        mean_rad[aeri_ds[hatch_name].values != value] = np.nan

    for ii in range(mean_rad.size):
        if np.isnan(mean_rad[ii]):
            continue
        else:
            try:
                aeri_irt_vals[ii] = brentq(
                    sum_function_irt,
                    temp_low,
                    temp_high,
                    args=(mean_rad[ii],),
                    xtol=tolerance,
                    maxiter=maxiter,
                )
            except ValueError:
                pass

    # Add new values to Xarray Dataset
    aeri_ds['aeri_irt_equiv_temperature'] = (
        'time',
        aeri_irt_vals,
        {'long_name': 'Derived IRT equivalent temperatrues from AERI', 'units': 'K'},
    )

    return aeri_ds