from typing import *

from approaches.abst_appr import AbstractAppr
from approaches.hat.model_hat import ModelHAT


class Appr(AbstractAppr):
    def __init__(self, device: str, list__ncls: List[int], inputsize: Tuple[int, ...],
                 lr: float, lr_factor: float, lr_min: float, epochs_max: int, patience_max: int,
                 smax: float, lamb: float,
                 drop1: float, drop2: float, ):
        super().__init__(device=device, list__ncls=list__ncls, inputsize=inputsize,
                         lr=lr, lr_factor=lr_factor, lr_min=lr_min,
                         epochs_max=epochs_max, patience_max=patience_max,
                         smax=smax, lamb=lamb)
        self.model = ModelHAT(list__ncls=list__ncls, inputsize=inputsize,
                              smax=smax, hat_enabled=True,
                              drop1=drop1, drop2=drop2).to(self.device)
    # enddef

    def complete_learning(self, idx_task: int) -> None:
        self.model.freeze_masks(idx_task)
    # enddef
# enclasss
