"""
This module provides tools for the top-level loop over epochs during training.
kubeGA expects the training program to be implemented as loop over several
epochs, each containing a series of loops over datasets (e.g. one loop over the
training set followed by one loop over the validation set). The program can be
interrupted between every iteration of any dataset loop, trigger a checkpoint
to be taken, and restarted using a different set of replicas.

**Due to checkpoint-restarts, parts of the training program may be executed
multiple times (e.g. once after each restart)!** To avoid incorrect execution,
ensure that your code is idempotent in the following locations:

1.  Immediately before any epoch loop (using :func:`remaining_epochs_until`).
2.  Immediately before any dataset loop (using
    :class:`kubega.torch.data.AdaptiveDataLoader`).

Your code may be non-idempotent in other locations.

.. code-block:: python

    ### IDEMPOTENT CODE ONLY ###
    for epoch in remaining_epochs_until(30):
        ### IDEMPOTENT CODE ONLY ###
        for batch in train_loader:
            # ... any code ...
        ### IDEMPOTENT CODE ONLY ###
        for batch in valid_loader:
            # ... any code ...
        # ... any code ...

    # ... any code ...

    ### END PROGRAM ###

For example, a common non-idempotent operation is learning-rate annealing:

.. code-block:: python

    for epoch in remaining_epochs_until(30):
        lr_scheduler.step()  # (A) WRONG!

        for batch in train_loader:
            # ...

        lr_scheduler.step()  # (B) WRONG!

        for batch in valid_loader:
            # ...

        lr_scheduler.step()  # (C) OK!

Location (A) will be executed again after any checkpoint-restart during either
the training or validation loop, resulting in the learning rate being annealed
several times in one epoch! Similarly, with location (B), if checkpoint-restart
happens during the validation loop.

Location (C) results in the correct behavior, because (1) an epoch will not be
repeated once it has finished, and (2) no checkpoint-restarts can occur between
the learning rate annealing and the end of the epoch.

.. _idempotent: https://stackoverflow.com/a/1077421
"""
import logging
import pickle

from kubega import checkpoint

logging.basicConfig(level=logging.INFO)
LOG = logging.getLogger(__name__)
LOG.setLevel(logging.INFO)

_EPOCH_STATE = None


def _epoch_state():
    global _EPOCH_STATE
    if _EPOCH_STATE is None:
        _EPOCH_STATE = _EpochState()
        checkpoint.load_state(_EPOCH_STATE)
    return _EPOCH_STATE


def current_epoch():
    """
    Get the current epoch while iterating with :func:`remaining_epochs_until`.

    Returns:
        int or None: The current epoch number if called from within a
        :func:`remaining_epochs_until` iteration, ``None`` otherwise.
    """
    return _epoch_state().current_epoch


def finished_epochs():
    """
    Get the number of epochs finished using :func:`remaining_epochs_until`.

    Returns:
        int: The number of finished epochs. Equal to :func:`current_epoch`
        if called from within a :func:`remaining_epochs_until` iteration.
    """
    return _epoch_state().finished_epochs


def remaining_epochs_until(epoch):
    """
    Iterate over epochs in a way that is consistent with checkpoint-restarts.
    For example:

    .. code-block:: python

        for epoch in remaining_epochs_until(30):
            print(current_epoch())  # Should print 0 through 29

        for epoch in remaining_epochs_until(60):
            print(current_epoch())  # Should print 30 through 59

    If a checkpoint-restart happens during an epoch, all previous epochs will
    be skipped after the program restarts.

    Arguments:
        epoch (int): The epoch number to end at (exclusively).

    Raises:
        RuntimeError: If invoked before a previous epoch loop has ended.
    """
    if current_epoch() is not None:
        raise RuntimeError("overlapping epoch loops detected")
    if finished_epochs() < epoch:
        LOG.info("starting at epoch %s", finished_epochs())
    else:
        LOG.info("skipping all epochs up to %s", epoch)
    while finished_epochs() < epoch:
        _epoch_state().current_epoch = finished_epochs()
        try:
            yield current_epoch()
        finally:
            # Try to catch any exits from epoch loop, including breaks and
            # Exceptions. See https://www.peterbe.com/plog/generatorexit.
            _epoch_state().finished_epochs += 1
            _epoch_state().current_epoch = None


class _EpochState(checkpoint.State):
    def __init__(self):
        super().__init__(".kubega-epoch")
        self.finished_epochs = 0
        self.current_epoch = None

    def save(self, fileobj):
        pickle.dump(self.finished_epochs, fileobj)

    def load(self, fileobj):
        self.finished_epochs = pickle.load(fileobj)
