from abc import ABC, abstractmethod
import pandas as pd
from time import perf_counter

from constants import TESTING_DIR, VALIDATION_DIR
from helper import LIAR_labelling, dataset_iterator, default_labelling
from labels import Label
from typing import Callable

class Model(ABC):
    def __init__(self, model_filename:str="", label_translator: Callable[[str], Label] = default_labelling) -> None:
        self.label_translator = label_translator
        if model_filename:
            self.load(model_filename)

    @abstractmethod
    def train(self, training_dataset:str, hyperparameters:dict[str, float]) -> None:
        pass

    @abstractmethod
    def classify(self, input:pd.Series) -> Label:
        pass

    @abstractmethod
    def save(self, filename:str) -> None:
        pass

    @abstractmethod
    def load(self, filename:str) -> None:
        pass

    def test(self, test_dataset: str, validate:bool=True) -> tuple[float, float, float, float]:
        TP = 0
        TN = 0
        FP = 0
        FN = 0

        if test_dataset == "LIAR.parquet":
            self.label_translator = LIAR_labelling

        dataset_dir = VALIDATION_DIR if validate else TESTING_DIR
        df = pd.read_parquet(f"{dataset_dir}/{test_dataset}")

        expected = df['type'].apply(self.label_translator)
        predicted = df.apply(self.classify, axis=1)

        TP = ((expected == Label.FAKE) & (predicted == Label.FAKE)).sum()
        FP = ((expected == Label.REAL) & (predicted == Label.FAKE)).sum()
        TN = ((expected == Label.REAL) & (predicted == Label.REAL)).sum()
        FN = ((expected == Label.FAKE) & (predicted == Label.REAL)).sum()

        accuracy = (TP + TN) / (TP + TN + FP + FN) if (TP + TN + FP + FN) > 0 else 0
        recall = (TP) / (TP + FN) if (TP + FN) > 0 else 0
        precision = (TP) / (TP + FP) if (TP + FP) > 0 else 0
        F1 = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0

        print(f"Accuracy {accuracy}")
        print(f"Recall {recall}")
        print(f"precision {precision}")
        print(f"F1-score {F1}")
        return (accuracy, recall, precision, F1)
