import pickle
from typing import override, Callable
from constants import TRAINING_DIR, MODEL_DIR
from labels import Label
from models.model import Model
from helper import dataset_iterator, default_labelling
import pandas as pd
from random import random

class Baseline_model(Model):
    def __init__(self, model_filename:str="", label_translator: Callable[[str], Label] = default_labelling) -> None:
        self.fake_probability = 0
        super().__init__(model_filename, label_translator)

    @override
    def train(self, training_dataset:str, hyperparameters:dict[str, float]={}) -> None:
        fake_amount = 0
        real_amount = 0
        total_amount = 0

        for chunk in dataset_iterator(f"{TRAINING_DIR}/{training_dataset}", columns=['type']):
            chunk_fake_amount = (chunk['type'].map(self.label_translator) == Label.FAKE).sum()

            fake_amount += chunk_fake_amount
            real_amount += len(chunk) - chunk_fake_amount
            total_amount += len(chunk)

        self.fake_probability = fake_amount/total_amount

    @override
    def classify(self, input:pd.Series) -> Label:
        if random() <= self.fake_probability:
            return Label.FAKE
        return Label.REAL

    @override
    def save(self, filename:str) -> None:
        data = {}
        data["label_translator"] = self.label_translator
        data["fake_probability"] = self.fake_probability

        with open(f"{MODEL_DIR}/{filename}", 'wb') as file:
            pickle.dump(data, file)

    @override
    def load(self, filename:str) -> None:
        with open(f"{MODEL_DIR}/{filename}", 'rb') as file:
            data = pickle.load(file)

        self.label_translator = data["label_translator"]
        self.fake_probability = data["fake_probability"]
