import pickle
from typing import override, Callable
from scipy.sparse import lil_array
from constants import TRAINING_DIR, MODEL_DIR
from labels import Label
from models.model import Model
from helper import dataset_iterator
from sklearn.linear_model import LogisticRegression
import pandas as pd
import numpy as np
from helper import default_labelling

class Logistic_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:
        token_counts:dict[str, int] = {}
        sorted_token_counts:dict[str, int] = {}
        token_id:dict[str, int] = {} # converts top 10K words to id's.

        domain_counts:dict[str, int] = {}
        sorted_domain_counts:dict[str, int] = {}
        domain_id:dict[str, int] = {} # converts top 500 domains to id's.

        self.consider_metadata = False
        if "metadata" in hyperparameters and hyperparameters["metadata"] == 1:
            self.consider_metadata = True

        columns = ["tokens", "domain"]
        rows_processed = 0
        for chunk in dataset_iterator(f"{TRAINING_DIR}/{training_dataset}", columns=columns):
            rows_processed += len(chunk)
            for _, row in chunk.iterrows():
                for token in row['tokens']:
                    if token not in token_counts:
                        token_counts[token] = 0
                    token_counts[token] += 1
                if row['domain'] not in domain_counts:
                    domain_counts[row['domain']] = 0
                domain_counts[row['domain']] += 1
        for token in sorted(token_counts, key=lambda token: token_counts[token], reverse=True):
            sorted_token_counts[token] = token_counts[token]
        for domain in sorted(domain_counts, key=lambda domain: domain_counts[domain], reverse=True):
            sorted_domain_counts[domain] = domain_counts[domain]
        
        idx = 0
        for token in sorted_token_counts:
            token_id[token] = idx
            idx += 1
            if idx >= 10000:
                break

        idx = 0
        for domain in sorted_domain_counts:
            domain_id[domain] = idx
            idx += 1
            if idx >= 1000:
                break

        if self.consider_metadata: # consider things other than tokens
            X = lil_array((rows_processed, 11000), dtype="float64")
        else:
            X = lil_array((rows_processed, 10000), dtype="float64") # non-sparse array uses 74GiB ram on 995,000_rows. Sklearn LogisticRegression supports sparse arrays though. It still uses 9+ now.

        Y = np.zeros(rows_processed, dtype=int)

        columns.append("type")
        article_num = 0
        for chunk in dataset_iterator(f"{TRAINING_DIR}/{training_dataset}", columns=columns):
            for _, row in chunk.iterrows():
                tokens = row['tokens']
                article_type = row['type']

                article_word_counts = np.zeros(10000)
                for token in tokens:
                    if token not in token_id:
                        continue # if they are not in top 10K vocab we can ignore them
                    article_word_counts[token_id[token]] += 1
                X[article_num, :10000] = article_word_counts
                if self.consider_metadata:
                    if row['domain'] in domain_id:
                        X[article_num, 10000+domain_id[row['domain']]] = 1

                Y[article_num] = self.label_translator(article_type).value
                article_num += 1

        self.regression_model = LogisticRegression(max_iter=10000, n_jobs = -1, class_weight="balanced").fit(X, Y)
        self.token_id = token_id
        self.domain_id = domain_id

    @override
    def classify(self, input: pd.Series) -> Label:
        if self.consider_metadata:
            x = np.zeros(11000)
        else:
            x = np.zeros(10000)

        for token in input['tokens']:
            if token not in self.token_id:
                continue
            x[self.token_id[token]] += 1

        if self.consider_metadata:
            if input['domain'] in self.domain_id:
                x[10000+self.domain_id[input['domain']]] = 1

        prediction = self.regression_model.predict([x])[0]
        return Label(prediction)

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

        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.regression_model = data["regression_model"]
        self.token_id = data["token_id"]
        self.domain_id = data["domain_id"]
        self.consider_metadata = data["consider_metadata"] 
