# Source code for qutip.control.stats

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

    """
    Statistics for the optimisation
    Note that some of the stats here are redundant copies from the optimiser
    used here for calculations
    """
    import numpy as np
    import datetime

    class Stats(object):
        """
        Base class for all optimisation statistics
        Used for configurations where all timeslots are updated each iteration
        e.g. exact gradients
        Note that all times are generated using timeit.default_timer() and are
        in seconds

        Attributes
        ----------
        dyn_gen_name : string
            Text used in some report functions.
            Makes sense to set it to 'Hamiltonian' when using unitary dynamics
            Default is simply 'dynamics generator'

        num_iter : integer
            Number of iterations of the optimisation algorithm

        wall_time_optim_start : float
            Start time for the optimisation

        wall_time_optim_end : float
            End time for the optimisation

        wall_time_optim : float
            Time elasped during the optimisation

        wall_time_dyn_gen_compute : float
            Total wall (elasped) time computing combined dynamics generator
            (for example combining drift and control Hamiltonians)

        wall_time_prop_compute : float
            Total wall (elasped) time computing propagators, that is the
            time evolution from one timeslot to the next
            Includes calculating the propagator gradient for exact gradients

        wall_time_fwd_prop_compute : float
            Total wall (elasped) time computing combined forward propagation,
            that is the time evolution from the start to a specific timeslot.
            Excludes calculating the propagators themselves

        wall_time_onwd_prop_compute : float
            Total wall (elasped) time computing combined onward propagation,
            that is the time evolution from a specific timeslot to the end time.
            Excludes calculating the propagators themselves

        wall_time_gradient_compute : float
            Total wall (elasped) time computing the fidelity error gradient.
            Excludes calculating the propagator gradients (in exact gradient
            methods)

        num_fidelity_func_calls : integer
            Number of calls to fidelity function by the optimisation algorithm

        num_grad_func_calls : integer
            Number of calls to gradient function by the optimisation algorithm

        num_tslot_recompute : integer
            Number of time the timeslot evolution is recomputed
            (It is only computed if any amplitudes changed since the last call)

        num_fidelity_computes : integer
            Number of time the fidelity is computed
            (It is only computed if any amplitudes changed since the last call)

        num_grad_computes : integer
            Number of time the gradient is computed
            (It is only computed if any amplitudes changed since the last call)

        num_ctrl_amp_updates : integer
            Number of times the control amplitudes are updated

        mean_num_ctrl_amp_updates_per_iter : float
            Mean number of control amplitude updates per iteration

        num_timeslot_changes : integer
            Number of times the amplitudes of a any control in a timeslot changes

        mean_num_timeslot_changes_per_update : float
            Mean average number of timeslot amplitudes that are changed per update

        num_ctrl_amp_changes : integer
            Number of times individual control amplitudes that are changed

        mean_num_ctrl_amp_changes_per_update : float
            Mean average number of control amplitudes that are changed per update
        """
        def __init__(self):
            self.reset()

        def reset(self):
            self.dyn_gen_name = "dynamics generator"
            self.clear()

        def clear(self):
            self.num_iter = 0
            # Duration attributes
            self.wall_time_optim_start = 0.0
            self.wall_time_optim_end = 0.0
            self.wall_time_optim = 0.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
            self.wall_time_gradient_compute = 0.0
            # Fidelity and gradient function calls and computes
            self.num_fidelity_func_calls = 0
            self.num_grad_func_calls = 0
            self.num_tslot_recompute = 0
            self.num_fidelity_computes = 0
            self.num_grad_computes = 0
            # Control amplitudes
            self.num_ctrl_amp_updates = 0
            self.mean_num_ctrl_amp_updates_per_iter = 0.0
            self.num_timeslot_changes = 0
            self.mean_num_timeslot_changes_per_update = 0.0
            self.num_ctrl_amp_changes = 0
            self.mean_num_ctrl_amp_changes_per_update = 0.0

        def calculate(self):
            """
            Perform the calculations (e.g. averages) that are required on the stats
            Should be called before calling report
            """
            # If the optimation is still running then the optimisation
            # time is the time so far
            if self.wall_time_optim_end > 0.0:
                self.wall_time_optim = \
                    self.wall_time_optim_end - self.wall_time_optim_start

            try:
                self.mean_num_ctrl_amp_updates_per_iter = \
                    self.num_ctrl_amp_updates / float(self.num_iter)
            except:
                self.mean_num_ctrl_amp_updates_per_iter = np.NaN

            try:
                self.mean_num_timeslot_changes_per_update = \
                    self.num_timeslot_changes / float(self.num_ctrl_amp_updates)
            except:
                self.mean_num_timeslot_changes_per_update = np.NaN

            try:
                self.mean_num_ctrl_amp_changes_per_update = \
                    self.num_ctrl_amp_changes / float(self.num_ctrl_amp_updates)
            except:
                self.mean_num_ctrl_amp_changes_per_update = np.NaN

        def _format_datetime(self, t, tot=0.0):
            dtStr = str(datetime.timedelta(seconds=t))
            if tot > 0:
                percent = 100*t/tot
                dtStr += " ({:03.2f}%)".format(percent)
            return dtStr

        def report(self):
            """
            Print a report of the stats to the console
            """
            print("\n------------------------------------"
                  "\n---- Control optimisation stats ----")
            self.report_timings()
            self.report_func_calls()
            self.report_amp_updates()
            print("------------------------------------")

        def report_timings(self):
            print("**** Timings (HH:MM:SS.US) ****")
            tot = self.wall_time_optim
            print("Total wall time elapsed during optimisation: " +
                  self._format_datetime(tot))
            print("Wall time computing Hamiltonians: " +
                  self._format_datetime(self.wall_time_dyn_gen_compute, tot))
            print("Wall time computing propagators: " +
                  self._format_datetime(self.wall_time_prop_compute, tot))
            print("Wall time computing forward propagation: " +
                  self._format_datetime(self.wall_time_fwd_prop_compute, tot))
            print("Wall time computing onward propagation: " +
                  self._format_datetime(self.wall_time_onwd_prop_compute, tot))
            print("Wall time computing gradient: " +
                  self._format_datetime(self.wall_time_gradient_compute, tot))
            print("")

        def report_func_calls(self):
            print("**** Iterations and function calls ****")
            print("Number of iterations: {}".format(self.num_iter))
            print("Number of fidelity function calls: "
                  "{}".format(self.num_fidelity_func_calls))
            print("Number of times fidelity is computed: "
                  "{}".format(self.num_fidelity_computes))
            print("Number of gradient function calls: "
                  "{}".format(self.num_grad_func_calls))
            print("Number of times gradients are computed: "
                  "{}".format(self.num_grad_computes))
            print("Number of times timeslot evolution is recomputed: "
                  "{}".format(self.num_tslot_recompute))
            print("")

        def report_amp_updates(self):
            print("**** Control amplitudes ****")
            print("Number of control amplitude updates: "
                  "{}".format(self.num_ctrl_amp_updates))
            print("Mean number of updates per iteration: "
                  "{}".format(self.mean_num_ctrl_amp_updates_per_iter))
            print("Number of timeslot values changed: "
                  "{}".format(self.num_timeslot_changes))
            print("Mean number of timeslot changes per update: "
                  "{}".format(self.mean_num_timeslot_changes_per_update))
            print("Number of amplitude values changed: "
                  "{}".format(self.num_ctrl_amp_changes))
            print("Mean number of amplitude changes per update: "
                  "{}".format(self.mean_num_ctrl_amp_changes_per_update))

    class StatsDynTsUpdate(Stats):
        """
        Optimisation stats class for configurations where all timeslots are not
        necessarily updated at each iteration. In this case it may be interesting
        to know how many Hamiltions etc are computed each ctrl amplitude update

        Attributes
        ----------
        num_dyn_gen_computes : integer
            Total number of dynamics generator (Hamiltonian) computations,
            that is combining drift and control dynamics to calculate the
            combined dynamics generator for the timeslot

        mean_num_dyn_gen_computes_per_update : float
            # Mean average number of dynamics generator computations per update

        mean_wall_time_dyn_gen_compute : float
            # Mean average time to compute a timeslot dynamics generator

        num_prop_computes : integer
            Total number of propagator (and propagator gradient for exact
            gradient types) computations

        mean_num_prop_computes_per_update : float
            Mean average number of propagator computations per update

        mean_wall_time_prop_compute : float
            Mean average time to compute a propagator (and its gradient)

        num_fwd_prop_step_computes : integer
            Total number of steps (matrix product) computing forward propagation

        mean_num_fwd_prop_step_computes_per_update : float
            Mean average number of steps computing forward propagation

        mean_wall_time_fwd_prop_compute : float
            Mean average time to compute forward propagation

        num_onwd_prop_step_computes : integer
            Total number of steps (matrix product) computing onward propagation

        mean_num_onwd_prop_step_computes_per_update : float
            Mean average number of steps computing onward propagation

        mean_wall_time_onwd_prop_compute
            Mean average time to compute onward propagation
        """

        def __init__(self):
            self.reset()

        def reset(self):
            Stats.reset(self)
            # Dynamics generators (Hamiltonians)
            self.num_dyn_gen_computes = 0
            self.mean_num_dyn_gen_computes_per_update = 0.0
            self.mean_wall_time_dyn_gen_compute = 0.0
            # **** Propagators *****
            self.num_prop_computes = 0
            self.mean_num_prop_computes_per_update = 0.0
            self.mean_wall_time_prop_compute = 0.0
            # **** Forward propagation ****
            self.num_fwd_prop_step_computes = 0
            self.mean_num_fwd_prop_step_computes_per_update = 0.0
            self.mean_wall_time_fwd_prop_compute = 0.0
            # **** onward propagation ****
            self.num_onwd_prop_step_computes = 0
            self.mean_num_onwd_prop_step_computes_per_update = 0.0
            self.mean_wall_time_onwd_prop_compute = 0.0

        def calculate(self):
            Stats.calculate(self)
            self.mean_num_dyn_gen_computes_per_update = \
                self.num_dyn_gen_computes / float(self.num_ctrl_amp_updates)

            self.mean_wall_time_dyn_gen_compute = \
                (self.wall_time_dyn_gen_compute /
                 float(self.num_dyn_gen_computes))

            self.mean_num_prop_computes_per_update = \
                self.num_prop_computes / float(self.num_ctrl_amp_updates)

            self.mean_wall_time_prop_compute = \
                self.wall_time_prop_compute / float(self.num_prop_computes)

            self.mean_num_fwd_prop_step_computes_per_update = \
                (self.num_fwd_prop_step_computes /
                 float(self.num_ctrl_amp_updates))

            self.mean_wall_time_fwd_prop_compute = \
                (self.wall_time_fwd_prop_compute /
                 float(self.num_fwd_prop_step_computes))

            self.mean_num_onwd_prop_step_computes_per_update = \
                (self.num_onwd_prop_step_computes /
                 float(self.num_ctrl_amp_updates))

            self.mean_wall_time_onwd_prop_compute = \
                (self.wall_time_onwd_prop_compute /
                 float(self.num_onwd_prop_step_computes))

        def report(self):
            """
            Print a report of the stats to the console
            """

            print("\n------------------------------------"
                  "\n---- Control optimisation stats ----")
            self.report_timings()
            self.report_func_calls()
            self.report_amp_updates()
            self.report_dyn_gen_comps()
            self.report_fwd_prop()
            self.report_onwd_prop()
            print("------------------------------------")

        def report_dyn_gen_comps(self):
            print("**** {} Computations ****".format(self.dyn_gen_name))
            print("Total number of {} computations: "
                  "{}".format(self.dyn_gen_name, self.num_dyn_gen_computes))
            print("Mean number of {} computations per update: "
                  "{}".format(self.dyn_gen_name,
                              self.mean_num_dyn_gen_computes_per_update))
            print("Mean wall time to compute {}s: "
                  "{} s".format(self.dyn_gen_name,
                                self.mean_wall_time_dyn_gen_compute))
            print("**** Propagator Computations ****")
            print("Total number of propagator computations: "
                  "{}".format(self.num_prop_computes))
            print("Mean number of propagator computations per update: "
                  "{}".format(self.mean_num_prop_computes_per_update))
            print("Mean wall time to compute propagator "
                  "{} s".format(self.mean_wall_time_prop_compute))

        def report_fwd_prop(self):
            print("**** Forward Propagation ****")
            print("Total number of forward propagation step computations: "
                  "{}".format(self.num_fwd_prop_step_computes))
            print("Mean number of forward propagation step computations"
                  " per update: "
                  "{}".format(self.mean_num_fwd_prop_step_computes_per_update))
            print("Mean wall time to compute forward propagation "
                  "{} s".format(self.mean_wall_time_fwd_prop_compute))

        def report_onwd_prop(self):
            print("**** Onward Propagation ****")
            print("Total number of onward propagation step computations: "
                  "{}".format(self.num_onwd_prop_step_computes))
            print("Mean number of onward propagation step computations"
                  " per update: "
                  "{}".format(self.mean_num_onwd_prop_step_computes_per_update))
            print("Mean wall time to compute onward propagation "
                  "{} s".format(self.mean_wall_time_onwd_prop_compute))