"""Routines that can be passed as `info_hook` to :func:`.optimize_pulses`"""
import sys
import time
import numpy as np
from uniseg.graphemecluster import grapheme_clusters

__all__ = ['chain', 'print_debug_information', 'print_table']


def _qobj_nbytes(qobj):
    """Return an estimate for the number of bytes of memory required to store a
    Qobj"""
    return (
        qobj.data.data.nbytes
        + qobj.data.indptr.nbytes
        + qobj.data.indices.nbytes
        + sys.getsizeof(qobj)
        + sum(sys.getsizeof(attr) for attr in qobj.__dict__)
    )


def chain(*hooks):
    """Chain multiple `info_hook` or `modify_params_after_iter` callables
    together.

    Example:

        >>> def print_fidelity(**kwargs):
        ...     F_re = np.average(np.array(kwargs['tau_vals']).real)
        ...     print("    F = %f" % F_re)
        >>> info_hook = chain(print_debug_information, print_fidelity)

    Note:

        Functions that are connected via :func:`chain` may use the
        `shared_data` share the same `shared_data` argument, which they can use
        to communicate down the chain.
    """

    def info_hook(**kwargs):
        result = []
        for hook in hooks:
            res = hook(**kwargs)
            if res is not None:
                result.append(res)
        if len(result) > 0:
            if len(result) == 1:
                return result[0]
            else:
                return tuple(result)
        else:
            return None

    return info_hook


