################################################################################
# Copyright (c) 2021 ContinualAI.                                              #
# Copyrights licensed under the MIT License.                                   #
# See the accompanying LICENSE file for terms.                                 #
#                                                                              #
# Date: 24/07/2021                                                             #
# Author(s): Antonio Carta                                                     #
# E-mail: contact@continualai.org                                              #
# Website: avalanche.continualai.org                                           #
################################################################################
from avalanche.training.plugins import StrategyPlugin


class Clock(StrategyPlugin):
    def __init__(self):
        """ Counter for strategy events. """
        super().__init__()
        # train
        self.train_iterations = 0
        """ Total number of training iterations. """

        self.train_exp_counter = 0
        """ Number of past training experiences. """

        self.train_exp_epochs = 0
        """ Number of training epochs for the current experience. """

        self.train_exp_iterations = 0
        """ Number of training iterations for the current experience. """

        self.train_epoch_iterations = 0
        """ Number of iterations for the current epoch. """

        self.total_iterations = 0
        """ Total number of iterations in training and eval mode. """

    def before_training_exp(self, strategy, **kwargs):
        self.train_exp_iterations = 0
        self.train_exp_epochs = 0

    def before_training_epoch(self, strategy, **kwargs):
        self.train_epoch_iterations = 0

    def after_training_iteration(self, strategy, **kwargs):
        self.train_epoch_iterations += 1
        self.train_exp_iterations += 1
        self.train_iterations += 1
        self.total_iterations += 1

    def after_training_epoch(self, strategy, **kwargs):
        self.train_exp_epochs += 1

    def after_training_exp(self, strategy, **kwargs):
        self.train_exp_counter += 1

    def after_eval_iteration(self, strategy, **kwargs):
        self.total_iterations += 1
