from __future__ import annotations

from abc import ABCMeta, abstractmethod
from typing import Tuple, TYPE_CHECKING

from torch import Tensor
from torch.nn import Module
from torch.optim import Optimizer

if TYPE_CHECKING:
    from ..xztrainer import XZTrainer, TrainContext
    from ..model import LRSchedulerProtocol, SchedulerType


class TrainingEngine(metaclass=ABCMeta):
    @abstractmethod
    def wrap_model(self, model: Module, optimizer: Optimizer, scheduler: LRSchedulerProtocol,
                   scheduler_type: SchedulerType) -> Tuple[Module, Optimizer, LRSchedulerProtocol]:
        ...

    @abstractmethod
    def backward_pass(self, context: TrainContext, batch_i: int, loss: Tensor):
        ...


class TrainingEngineConfig(metaclass=ABCMeta):
    @abstractmethod
    def create_engine(self, trainer: XZTrainer) -> TrainingEngine:
        pass
