# -------------------------------------------------------------------
# Pimcore
#
# This source file is available under two different licenses:
#  - GNU General Public License version 3 (GPLv3)
#  - Pimcore Commercial License (PCL)
# Full copyright and license information is available in
# LICENSE.md which is distributed with this source code.
#
#  @copyright  Copyright (c) Pimcore GmbH (http://www.pimcore.org)
#  @license    http://www.pimcore.org/license     GPLv3 and PCL
# -------------------------------------------------------------------


import logging
from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl

from .training_status import TrainingStatus


logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

class ProgressCallback(TrainerCallback):

    __trainingStatus: TrainingStatus = None
    __startPercentage: int = None
    __endPercentage: int = None

    def __init__(self, trainingStatus: TrainingStatus, startPercentage: int, endPercentage: int):
        self.__trainingStatus = trainingStatus
        self.__startPercentage = startPercentage
        self.__endPercentage = endPercentage

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        logger.info(f"Completed step {state.global_step} of {state.max_steps}")

        if self.__trainingStatus.is_training_aborted(): 
            control.should_training_stop = True
            logger.info("Training aborted")
            return

        scope = self.__endPercentage - self.__startPercentage
        progress = round(self.__startPercentage + (state.global_step / state.max_steps) * scope, 2)

        self.__trainingStatus.update_status(progress, f"Training model, completed step {state.global_step} of {state.max_steps}")