import copy

from torch.nn import Module


class ModelStateManager:
    def __init__(
        self,
        mode: str = "none",
        alpha: float = 0.0,
        patience: int = 1,
        min_delta: float = 0.0,
    ) -> None:
        if mode not in ("none", "trend", "plateau"):
            raise ValueError("mode must be one of: none, trend, plateau")
        if patience < 1:
            raise ValueError("patience must be >= 1")
        if min_delta < 0.0:
            raise ValueError("min_delta must be >= 0")

        self.mode = mode
        self.early_stopping = mode != "none"

        self.best_loss: float = float("inf")
        self.best_model_state: dict = {}

        self.alpha: float = alpha
        self.patience: int = patience
        self.min_delta: float = min_delta
        self.stop: bool = False

        self.prev_loss: float = float("inf")
        self.derivative: float = 0.0
        self.positive_derivative_streak: int = 0
        self.non_improving_streak: int = 0

    def __call__(self, model: Module, loss: float) -> None:
        if not self.best_model_state:
            self.best_model_state = copy.deepcopy(model.state_dict())

        previous_best_loss = self.best_loss
        if loss < previous_best_loss:
            self.best_loss = loss
            self.best_model_state = copy.deepcopy(model.state_dict())

        if not self.early_stopping:
            return

        if self.mode == "plateau":
            if loss < (previous_best_loss - self.min_delta):
                self.non_improving_streak = 0
            else:
                self.non_improving_streak += 1

            if self.non_improving_streak >= self.patience:
                self.stop = True
            return

        unfiltered_derivate: float = loss - self.prev_loss if self.prev_loss != float("inf") else 0.0
        self.derivative = self.alpha * self.derivative + (1 - self.alpha) * unfiltered_derivate

        self.prev_loss = loss

        if self.derivative > 0:
            self.positive_derivative_streak += 1
        else:
            self.positive_derivative_streak = 0

        if self.positive_derivative_streak >= self.patience:
            self.stop = True

    def load_best_model_state(self, model: Module) -> None:
        model.load_state_dict(self.best_model_state)
