import os.path
import pickle
import pandas as pd
from transformers import AutoTokenizer
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from xgboost import XGBClassifier
from sklearn.preprocessing import LabelEncoder

class MentalHealthClassifier:
    def __init__(self, data_path, model_path):
        self.data = pd.read_csv(data_path, skip_blank_lines=True)
        self.data['category'] = ['anxiety' if x == 'axienty' else x for x in self.data['category']]
        self.data.dropna(subset=['text'], inplace=True)
        self.data.dropna(subset=['clean_text'], inplace=True)
        self.data_selected = self.data[['clean_text', 'category']]
        self.df = pd.DataFrame(self.data_selected)
        self.label_encoder = LabelEncoder()
        self.df['category_encoded'] = self.label_encoder.fit_transform(self.df['category'])
        self.tokenizer = None
        self.vectorizer = CountVectorizer()
        self.model_path = model_path
        self.model = self.load_model() if os.path.exists(model_path) else XGBClassifier()

    def preprocess_data(self):
        tokenized_texts = [self.tokenizer.tokenize(text, padding=True, truncation=True) for text in self.df['clean_text']]
        X = self.vectorizer.fit_transform([' '.join(tokens) for tokens in tokenized_texts]).toarray()
        return X, self.df['category_encoded']

    def train_model(self, X, y):
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        self.model.fit(X_train, y_train)
        y_pred = self.model.predict(X_test)
        return y_test, y_pred

    def predict_category(self, input_text):
        if self.tokenizer is None:
            raise ValueError("Tokenizer not initialized. Call 'initialize_tokenizer' first.")
        tokenized_input = self.tokenizer.tokenize(input_text, padding=True, truncation=True)
        input_feature_vector = self.vectorizer.transform([' '.join(tokenized_input)]).toarray()
        predicted_category_encoded = self.model.predict(input_feature_vector)
        predicted_category = self.label_encoder.inverse_transform(predicted_category_encoded)
        return predicted_category[0]

    def initialize_tokenizer(self, model_name):
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)

    def save_model(self):
        print("saving model...to pickle...")
        with open(self.model_path, 'wb') as f:
            pickle.dump(self.model, f)

    def load_model(self):
        print("loading model...from pickle...")
        with open(self.model_path, 'rb') as f:
            return pickle.load(f)

if __name__ == "__main__":
    tokenizer_model_name = "nlptown/bert-base-multilingual-uncased-sentiment"
    data_path = 'data/data.csv'
    model_path = 'mental_health_model.pkl'
    mental_classifier = MentalHealthClassifier(data_path, model_path)

    if not os.path.exists(model_path):
        mental_classifier.initialize_tokenizer(tokenizer_model_name)
        X, y = mental_classifier.preprocess_data()
        y_test, y_pred = mental_classifier.train_model(X, y)
        mental_classifier.save_model()
    else:
        mental_classifier.load_model(model_path)
        mental_classifier.initialize_tokenizer(tokenizer_model_name)  # Ensure tokenizer is initialized if loading model from pickle
        mental_classifier.preprocess_data()

    # input_text = "I feel anxiety whenever i am doing nothing."
    # predicted_category = mental_classifier.predict_category(input_text)
    # print("Predicted mental health condition:", predicted_category)