def print_debug_information(
    *,
    objectives,
    adjoint_objectives,
    backward_states,
    forward_states,
    forward_states0,
    guess_pulses,
    optimized_pulses,
    g_a_integrals,
    lambda_vals,
    shape_arrays,
    fw_states_T,
    tlist,
    tau_vals,
    start_time,
    stop_time,
    iteration,
    info_vals,
    shared_data,
    out=sys.stdout
):
    r"""Print full debug information about the current Krotov iteration

    This routine is intended to be passed to :func:`.optimize_pulses` as
    `info_hook`, and it exemplifies the full signature of a routine suitable
    for this purpose.

    Keyword Args:
        objectives (list[Objective]): list of the objectives
        adjoint_objectives (list[Objective]): list of the adjoint objectives
        backward_states (list): If available, for each objective, an array-like
            object containing the states of the Krotov backward propagation.
        forward_states: If available (second order only), for each objective,
            an array-like object containing the forward-propagated states under
            the optimized pulses. None otherwise.
        forward_states0: If available (second order only), for each objective,
            an array-like object containing the forward-propagated states under
            the guess pulses. None otherwise.
        guess_pulses (list[numpy.ndarray]): list of guess pulses
        optimized_pulses (list[numpy.ndarray]): list of optimized pulses
        g_a_integrals (numpy.ndarray): array of values
            :math:`\int_0^T g_a(t) \dd t = \int_0^T \frac{\lambda_a}{S(t)}
            \Abs{\Delta \epsilon(t)}^2 \dd t`, for each pulse $\epsilon(t)$.
            The pulse updates $\Delta \epsilon(t)$ are the differences of the
            `optimized_pulses` and the `guess_pulses` (zero in the zeroth
            iteration that only performs a forward-propagation of the guess
            pulses). The quantity $\int g_a(t) \dd t$ is a very useful measure
            of how much the pulse amplitudes changes in each iteration. This
            tells us whether we've chosen good values for $\lambda_a$. Values
            that are too small cause "pulse explosions" which immediately show
            up in $\int_0^T g_a(t) \dd t$.  Also, whether $\int g_a(t) \dd t$
            is increasing or decreasing between iterations gives an indication
            whether the optimization is "speeding up" or "slowing down", and
            thus whether convergence is reached (negligible pulse updates).
        lambda_vals (numpy.ndarray): for each pulse, the value of the
            $\lambda_a$ parameter
        shape_arrays (list[numpy.ndarray]): for each pulse, the array of
            update-shape values $S(t)$
        fw_states_T (list): for each objective, the forward-propagated state
        tlist (numpy.ndarray): array of time grid values on which the states
            are defined
        tau_vals (numpy.ndarray): for each objective, the complex overlap for
            the forward-propagated state with the target state, or None if no
            target state is defined.
        start_time (float): The time at which the iteration started, in epoch
            seconds
        stop_time (float): The time at which the iteration started, in epoch
            seconds
        iteration (int): The current iteration number. For the initial
            propagation of the guess controls, zero.
        shared_data (dict): Dict of data shared between any
            `modify_params_after_iter` and any `info_hook` functions chained
            together via :func:`chain`.
        info_vals (list): List of the return values of the info_hook from
            previous iterations
        out: An open file handle where to write the information. This parameter
            is not part of the `info_hook` interface, and defaults to stdout.
            Use :func:`functools.partial` to pass a different value.

    Note:
        This routine implements the full signature of an `info_hook` in
        :func:`.optimize_pulses`, excluding `out`. However, since
        the `info_hook` only allows for keyword arguments, it is usually much
        simpler to use Python's variable keyword arguments syntax
        (``**kwargs``). For example, consider the following `info_hook` that
        prints (and stores) the value of the real-part gate fidelity::

            def print_fidelity(**kwargs):
                F_re = np.average(np.array(kwargs['tau_vals']).real)
                print("    F = %f" % F_re)
                return F_re
    """
    out.write("Iteration %d\n" % iteration)
    if iteration == 0:
        out.write("    objectives:\n")
        for (i, obj) in enumerate(objectives):
            out.write("        %d:%s\n" % (i + 1, obj))
        out.write("    adjoint objectives:\n")
        for (i, obj) in enumerate(adjoint_objectives):
            out.write("        %d:%s\n" % (i + 1, obj))
        out.write(
            "    S(t) (ranges): %s\n"
            % (
                ", ".join(
                    ["[%f, %f]" % (np.min(S), np.max(S)) for S in shape_arrays]
                )
            )
        )
    out.write(
        "    duration: %.1f secs (started at %s)\n"
        % (
            (stop_time - start_time),
            time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(start_time)),
        )
    )
    out.write(
        "    optimized pulses (ranges): %s\n"
        % (", ".join([_pulse_range(pulse) for pulse in optimized_pulses]))
    )
    out.write(
        "    ∫gₐ(t)dt: %s\n" % (", ".join(["%.2e" % v for v in g_a_integrals]))
    )
    out.write(
        "    λₐ: %s\n" % (", ".join(["%.2e" % λ for λ in lambda_vals]))
    )
    MB_per_timeslot = (
        sum(_qobj_nbytes(state) for state in fw_states_T) / (1024 ** 2)
    )
    out.write("    storage (bw, fw, fw0): ")
    if backward_states is None:
        out.write("None, ")
    else:
        n = len(backward_states)
        out.write(
            "[%d * %s(%d)] (%.1f MB), "
            % (
                n,
                backward_states[0].__class__.__name__,
                len(backward_states[0]),
                len(backward_states[0]) * MB_per_timeslot,
            )
        )
    if forward_states is None:
        out.write("None, ")
    else:
        n = len(forward_states)
        out.write(
            "[%d * %s(%d)] (%.1f MB), "
            % (
                n,
                forward_states[0].__class__.__name__,
                len(forward_states[0]),
                len(forward_states[0]) * MB_per_timeslot,
            )
        )
    if forward_states0 is None:
        out.write("None\n")
    else:
        n = len(forward_states0)
        out.write(
            "[%d * %s(%d)] (%.1f MB)\n"
            % (
                n,
                forward_states0[0].__class__.__name__,
                len(forward_states0[0]),
                len(forward_states0[0]) * MB_per_timeslot,
            )
        )
    out.write(
        "    fw_states_T norm: %s\n"
        % (", ".join(["%f" % state.norm() for state in fw_states_T]))
    )
    if not np.any(tau_vals == None):  # noqa
        out.write(
            "    τ: %s\n"
            % (
                ", ".join(
                    [
                        "(%.2e:%.2fπ)" % (abs(z), np.angle(z) / np.pi)
                        for z in tau_vals
                    ]
                )
            )
        )


def _grapheme_len(text):
    """Number of graphemes in `text`

    This is the length of the `text` when printed::
        >>> s = 'Â'
        >>> len(s)
        2
        >>> _grapheme_len(s)
        1
    """
    return len(list(grapheme_clusters(text)))


def _rjust(text, width, fillchar=' '):
    """Right-justify text for a total of `width` graphemes

    The `width` is based on graphemes::

        >>> s = 'Â'
        >>> s.rjust(2)
        'Â'
        >>> _rjust(s, 2)
        ' Â'
    """
    len_text = _grapheme_len(text)
    return fillchar * (width - len_text) + text


