from constants import TRAINING_DIR, MODEL_DIR
from models.model import Model
from labels import Label
from helper import default_labelling
from typing import override, Callable
import pandas as pd
import pickle
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.ensemble import GradientBoostingClassifier

def no_tokenization(str):
    return str.split(" ")

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

    @override
    def train(self, training_dataset: str, hyperparameters: dict[str, float] = {}) -> None:
        print("this model takes around 10 hours to train")

        X = pd.read_parquet(f"{TRAINING_DIR}/{training_dataset}", columns=['tokens'])['tokens']
        X = X.apply(lambda token_list: " ".join(token_list))
        Y = pd.read_parquet(f"{TRAINING_DIR}/{training_dataset}", columns=['type'])['type']
        Y = Y.apply(lambda label: self.label_translator(label).value)

        X = X[:250000]
        Y = Y[:250000]

        model = Pipeline([
            ("L string", TfidfVectorizer(tokenizer=no_tokenization)),
            ("forest", GradientBoostingClassifier(random_state=0, n_estimators=4000))
        ])

        model.fit(X, Y)
        self.model = model

    @override
    def classify(self, input: pd.Series) -> Label:
        X = " ".join(input['tokens'])
        return Label(self.model.predict([X])[0])

    @override
    def save(self, filename: str) -> None:
        data = {}
        data["label_translator"] = self.label_translator
        data["model"] = self.model
        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.model = data["model"]
