# Source code for qutip.control.tslotcomp

    # -*- coding: utf-8 -*-
    # @author: Alexander Pitchford
    # @email1: agp1@aber.ac.uk
    # @email2: alex.pitchford@gmail.com
    # @organization: Aberystwyth University
    # @supervisor: Daniel Burgarth

    """
    Timeslot Computer
    These classes determine which dynamics generators, propagators and evolutions
    are recalculated when there is a control amplitude update.
    The timeslot computer processes the lists held by the dynamics object

    The default (UpdateAll) updates all of these each amp update, on the
    assumption that all amplitudes are changed each iteration. This is typical
    when using optimisation methods like BFGS in the GRAPE algorithm

    The alternative (DynUpdate) assumes that only a subset of amplitudes
    are updated each iteration and attempts to minimise the number of expensive
    calculations accordingly. This would be the appropriate class for Krotov type
    methods. Note that the Stats_DynTsUpdate class must be used for stats
    in conjunction with this class.
    NOTE: AJGP 2011-10-2014: This _DynUpdate class currently has some bug,
    no pressing need to fix it presently

    If all amplitudes change at each update, then the behavior of the classes is
    equivalent. _UpdateAll is easier to understand and potentially slightly faster
    in this situation.

    Note the methods in the _DynUpdate class were inspired by:
    DYNAMO - Dynamic Framework for Quantum Optimal Control
    See Machnes et.al., arXiv.1011.4874
    """

    import warnings
    import numpy as np
    import timeit
    # QuTiP
    from qutip.qobj import Qobj
    # QuTiP control modules
    import qutip.control.errors as errors
    import qutip.control.dump as qtrldump
    # QuTiP logging
    import qutip.logging_utils as logging
    logger = logging.get_logger()

    def _func_deprecation(message, stacklevel=3):
        """
        Issue deprecation warning
        Using stacklevel=3 will ensure message refers the function
        calling with the deprecated parameter,
        """
        warnings.warn(message, DeprecationWarning, stacklevel=stacklevel)

    class TimeslotComputer(object):
        """
        Base class for all Timeslot Computers
        Note: this must be instantiated with a Dynamics object, that is the
        container for the data that the methods operate on

        Attributes
        ----------
        log_level : integer
            level of messaging output from the logger.
            Options are attributes of qutip.logging_utils,
            in decreasing levels of messaging, are:
            DEBUG_INTENSE, DEBUG_VERBOSE, DEBUG, INFO, WARN, ERROR, CRITICAL
            Anything WARN or above is effectively 'quiet' execution,
            assuming everything runs as expected.
            The default NOTSET implies that the level will be taken from
            the QuTiP settings file, which by default is WARN

        evo_comp_summary : EvoCompSummary
            A summary of the most recent evolution computation
            Used in the stats and dump
            Will be set to None if neither stats or dump are set
        """
        def __init__(self, dynamics, params=None):
            from qutip.control.dynamics import Dynamics
            if not isinstance(dynamics, Dynamics):
                raise TypeError("Must instantiate with {} type".format(
                                            Dynamics))
            self.parent = dynamics
            self.params = params
            self.reset()

        def reset(self):
            self.log_level = self.parent.log_level
            self.id_text = 'TS_COMP_BASE'
            self.evo_comp_summary = None

        def apply_params(self, params=None):
            """
            Set object attributes based on the dictionary (if any) passed in the
            instantiation, or passed as a parameter
            This is called during the instantiation automatically.
            The key value pairs are the attribute name and value
            Note: attributes are created if they do not exist already,
            and are overwritten if they do.
            """
            if not params:
                params = self.params

            if isinstance(params, dict):
                self.params = params
                for key in params:
                    setattr(self, key, params[key])

        def flag_all_calc_now(self):
            pass

        def init_comp(self):
            pass

        @property
        def log_level(self):
            return logger.level

        @log_level.setter
        def log_level(self, lvl):
            """
            Set the log_level attribute and set the level of the logger
            that is call logger.setLevel(lvl)
            """
            logger.setLevel(lvl)

        def dump_current(self):
            """Store a copy of the current time evolution"""
            dyn = self.parent
            dump = dyn.dump
            if not isinstance(dump, qtrldump.DynamicsDump):
                raise RuntimeError("Cannot dump current evolution, "
                    "as dynamics.dump is not set")

            anything_dumped = False
            item_idx = None
            if dump.dump_any:
                dump_item = dump.add_evo_dump()
                item_idx = dump_item.idx
                anything_dumped = True

            if dump.dump_summary:
                dump.add_evo_comp_summary(dump_item_idx=item_idx)
                anything_dumped = True

            if not anything_dumped:
                logger.warning("Dump set, but nothing dumped, check dump config")

    class TSlotCompUpdateAll(TimeslotComputer):
        """
        Timeslot Computer - Update All
        Updates all dynamics generators, propagators and evolutions when
        ctrl amplitudes are updated
        """
        def reset(self):
            TimeslotComputer.reset(self)
            self.id_text = 'ALL'
            self.apply_params()

        def compare_amps(self, new_amps):
            """
            Determine if any amplitudes have changed. If so, then mark the
            timeslots as needing recalculation
            Returns: True if amplitudes are the same, False if they have changed
            """
            changed = False
            dyn = self.parent

            if (dyn.stats or dyn.dump):
                if self.evo_comp_summary:
                    self.evo_comp_summary.reset()
                else:
                    self.evo_comp_summary = EvoCompSummary()
            ecs = self.evo_comp_summary

            if dyn.ctrl_amps is None:
                # Flag fidelity and gradients as needing recalculation
                changed = True
                if ecs:
                    ecs.num_amps_changed = len(new_amps.flat)
                    ecs.num_timeslots_changed = new_amps.shape[0]
            else:
                # create boolean array with same shape as ctrl_amps
                # True where value in new_amps differs, otherwise false
                changed_amps = dyn.ctrl_amps != new_amps
                if np.any(changed_amps):
                    # Flag fidelity and gradients as needing recalculation
                    changed = True
                    if self.log_level <= logging.DEBUG:
                        logger.debug("{} amplitudes changed".format(
                            changed_amps.sum()))

                    if ecs:
                        ecs.num_amps_changed = changed_amps.sum()
                        ecs.num_timeslots_changed = np.any(changed_amps, 1).sum()

                else:
                    if self.log_level <= logging.DEBUG:
                        logger.debug("No amplitudes changed")

            # *** update stats ***
            if dyn.stats:
                dyn.stats.num_ctrl_amp_updates += bool(ecs.num_amps_changed)
                dyn.stats.num_ctrl_amp_changes += ecs.num_amps_changed
                dyn.stats.num_timeslot_changes += ecs.num_timeslots_changed

            if changed:
                dyn.ctrl_amps = new_amps
                dyn.flag_system_changed()
                return False
            else:
                return True

        def recompute_evolution(self):
            """
            Recalculates the evolution operators.
            Dynamics generators (e.g. Hamiltonian) and
            prop (propagators) are calculated as necessary
            """

            dyn = self.parent
            prop_comp = dyn.prop_computer
            n_ts = dyn.num_tslots
            n_ctrls = dyn.num_ctrls

            # Clear the public lists
            # These are only set if (external) users access them
            dyn._dyn_gen_qobj = None
            dyn._prop_qobj = None
            dyn._prop_grad_qobj = None
            dyn._fwd_evo_qobj = None
            dyn._onwd_evo_qobj = None
            dyn._onto_evo_qobj = None

            if (dyn.stats or dyn.dump) and not self.evo_comp_summary:
                self.evo_comp_summary = EvoCompSummary()
            ecs = self.evo_comp_summary

            if dyn.stats is not None:
                dyn.stats.num_tslot_recompute += 1
                if self.log_level <= logging.DEBUG:
                    logger.log(logging.DEBUG, "recomputing evolution {} "
                               "(UpdateAll)".format(
                                   dyn.stats.num_tslot_recompute))

            # calculate the Hamiltonians
            if ecs: time_start = timeit.default_timer()
            for k in range(n_ts):
                dyn._combine_dyn_gen(k)
                if dyn._decomp_curr is not None:
                    dyn._decomp_curr[k] = False

            if ecs:
                ecs.wall_time_dyn_gen_compute = \
                    timeit.default_timer() - time_start

            # calculate the propagators and the propagotor gradients
            if ecs: time_start = timeit.default_timer()
            for k in range(n_ts):
                if prop_comp.grad_exact and dyn.cache_prop_grad:
                    for j in range(n_ctrls):
                        if j == 0:
                            dyn._prop[k], dyn._prop_grad[k, j] = \
                                        prop_comp._compute_prop_grad(k, j)
                            if self.log_level <= logging.DEBUG_INTENSE:
                                logger.log(logging.DEBUG_INTENSE,
                                           "propagator {}:\n{:10.3g}".format(
                                               k, self._prop[k]))
                        else:
                            dyn._prop_grad[k, j] = \
                                prop_comp._compute_prop_grad(k, j,
                                                             compute_prop=False)
                else:
                    dyn._prop[k] = prop_comp._compute_propagator(k)

            if ecs:
                ecs.wall_time_prop_compute = \
                    timeit.default_timer() - time_start

            if ecs: time_start = timeit.default_timer()
            # compute the forward propagation
            R = range(n_ts)
            for k in R:
                if dyn.oper_dtype == Qobj:
                    dyn._fwd_evo[k+1] = dyn._prop[k]*dyn._fwd_evo[k]
                else:
                    dyn._fwd_evo[k+1] = dyn._prop[k].dot(dyn._fwd_evo[k])

            if ecs:
                ecs.wall_time_fwd_prop_compute = \
                            timeit.default_timer() - time_start
                time_start = timeit.default_timer()
            # compute the onward propagation
            if dyn.fid_computer.uses_onwd_evo:
                dyn._onwd_evo[n_ts-1] = dyn._prop[n_ts-1]
                R = range(n_ts-2, -1, -1)
                for k in R:
                    if dyn.oper_dtype == Qobj:
                        dyn._onwd_evo[k] = dyn._onwd_evo[k+1]*dyn._prop[k]
                    else:
                        dyn._onwd_evo[k] = dyn._onwd_evo[k+1].dot(dyn._prop[k])

            if dyn.fid_computer.uses_onto_evo:
                #R = range(n_ts-1, -1, -1)
                R = range(n_ts-1, -1, -1)
                for k in R:
                    if dyn.oper_dtype == Qobj:
                        dyn._onto_evo[k] = dyn._onto_evo[k+1]*dyn._prop[k]
                    else:
                        dyn._onto_evo[k] = dyn._onto_evo[k+1].dot(dyn._prop[k])

            if ecs:
                ecs.wall_time_onwd_prop_compute = \
                                timeit.default_timer() - time_start

            if dyn.stats:
                dyn.stats.wall_time_dyn_gen_compute += \
                                        ecs.wall_time_dyn_gen_compute
                dyn.stats.wall_time_prop_compute += \
                                        ecs.wall_time_prop_compute
                dyn.stats.wall_time_fwd_prop_compute += \
                                        ecs.wall_time_fwd_prop_compute
                dyn.stats.wall_time_onwd_prop_compute += \
                                        ecs.wall_time_onwd_prop_compute

            if dyn.unitarity_check_level:
                dyn.check_unitarity()

            if dyn.dump:
                self.dump_current()

        def get_timeslot_for_fidelity_calc(self):
            """
            Returns the timeslot index that will be used calculate current fidelity
            value.
            This (default) method simply returns the last timeslot
            """
            _func_deprecation("'get_timeslot_for_fidelity_calc' is deprecated. "
                            "Use '_get_timeslot_for_fidelity_calc'")
            return self._get_timeslot_for_fidelity_calc

        def _get_timeslot_for_fidelity_calc(self):
            """
            Returns the timeslot index that will be used calculate current fidelity
            value.
            This (default) method simply returns the last timeslot
            """
            return self.parent.num_tslots

    class TSlotCompDynUpdate(TimeslotComputer):
        """
        Timeslot Computer - Dynamic Update
        ********************************
        ***** CURRENTLY HAS ISSUES *****
        ***** AJGP 2014-10-02
        ***** and is therefore not being maintained
        ***** i.e. changes made to _UpdateAll are not being implemented here
        ********************************
        Updates only the dynamics generators, propagators and evolutions as
        required when a subset of the ctrl amplitudes are updated.
        Will update all if all amps have changed.
        """

        def reset(self):
            self.dyn_gen_recalc = None
            self.prop_recalc = None
            self.evo_init2t_recalc = None
            self.evo_t2targ_recalc = None
            self.dyn_gen_calc_now = None
            self.prop_calc_now = None
            self.evo_init2t_calc_now = None
            self.evo_t2targ_calc_now = None
            TimeslotComputer.reset(self)
            self.id_text = 'DYNAMIC'
            self.apply_params()

        def init_comp(self):
            """
            Initialise the flags
            """
            ####
            # These maps are used to determine what needs to be updated
            ####
            # Note _recalc means the value needs updating at some point
            # e.g. here no values have been set, except the initial and final
            # evolution operator vals (which never change) and hence all other
            # values are set as requiring calculation.
            n_ts = self.parent.num_tslots
            self.dyn_gen_recalc = np.ones(n_ts, dtype=bool)
            # np.ones(n_ts, dtype=bool)
            self.prop_recalc = np.ones(n_ts, dtype=bool)
            self.evo_init2t_recalc = np.ones(n_ts + 1, dtype=bool)
            self.evo_init2t_recalc[0] = False
            self.evo_t2targ_recalc = np.ones(n_ts + 1, dtype=bool)
            self.evo_t2targ_recalc[-1] = False

            # The _calc_now map is used to during the calcs to specify
            # which values need updating immediately
            self.dyn_gen_calc_now = np.zeros(n_ts, dtype=bool)
            self.prop_calc_now = np.zeros(n_ts, dtype=bool)
            self.evo_init2t_calc_now = np.zeros(n_ts + 1, dtype=bool)
            self.evo_t2targ_calc_now = np.zeros(n_ts + 1, dtype=bool)

        def compare_amps(self, new_amps):
            """
            Determine which timeslots will have changed Hamiltonians
            i.e. any where control amplitudes have changed for that slot
            and mark (using masks) them and corresponding exponentiations and
            time evo operators for update
            Returns: True if amplitudes are the same, False if they have changed
            """
            dyn = self.parent
            n_ts = dyn.num_tslots
            # create boolean array with same shape as ctrl_amps
            # True where value in New_amps differs, otherwise false
            if self.parent.ctrl_amps is None:
                changed_amps = np.ones(new_amps.shape, dtype=bool)
            else:
                changed_amps = self.parent.ctrl_amps != new_amps

            if self.log_level <= logging.DEBUG_VERBOSE:
                logger.log(logging.DEBUG_VERBOSE, "changed_amps:\n{}".format(
                    changed_amps))
            # create Boolean vector with same length as number of timeslots
            # True where any of the amplitudes have changed, otherwise false
            changed_ts_mask = np.any(changed_amps, 1)
            # if any of the amplidudes have changed then mark for recalc
            if np.any(changed_ts_mask):
                self.dyn_gen_recalc[changed_ts_mask] = True
                self.prop_recalc[changed_ts_mask] = True
                dyn.ctrl_amps = new_amps
                if self.log_level <= logging.DEBUG:
                    logger.debug("Control amplitudes updated")
                # find first and last changed dynamics generators
                first_changed = None
                for i in range(n_ts):
                    if changed_ts_mask[i]:
                        last_changed = i
                        if first_changed is None:
                            first_changed = i

                # set all fwd evo ops after first changed Ham to be recalculated
                self.evo_init2t_recalc[first_changed + 1:] = True
                # set all bkwd evo ops up to (incl) last changed Ham to be
                # recalculated
                self.evo_t2targ_recalc[:last_changed + 1] = True

                # Flag fidelity and gradients as needing recalculation
                dyn.flag_system_changed()

                # *** update stats ***
                if dyn.stats is not None:
                    dyn.stats.num_ctrl_amp_updates += 1
                    dyn.stats.num_ctrl_amp_changes += changed_amps.sum()
                    dyn.stats.num_timeslot_changes += changed_ts_mask.sum()

                return False
            else:
                return True

        def flag_all_calc_now(self):
            """
            Flags all Hamiltonians, propagators and propagations to be
            calculated now
            """
            # set flags for calculations
            self.dyn_gen_calc_now[:] = True
            self.prop_calc_now[:] = True
            self.evo_init2t_calc_now[:-1] = True
            self.evo_t2targ_calc_now[1:] = True

        def recompute_evolution(self):
            """
            Recalculates the evo_init2t (forward) and evo_t2targ (onward) time
            evolution operators
            DynGen (Hamiltonians etc) and prop (propagator) are calculated
            as necessary
            """
            if self.log_level <= logging.DEBUG_VERBOSE:
                logger.log(logging.DEBUG_VERBOSE, "recomputing evolution "
                           "(DynUpdate)")

            dyn = self.parent
            n_ts = dyn.num_tslots
            # find the op slots that have been marked for update now
            # and need recalculation
            evo_init2t_recomp_now = self.evo_init2t_calc_now & \
                self.evo_init2t_recalc
            evo_t2targ_recomp_now = self.evo_t2targ_calc_now & \
                self.evo_t2targ_recalc

            # to recomupte evo_init2t, will need to start
            #  at a cell that has been computed
            if np.any(evo_init2t_recomp_now):
                for k in range(n_ts, 0, -1):
                    if evo_init2t_recomp_now[k] and self.evo_init2t_recalc[k-1]:
                        evo_init2t_recomp_now[k-1] = True

            # for evo_t2targ, will also need to start
            #  at a cell that has been computed
            if np.any(evo_t2targ_recomp_now):
                for k in range(0, n_ts):
                    if evo_t2targ_recomp_now[k] and self.evo_t2targ_recalc[k+1]:
                        evo_t2targ_recomp_now[k+1] = True

            # determine which dyn gen and prop need recalculating now in order to
            # calculate the forwrd and onward evolutions
            prop_recomp_now = (evo_init2t_recomp_now[1:]
                               | evo_t2targ_recomp_now[:-1]
                               | self.prop_calc_now[:]) & self.prop_recalc[:]
            dyn_gen_recomp_now = (prop_recomp_now[:] | self.dyn_gen_calc_now[:]) \
                & self.dyn_gen_recalc[:]

            if np.any(dyn_gen_recomp_now):
                time_start = timeit.default_timer()
                for k in range(n_ts):
                    if dyn_gen_recomp_now[k]:
                        # calculate the dynamics generators
                        dyn.dyn_gen[k] = dyn.compute_dyn_gen(k)
                        self.dyn_gen_recalc[k] = False
                if dyn.stats is not None:
                    dyn.stats.num_dyn_gen_computes += dyn_gen_recomp_now.sum()
                    dyn.stats.wall_time_dyn_gen_compute += \
                        timeit.default_timer() - time_start

            if np.any(prop_recomp_now):
                time_start = timeit.default_timer()
                for k in range(n_ts):
                    if prop_recomp_now[k]:
                        # calculate exp(H) and other per H computations needed for
                        # the gradient function
                        dyn.prop[k] = dyn._compute_propagator(k)
                        self.prop_recalc[k] = False
                if dyn.stats is not None:
                    dyn.stats.num_prop_computes += prop_recomp_now.sum()
                    dyn.stats.wall_time_prop_compute += \
                        timeit.default_timer() - time_start

            # compute the forward propagation
            if np.any(evo_init2t_recomp_now):
                time_start = timeit.default_timer()
                R = range(1, n_ts + 1)
                for k in R:
                    if evo_init2t_recomp_now[k]:
                        dyn.evo_init2t[k] = \
                            dyn.prop[k-1].dot(dyn.evo_init2t[k-1])
                        self.evo_init2t_recalc[k] = False
                if dyn.stats is not None:
                    dyn.stats.num_fwd_prop_step_computes += \
                        evo_init2t_recomp_now.sum()
                    dyn.stats.wall_time_fwd_prop_compute += \
                        timeit.default_timer() - time_start

            if np.any(evo_t2targ_recomp_now):
                time_start = timeit.default_timer()
                # compute the onward propagation
                R = range(n_ts-1, -1, -1)
                for k in R:
                    if evo_t2targ_recomp_now[k]:
                        dyn.evo_t2targ[k] = dyn.evo_t2targ[k+1].dot(dyn.prop[k])
                        self.evo_t2targ_recalc[k] = False
                if dyn.stats is not None:
                    dyn.stats.num_onwd_prop_step_computes += \
                        evo_t2targ_recomp_now.sum()
                    dyn.stats.wall_time_onwd_prop_compute += \
                        timeit.default_timer() - time_start

            # Clear calc now flags
            self.dyn_gen_calc_now[:] = False
            self.prop_calc_now[:] = False
            self.evo_init2t_calc_now[:] = False
            self.evo_t2targ_calc_now[:] = False

        def get_timeslot_for_fidelity_calc(self):
            """
            Returns the timeslot index that will be used calculate current fidelity
            value. Attempts to find a timeslot where the least number of propagator
            calculations will be required.
            Flags the associated evolution operators for calculation now
            """
            dyn = self.parent
            n_ts = dyn.num_tslots
            kBothEvoCurrent = -1
            kFwdEvoCurrent = -1
            kUse = -1
            # If no specific timeslot set in config, then determine dynamically
            if kUse < 0:
                for k in range(n_ts):
                    # find first timeslot where both evo_init2t and
                    # evo_t2targ are current
                    if not self.evo_init2t_recalc[k]:
                        kFwdEvoCurrent = k
                        if not self.evo_t2targ_recalc[k]:
                            kBothEvoCurrent = k
                            break

                if kBothEvoCurrent >= 0:
                    kUse = kBothEvoCurrent
                elif kFwdEvoCurrent >= 0:
                    kUse = kFwdEvoCurrent
                else:
                    raise errors.FunctionalError("No timeslot found matching "
                                                 "criteria")

            self.evo_init2t_calc_now[kUse] = True
            self.evo_t2targ_calc_now[kUse] = True
            return kUse

    class EvoCompSummary(qtrldump.DumpSummaryItem):
        """
        A summary of the most recent time evolution computation
        Used in stats calculations and for data dumping

        Attributes
        ----------
        evo_dump_idx : int
            Index of the linked :class:`dump.EvoCompDumpItem`
            None if no linked item

        iter_num : int
            Iteration number of the pulse optimisation
            None if evolution compute outside of a pulse optimisation

        fid_func_call_num : int
            Fidelity function call number of the pulse optimisation
            None if evolution compute outside of a pulse optimisation

        grad_func_call_num : int
            Gradient function call number of the pulse optimisation
            None if evolution compute outside of a pulse optimisation

        num_amps_changed : int
            Number of control timeslot amplitudes changed since previous
            evolution calculation

        num_timeslots_changed : int
            Number of timeslots in which any amplitudes changed since previous
            evolution calculation

        wall_time_dyn_gen_compute : float
            Time spent computing dynamics generators
            (in seconds of elapsed time)

        wall_time_prop_compute : float
            Time spent computing propagators (including and propagator gradients)
            (in seconds of elapsed time)

        wall_time_fwd_prop_compute : float
            Time spent computing the forward evolution of the system
            see :property:`dynamics.fwd_evo`
            (in seconds of elapsed time)

        wall_time_onwd_prop_compute : float
            Time spent computing the 'backward' evolution of the system
            see :property:`dynamics.onwd_evo` and :property:`dynamics.onto_evo`
            (in seconds of elapsed time)
        """

        min_col_width = 11
        summary_property_names = (
            "idx", "evo_dump_idx",
            "iter_num", "fid_func_call_num", "grad_func_call_num",
            "num_amps_changed", "num_timeslots_changed",
            "wall_time_dyn_gen_compute", "wall_time_prop_compute",
            "wall_time_fwd_prop_compute", "wall_time_onwd_prop_compute")

        summary_property_fmt_type = (
            'd', 'd',
            'd', 'd', 'd',
            'd', 'd',
            'g', 'g',
            'g', 'g'
            )

        summary_property_fmt_prec = (
            0, 0,
            0, 0, 0,
            0, 0,
            3, 3,
            3, 3
            )

        def __init__(self):
            self.reset()

        def reset(self):
            qtrldump.DumpSummaryItem.reset(self)
            self.evo_dump_idx = None
            self.iter_num = None
            self.fid_func_call_num = None
            self.grad_func_call_num = None
            self.num_amps_changed = 0
            self.num_timeslots_changed = 0
            self.wall_time_dyn_gen_compute = 0.0
            self.wall_time_prop_compute = 0.0
            self.wall_time_fwd_prop_compute = 0.0
            self.wall_time_onwd_prop_compute = 0.0