def print_table(
    J_T, show_g_a_int_per_pulse=False, J_T_prev=None, unicode=True
):
    r"""Print a tabular overview of the functional values in the iteration

    An example output is:

    .. code-block:: console

        iter.        J_T    ∫gₐ(t)dt          J       ΔJ_T         ΔJ  secs
            0   1.00e+00    0.00e+00   1.00e+00        n/a        n/a     0
            1   7.65e-01    2.33e-02   7.88e-01  -2.35e-01  -2.12e-01     1
            2   5.56e-01    2.07e-02   5.77e-01  -2.09e-01  -1.88e-01     1

    The table has the following columns:

    * iteration number
    * value of the final-time functional $J_T$
    * If `show_g_a_int_per_pulse` is True and there is more than one control
      pulse: one column for each pulse containing the value of :math:`\int_0^T
      \frac{\lambda_{a, i}}{S_i(t)} \Abs{\Delta \epsilon_i(t)}^2 \dd t`
    * The value of :math:`\sum_i \int_0^T g_a(\epsilon_i(t)) \dd t =
      \sum_i \int_0^T \frac{\lambda_{a, i}}{S_i(t)}
      \Abs{\Delta \epsilon_i(t)}^2 \dd t`, or just :math:`\int_0^T
      \frac{\lambda_{a}}{S(t)} \Abs{\Delta \epsilon(t)}^2 \dd t` if there
      is only a single control pulse (as in the above example output).
      This value (respectively the individual values with
      `show_g_a_int_per_pulse`) should always be at least three orders of
      magnitude smaller than the pulse fluence :math:`\sum_i\int_0^T
      \Abs{\epsilon_i(t)}^2 \dd t`.  Larger changes in the pulse amplitude may
      be a sign of a "pulse explosion" due to values for $\lambda_{a,i}$ that
      are too small. Changes in :math:`\sum_i \int_0^T \frac{\lambda_{a,
      i}}{S_i(t)}` are often a better indicator of whether the optimization is
      "speeding up"/"slowing down"/reaching convergence than the values of
      $J_T$.
    * The value of the total functional :math:`J = J_T +
      \sum_i \int_0^T g_a(\epsilon_i(t)) \dd t`
    * The change $\Delta J_T$ in the final time functional compared to the
      previous iteration. This should be a negative value, indicating monotonic
      convergence in a minimization of $J_T$.
    * The change $\Delta J$ in the total functional compared to the previous
      iteration. This is evaluated as :math:`\Delta J = \Delta J_T +
      \sum_i \int_0^T g_a(\epsilon_i(t)) \dd t`. Somewhat counter-intuitively,
      $\Delta J$ does not contain a contribution from the $g_a(t)$ of the
      previous iteration. This is because the $\Delta \epsilon_i(T)$ on which
      $g_a(t)$ depends must be evaluated with respect to the same reference
      field (the guess pulse of the *current* iteration), to that
      $\Delta \epsilon_i(T) = 0$ when evaluated with the optimized pulse of the
      previous iteration (i.e., the same guess pulse of the current iteration).
    * The number of seconds in wallclock time spent on the iteration
    * An indicator ``*`` or ``**`` if there is a loss of monotonic convergence
      in $\Delta J_T$ and/or $\Delta J$. Krotov's method mathematically
      guarantees a negative $\Delta J$ in the continuous limit. Assuming there
      are no errors in the time propagation, or in the `chi_constructor` passed
      to :func:`.optimize_pulses`, a loss of monotonic convergence is due to
      the $\lambda_a$ associated with the pulses (via `pulse_options` in
      :func:`.optimize_pulses`) being too small. In practice, we usually don't
      care too much about a loss of monotonic convergence in $\Delta J$, but a
      loss of convergence in $\Delta J_T$ is a serious sign of trouble. It is
      often associated with sharp discontinuous spikes in the optimized pulses,
      or a dramatic increase in the pulse amplitude.

    Args:
        J_T (callable): A function that extract the value of the final time
            functional from the keyword-arguments passed to the `info_hook`.
        show_g_a_int_per_pulse (callable): If True, print a column with the
            value of :math:`\int_0^T g_a(\epsilon_i(t)) \dd t = \int_0^T
            \frac{\lambda_{a, i}}{S_i(t)} \Abs{\Delta \epsilon_i(t)}^2 \dd t`
            for every pulse $\epsilon_i(t)$. Otherwise, only print the sum over
            those integrals for all pulses.
        J_T_prev (None or callable): A function that extract the value of the
            final time functional *from the previous iteration*. If None, use
            the last values from the `info_vals` passed to the `info_hook`.
        unicode (bool): Whether to use unicode symbols for the column headers.
            Some systems have broken monospace fonts in the Jupyter notebook
            that cause the headers not to line up as intended.
    """

    def get_J_T_prev(**kwargs):
        try:
            return kwargs['info_vals'][-1]
        except IndexError:
            return 0

    if J_T_prev is None:
        J_T_prev = get_J_T_prev

    def info_hook(**kwargs):
        iteration = kwargs['iteration']
        n_pulses = len(kwargs['guess_pulses'])
        if iteration == 0:
            print(_rjust("iter.", 7), end="")
            print(_rjust("J_T", 11), end="")
            if n_pulses > 1:
                if show_g_a_int_per_pulse:
                    for i in range(n_pulses):
                        if unicode:
                            print(_rjust("∫gₐ(ϵ%d)dt" % (i + 1), 12), end="")
                        else:
                            print(_rjust("g_a_int_%d" % (i + 1), 12), end="")
                if unicode:
                    print(_rjust("∑∫gₐ(t)dt", 12), end="")
                else:
                    print(_rjust("g_a_int", 12), end="")
            else:
                if unicode:
                    print(_rjust("∫gₐ(t)dt", 12), end="")
                else:
                    print(_rjust("g_a_int", 12), end="")
            print(_rjust("J", 11), end="")
            if unicode:
                print(_rjust("ΔJ_T", 11), end="")
            else:
                print(_rjust("Delta_J_T", 11), end="")
            if unicode:
                print(_rjust("ΔJ", 11), end="")
            else:
                print(_rjust("Delta J", 11), end="")
            print(_rjust("secs", 6))
        J_T_val = J_T(**kwargs)
        Σgₐdt = np.sum(kwargs['g_a_integrals'])
        J = J_T_val + Σgₐdt
        if iteration > 0:
            J_T_prev_val = J_T_prev(**kwargs)
            ΔJ_T = J_T_val - J_T_prev_val
            ΔJ = ΔJ_T + Σgₐdt
        secs = int(kwargs['stop_time'] - kwargs['start_time'])
        print("%7d" % iteration, end="")
        print(" %10.2e" % J_T_val, end="")
        if n_pulses > 1:
            if show_g_a_int_per_pulse:
                for i in range(n_pulses):
                    print(" %11.2e" % kwargs['g_a_integrals'][i], end="")
        print(" %11.2e" % Σgₐdt, end="")
        print(" %10.2e" % J, end="")
        if iteration == 0:
            print(" %10s" % "n/a", end="")
            print(" %10s" % "n/a", end="")
        else:
            print(" %10.2e" % ΔJ_T, end="")
            print(" %10.2e" % ΔJ, end="")
        print(" %5d" % secs, end="")
        if iteration > 0:
            if (ΔJ_T > 0) or (ΔJ > 0):
                print(" ", end="")
                if ΔJ_T > 0:
                    print("*", end="")
                if ΔJ > 0:
                    print("*", end="")
        print("")
        return J_T_val

    return info_hook


def _pulse_range(pulse):
    """Return the range information about the given pulse value array

    Example:

        >>> _pulse_range(np.array([-1, 1, 5]))
        '[-1.00, 5.00]'
        >>> _pulse_range(np.array([-1, 1+5j, 5]))
        '[(r:-1.00, i:0.00), (r:5.00, i:5.00)]'
    """
    if np.iscomplexobj(pulse):
        pulse_real = pulse.real
        pulse_imag = pulse.imag
        r0, r1 = np.min(pulse_real), np.max(pulse_real)
        i0, i1 = np.min(pulse_imag), np.max(pulse_imag)
        return "[(r:%.2f, i:%.2f), (r:%.2f, i:%.2f)]" % (r0, i0, r1, i1)
    else:
        return "[%.2f, %.2f]" % (np.min(pulse), np.max(pulse))
