# Source code for qutip.control.fidcomp

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

    """
    Fidelity Computer

    These classes calculate the fidelity error - function to be minimised
    and fidelity error gradient, which is used to direct the optimisation

    They may calculate the fidelity as an intermediary step, as in some case
    e.g. unitary dynamics, this is more efficient

    The idea is that different methods for computing the fidelity can be tried
    and compared using simple configuration switches.

    Note the methods in these classes were inspired by:
    DYNAMO - Dynamic Framework for Quantum Optimal Control
    See Machnes et.al., arXiv.1011.4874
    The unitary dynamics fidelity is taken directly frm DYNAMO
    The other fidelity measures are extensions, and the sources are given
    in the class descriptions.
    """

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

    warnings.simplefilter('always', DeprecationWarning) #turn off filter
    def _attrib_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)

    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)

    def _trace(A):
        """wrapper for calculating the trace"""
        # input is an operator (Qobj, array, sparse etc), so
        if isinstance(A, Qobj):
            return A.tr()
        elif isinstance(A, np.ndarray):
            return np.trace(A)
        else:
            #Assume A some sparse matrix
            return np.sum(A.diagonal())

    class FidelityComputer(object):
        """
        Base class for all Fidelity Computers.
        This cannot be used directly. See subclass descriptions and choose
        one appropriate for the application
        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

        dimensional_norm : float
            Normalisation constant

        fid_norm_func : function
            Used to normalise the fidelity
            See SU and PSU options for the unitary dynamics

        grad_norm_func : function
            Used to normalise the fidelity gradient
            See SU and PSU options for the unitary dynamics

        uses_onwd_evo : boolean
            flag to specify whether the onwd_evo evolution operator
            (see Dynamics) is used by the FidelityComputer

        uses_onto_evo : boolean
            flag to specify whether the onto_evo evolution operator
             (see Dynamics) is used by the FidelityComputer

        fid_err : float
            Last computed value of the fidelity error

        fidelity : float
            Last computed value of the normalised fidelity

        fidelity_current : boolean
            flag to specify whether the fidelity / fid_err are based on the
            current amplitude values. Set False when amplitudes change

        fid_err_grad: array[num_tslot, num_ctrls] of float
            Last computed values for the fidelity error gradients wrt the
            control in the timeslot

        grad_norm : float
            Last computed value for the norm of the fidelity error gradients
            (sqrt of the sum of the squares)

        fid_err_grad_current : boolean
            flag to specify whether the fidelity / fid_err are based on the
            current amplitude values. Set False when amplitudes change
        """
        def __init__(self, dynamics, params=None):
            self.parent = dynamics
            self.params = params
            self.reset()

        def reset(self):
            """
            reset any configuration data and
            clear any temporarily held status data
            """
            self.log_level = self.parent.log_level
            self.id_text = 'FID_COMP_BASE'
            self.dimensional_norm = 1.0
            self.fid_norm_func = None
            self.grad_norm_func = None
            self.uses_onwd_evo = False
            self.uses_onto_evo = False
            self.apply_params()
            self.clear()

        def clear(self):
            """
            clear any temporarily held status data
            """
            self.fid_err = None
            self.fidelity = None
            self.fid_err_grad = None
            self.grad_norm = np.inf
            self.fidelity_current = False
            self.fid_err_grad_current = False
            self.grad_norm = 0.0

        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])

        @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 init_comp(self):
            """
            initialises the computer based on the configuration of the Dynamics
            """
            # optionally implemented in subclass
            pass

        def get_fid_err(self):
            """
            returns the absolute distance from the maximum achievable fidelity
            """
            # must be implemented by subclass
            raise errors.UsageError(
                "No method defined for getting fidelity error."
                " Suspect base class was used where sub class should have been")

        def get_fid_err_gradient(self):
            """
            Returns the normalised gradient of the fidelity error
            in a (nTimeslots x n_ctrls) array wrt the timeslot control amplitude
            """
            # must be implemented by subclass
            raise errors.UsageError("No method defined for getting fidelity"
                                    " error gradient. Suspect base class was"
                                    " used where sub class should have been")

        def flag_system_changed(self):
            """
            Flag fidelity and gradients as needing recalculation
            """
            self.fidelity_current = False
            # Flag gradient as needing recalculating
            self.fid_err_grad_current = False

        @property
        def uses_evo_t2end(self):
            _attrib_deprecation(
                    "'uses_evo_t2end' has been replaced by 'uses_onwd_evo'")
            return self.uses_onwd_evo

        @uses_evo_t2end.setter
        def uses_evo_t2end(self, value):
            _attrib_deprecation(
                    "'uses_evo_t2end' has been replaced by 'uses_onwd_evo'")
            self.uses_onwd_evo = value

        @property
        def uses_evo_t2targ(self):
            _attrib_deprecation(
                    "'uses_evo_t2targ' has been replaced by 'uses_onto_evo'")
            return self.uses_onto_evo

        @uses_evo_t2targ.setter
        def uses_evo_t2targ(self, value):
            _attrib_deprecation(
                    "'uses_evo_t2targ' has been replaced by 'uses_onto_evo'")
            self.uses_onto_evo = value

    class FidCompUnitary(FidelityComputer):
        """
        Computes fidelity error and gradient assuming unitary dynamics, e.g.
        closed qubit systems
        Note fidelity and gradient calculations were taken from DYNAMO
        (see file header)

        Attributes
        ----------
        phase_option : string
            determines how global phase is treated in fidelity calculations:
                PSU - global phase ignored
                SU - global phase included

        fidelity_prenorm : complex
            Last computed value of the fidelity before it is normalised
            It is stored to use in the gradient normalisation calculation

        fidelity_prenorm_current : boolean
            flag to specify whether fidelity_prenorm are based on the
            current amplitude values. Set False when amplitudes change
        """

        def reset(self):
            FidelityComputer.reset(self)
            self.id_text = 'UNIT'
            self.uses_onto_evo = True
            self._init_phase_option('PSU')
            self.apply_params()

        def clear(self):
            FidelityComputer.clear(self)
            self.fidelity_prenorm = None
            self.fidelity_prenorm_current = False

        def set_phase_option(self, phase_option=None):
            """
            Deprecated - use phase_option
            Phase options are
            SU - global phase important
            PSU - global phase is not important
            """
            _func_deprecation("'set_phase_option' is deprecated. "
                                "Use phase_option property")
            self._init_phase_option(phase_option)

        @property
        def phase_option(self):
            return self._phase_option

        @phase_option.setter
        def phase_option(self, value):
            """
            # Phase options are
            #  SU - global phase important
            #  PSU - global phase is not important
            """
            self._init_phase_option(value)

        def _init_phase_option(self, value):
            self._phase_option = value
            if value == 'PSU':
                self.fid_norm_func = self.normalize_PSU
                self.grad_norm_func = self.normalize_gradient_PSU
            elif value == 'SU':
                self.fid_norm_func = self.normalize_SU
                self.grad_norm_func = self.normalize_gradient_SU
            elif value is None:
                raise errors.UsageError("phase_option cannot be set to None"
                                        " for this FidelityComputer.")
            else:
                raise errors.UsageError(
                        "No option for phase_option '{}'".format(value))

        def init_comp(self):
            """
            Check configuration and initialise the normalisation
            """
            if self.fid_norm_func is None or self.grad_norm_func is None:
                raise errors.UsageError("The phase_option must be be set"
                                        "for this fidelity computer")
            self.init_normalization()

        def flag_system_changed(self):
            """
            Flag fidelity and gradients as needing recalculation
            """
            FidelityComputer.flag_system_changed(self)
            # Flag the fidelity (prenormalisation) value as needing calculation
            self.fidelity_prenorm_current = False

        def init_normalization(self):
            """
            Calc norm of <Ufinal | Ufinal> to scale subsequent norms
            When considering unitary time evolution operators, this basically
            results in calculating the trace of the identity matrix
            and is hence equal to the size of the target matrix
            There may be situations where this is not the case, and hence it
            is not assumed to be so.
            The normalisation function called should be set to either the
            PSU - global phase ignored
            SU  - global phase respected
            """
            dyn = self.parent
            self.dimensional_norm = 1.0
            self.dimensional_norm = \
                self.fid_norm_func(dyn.target.dag()*dyn.target)

        def normalize_SU(self, A):
            """

            """
            try:
                if A.shape[0] == A.shape[1]:
                   # input is an operator (Qobj, array, sparse etc), so
                    norm = _trace(A)
                else:
                    raise TypeError("Cannot compute trace (not square)")
            except:
                # assume input is already scalar and hence assumed
                # to be the prenormalised scalar value, e.g. fidelity
                norm = A
            return np.real(norm) / self.dimensional_norm

        def normalize_gradient_SU(self, grad):
            """
            Normalise the gradient matrix passed as grad
            This SU version respects global phase
            """
            grad_normalized = np.real(grad) / self.dimensional_norm

            return grad_normalized

        def normalize_PSU(self, A):
            """

            """
            try:
                if A.shape[0] == A.shape[1]:
                   # input is an operator (Qobj, array, sparse etc), so
                    norm = _trace(A)
                else:
                    raise TypeError("Cannot compute trace (not square)")
            except:
                # assume input is already scalar and hence assumed
                # to be the prenormalised scalar value, e.g. fidelity
                norm = A
            return np.abs(norm) / self.dimensional_norm

        def normalize_gradient_PSU(self, grad):
            """
            Normalise the gradient matrix passed as grad
            This PSU version is independent of global phase
            """
            fid_pn = self.get_fidelity_prenorm()
            grad_normalized = np.real(grad * np.exp(-1j * np.angle(fid_pn)) /
                                      self.dimensional_norm)
            return grad_normalized

        def get_fid_err(self):
            """
            Gets the absolute error in the fidelity
            """
            return np.abs(1 - self.get_fidelity())

        def get_fidelity(self):
            """
            Gets the appropriately normalised fidelity value
            The normalisation is determined by the fid_norm_func pointer
            which should be set in the config
            """
            if not self.fidelity_current:
                self.fidelity = \
                    self.fid_norm_func(self.get_fidelity_prenorm())
                self.fidelity_current = True
                if self.log_level <= logging.DEBUG:
                    logger.debug("Fidelity (normalised): {}".format(self.fidelity))

            return self.fidelity

        def get_fidelity_prenorm(self):
            """
            Gets the current fidelity value prior to normalisation
            Note the gradient function uses this value
            The value is cached, because it is used in the gradient calculation
            """
            if not self.fidelity_prenorm_current:
                dyn = self.parent
                k = dyn.tslot_computer._get_timeslot_for_fidelity_calc()
                dyn.compute_evolution()
                if dyn.oper_dtype == Qobj:
                    f = (dyn._onto_evo[k]*dyn._fwd_evo[k]).tr()
                else:
                    f = _trace(dyn._onto_evo[k].dot(dyn._fwd_evo[k]))
                self.fidelity_prenorm = f
                self.fidelity_prenorm_current = True
                if dyn.stats is not None:
                        dyn.stats.num_fidelity_computes += 1
                if self.log_level <= logging.DEBUG:
                    logger.debug("Fidelity (pre normalisation): {}".format(
                        self.fidelity_prenorm))
            return self.fidelity_prenorm

        def get_fid_err_gradient(self):
            """
            Returns the normalised gradient of the fidelity error
            in a (nTimeslots x n_ctrls) array
            The gradients are cached in case they are requested
            mutliple times between control updates
            (although this is not typically found to happen)
            """
            if not self.fid_err_grad_current:
                dyn = self.parent
                grad_prenorm = self.compute_fid_grad()
                if self.log_level <= logging.DEBUG_INTENSE:
                    logger.log(logging.DEBUG_INTENSE, "pre-normalised fidelity "
                               "gradients:\n{}".format(grad_prenorm))
                # AJGP: Note this check should not be necessary if dynamics are
                #       unitary. However, if they are not then this gradient
                #       can still be used, however the interpretation is dubious
                if self.get_fidelity() >= 1:
                    self.fid_err_grad = self.grad_norm_func(grad_prenorm)
                else:
                    self.fid_err_grad = -self.grad_norm_func(grad_prenorm)

                self.fid_err_grad_current = True
                if dyn.stats is not None:
                    dyn.stats.num_grad_computes += 1

                self.grad_norm = np.sqrt(np.sum(self.fid_err_grad**2))
                if self.log_level <= logging.DEBUG_INTENSE:
                    logger.log(logging.DEBUG_INTENSE, "Normalised fidelity error "
                               "gradients:\n{}".format(self.fid_err_grad))

                if self.log_level <= logging.DEBUG:
                    logger.debug("Gradient (sum sq norm): "
                                 "{} ".format(self.grad_norm))

            return self.fid_err_grad

        def compute_fid_grad(self):
            """
            Calculates exact gradient of function wrt to each timeslot
            control amplitudes. Note these gradients are not normalised
            These are returned as a (nTimeslots x n_ctrls) array
            """
            dyn = self.parent
            n_ctrls = dyn.num_ctrls
            n_ts = dyn.num_tslots

            # create n_ts x n_ctrls zero array for grad start point
            grad = np.zeros([n_ts, n_ctrls], dtype=complex)

            dyn.tslot_computer.flag_all_calc_now()
            dyn.compute_evolution()

            # loop through all ctrl timeslots calculating gradients
            time_st = timeit.default_timer()
            for j in range(n_ctrls):
                for k in range(n_ts):
                    fwd_evo = dyn._fwd_evo[k]
                    onto_evo = dyn._onto_evo[k+1]
                    if dyn.oper_dtype == Qobj:
                        g = (onto_evo*dyn._get_prop_grad(k, j)*fwd_evo).tr()
                    else:
                        g = _trace(onto_evo.dot(
                                    dyn._get_prop_grad(k, j)).dot(fwd_evo))
                    grad[k, j] = g
            if dyn.stats is not None:
                dyn.stats.wall_time_gradient_compute += \
                    timeit.default_timer() - time_st
            return grad

    class FidCompTraceDiff(FidelityComputer):
        """
        Computes fidelity error and gradient for general system dynamics
        by calculating the the fidelity error as the trace of the overlap
        of the difference between the target and evolution resulting from
        the pulses with the transpose of the same.
        This should provide a distance measure for dynamics described by matrices
        Note the gradient calculation is taken from:
        'Robust quantum gates for open systems via optimal control:
        Markovian versus non-Markovian dynamics'
        Frederik F Floether, Pierre de Fouquieres, and Sophie G Schirmer

        Attributes
        ----------
        scale_factor : float
            The fidelity error calculated is of some arbitary scale. This
            factor can be used to scale the fidelity error such that it may
            represent some physical measure
            If None is given then it is caculated as 1/2N, where N
            is the dimension of the drift, when the Dynamics are initialised.
        """

        def reset(self):
            FidelityComputer.reset(self)
            self.id_text = 'TRACEDIFF'
            self.scale_factor = None
            self.uses_onwd_evo = True
            if not self.parent.prop_computer.grad_exact:
                raise errors.UsageError(
                    "This FidelityComputer can only be"
                    " used with an exact gradient PropagatorComputer.")
            self.apply_params()

        def init_comp(self):
            """
            initialises the computer based on the configuration of the Dynamics
            Calculates the scale_factor is not already set
            """
            if self.scale_factor is None:
                self.scale_factor = 1.0 / (2.0*self.parent.get_drift_dim())
                if self.log_level <= logging.DEBUG:
                    logger.debug("Scale factor calculated as {}".format(
                        self.scale_factor))

        def get_fid_err(self):
            """
            Gets the absolute error in the fidelity
            """
            if not self.fidelity_current:
                dyn = self.parent
                dyn.compute_evolution()
                n_ts = dyn.num_tslots
                evo_final = dyn._fwd_evo[n_ts]
                evo_f_diff = dyn._target - evo_final
                if self.log_level <= logging.DEBUG_VERBOSE:
                    logger.log(logging.DEBUG_VERBOSE, "Calculating TraceDiff "
                               "fidelity...\n Target:\n{}\n Evo final:\n{}\n"
                               "Evo final diff:\n{}".format(dyn._target, evo_final,
                                                            evo_f_diff))

                # Calculate the fidelity error using the trace difference norm
                # Note that the value should have not imagnary part, so using
                # np.real, just avoids the complex casting warning
                if dyn.oper_dtype == Qobj:
                    self.fid_err = self.scale_factor*np.real(
                            (evo_f_diff.dag()*evo_f_diff).tr())
                else:
                    self.fid_err = self.scale_factor*np.real(_trace(
                            evo_f_diff.conj().T.dot(evo_f_diff)))

                if np.isnan(self.fid_err):
                    self.fid_err = np.Inf

                if dyn.stats is not None:
                        dyn.stats.num_fidelity_computes += 1

                self.fidelity_current = True
                if self.log_level <= logging.DEBUG:
                    logger.debug("Fidelity error: {}".format(self.fid_err))

            return self.fid_err

        def get_fid_err_gradient(self):
            """
            Returns the normalised gradient of the fidelity error
            in a (nTimeslots x n_ctrls) array
            The gradients are cached in case they are requested
            mutliple times between control updates
            (although this is not typically found to happen)
            """
            if not self.fid_err_grad_current:
                dyn = self.parent
                self.fid_err_grad = self.compute_fid_err_grad()
                self.fid_err_grad_current = True
                if dyn.stats is not None:
                    dyn.stats.num_grad_computes += 1

                self.grad_norm = np.sqrt(np.sum(self.fid_err_grad**2))
                if self.log_level <= logging.DEBUG_INTENSE:
                    logger.log(logging.DEBUG_INTENSE, "fidelity error gradients:\n"
                               "{}".format(self.fid_err_grad))

                if self.log_level <= logging.DEBUG:
                    logger.debug("Gradient norm: "
                                 "{} ".format(self.grad_norm))

            return self.fid_err_grad

        def compute_fid_err_grad(self):
            """
            Calculate exact gradient of the fidelity error function
            wrt to each timeslot control amplitudes.
            Uses the trace difference norm fidelity
            These are returned as a (nTimeslots x n_ctrls) array
            """
            dyn = self.parent
            n_ctrls = dyn.num_ctrls
            n_ts = dyn.num_tslots

            # create n_ts x n_ctrls zero array for grad start point
            grad = np.zeros([n_ts, n_ctrls])

            dyn.tslot_computer.flag_all_calc_now()
            dyn.compute_evolution()

            # loop through all ctrl timeslots calculating gradients
            time_st = timeit.default_timer()

            evo_final = dyn._fwd_evo[n_ts]
            evo_f_diff = dyn._target - evo_final
            for j in range(n_ctrls):
                for k in range(n_ts):
                    fwd_evo = dyn._fwd_evo[k]
                    if dyn.oper_dtype == Qobj:
                        evo_grad = dyn._get_prop_grad(k, j)*fwd_evo
                        if k+1 < n_ts:
                            evo_grad = dyn._onwd_evo[k+1]*evo_grad
                        # Note that the value should have not imagnary part, so
                        # using np.real, just avoids the complex casting warning
                        g = -2*self.scale_factor*np.real(
                                        (evo_f_diff.dag()*evo_grad).tr())
                    else:
                        evo_grad = dyn._get_prop_grad(k, j).dot(fwd_evo)
                        if k+1 < n_ts:
                            evo_grad = dyn._onwd_evo[k+1].dot(evo_grad)
                        g = -2*self.scale_factor*np.real(_trace(
                                        evo_f_diff.conj().T.dot(evo_grad)))
                    if np.isnan(g):
                        g = np.Inf

                    grad[k, j] = g
            if dyn.stats is not None:
                dyn.stats.wall_time_gradient_compute += \
                    timeit.default_timer() - time_st
            return grad

    class FidCompTraceDiffApprox(FidCompTraceDiff):
        """
        As FidCompTraceDiff, except uses the finite difference method to
        compute approximate gradients

        Attributes
        ----------
        epsilon : float
            control amplitude offset to use when approximating the gradient wrt
            a timeslot control amplitude

        """

        def reset(self):
            FidelityComputer.reset(self)
            self.id_text = 'TDAPPROX'
            self.uses_onwd_evo = True
            self.scale_factor = None
            self.epsilon = 0.001
            self.apply_params()

        def compute_fid_err_grad(self):
            """
            Calculates gradient of function wrt to each timeslot
            control amplitudes. Note these gradients are not normalised
            They are calulated
            These are returned as a (nTimeslots x n_ctrls) array
            """
            dyn = self.parent
            prop_comp = dyn.prop_computer
            n_ctrls = dyn.num_ctrls
            n_ts = dyn.num_tslots

            if self.log_level >= logging.DEBUG:
                logger.debug("Computing fidelity error gradient")
            # create n_ts x n_ctrls zero array for grad start point
            grad = np.zeros([n_ts, n_ctrls])

            dyn.tslot_computer.flag_all_calc_now()
            dyn.compute_evolution()
            curr_fid_err = self.get_fid_err()

            # loop through all ctrl timeslots calculating gradients
            time_st = timeit.default_timer()

            for j in range(n_ctrls):
                for k in range(n_ts):
                    fwd_evo = dyn._fwd_evo[k]
                    prop_eps = prop_comp._compute_diff_prop(k, j, self.epsilon)
                    if dyn.oper_dtype == Qobj:
                        evo_final_eps = fwd_evo*prop_eps
                        if k+1 < n_ts:
                            evo_final_eps = evo_final_eps*dyn._onwd_evo[k+1]
                        evo_f_diff_eps = dyn._target - evo_final_eps
                        # Note that the value should have not imagnary part, so
                        # using np.real, just avoids the complex casting warning
                        fid_err_eps = self.scale_factor*np.real(
                            (evo_f_diff_eps.dag()*evo_f_diff_eps).tr())
                    else:
                        evo_final_eps = fwd_evo.dot(prop_eps)
                        if k+1 < n_ts:
                            evo_final_eps = evo_final_eps.dot(dyn._onwd_evo[k+1])
                        evo_f_diff_eps = dyn._target - evo_final_eps
                        fid_err_eps = self.scale_factor*np.real(_trace(
                            evo_f_diff_eps.conj().T.dot(evo_f_diff_eps)))

                    g = (fid_err_eps - curr_fid_err)/self.epsilon
                    if np.isnan(g):
                        g = np.Inf

                    grad[k, j] = g

            if dyn.stats is not None:
                dyn.stats.wall_time_gradient_compute += \
                    timeit.default_timer() - time_st

            return grad