import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
from datasets import load_dataset
import joblib
import os
import numpy as np

# Define paths for the Random Forest model, TF-IDF vectorizer, and label encoder
rf_model_path = 'random_forest_model.pkl'
vectorizer_path = "tfidf_vectorizer.pkl"
label_encoder_path = "label_encoder.pkl"
multi_rf_model_path= "random_forest_multi_model.pkl"

# Check if models and encoder exist
if os.path.exists(rf_model_path) and os.path.exists(vectorizer_path) and os.path.exists(label_encoder_path) and os.path.exists(multi_rf_model_path):
    # Load the models if they already exist
    rf_single = joblib.load(rf_model_path)
    vectorizer = joblib.load(vectorizer_path)
    le = joblib.load(label_encoder_path)
    rf_multi = joblib.load(multi_rf_model_path)
    print("Random Forest model, vectorizer, and label encoder loaded from disk.")
else:
    # Load the dataset
    ds = load_dataset('ahmedheakl/resume-atlas', cache_dir="C:/Users/dell/.cache/huggingface/datasets")

    # Create a DataFrame from the 'train' split
    df_train = pd.DataFrame(ds['train'])

    # Initialize the Label Encoder and encode the 'Category' labels
    le = LabelEncoder()
    df_train['Category_encoded'] = le.fit_transform(df_train['Category'])

    # Split the dataset into training and test sets
    X_train, X_test, y_train, y_test = train_test_split(
        df_train['Text'], df_train['Category_encoded'], test_size=0.2, random_state=42)

    # Initialize TF-IDF Vectorizer and transform the text data
    vectorizer = TfidfVectorizer(max_features=1000)
    X_train_tfidf = vectorizer.fit_transform(X_train)
    X_test_tfidf = vectorizer.transform(X_test)

    # Initialize and train the Random Forest models
    rf_single = RandomForestClassifier(n_estimators=100, random_state=42)
    rf_single.fit(X_train_tfidf, y_train)
    
    rf_multi = RandomForestClassifier(n_estimators=100, random_state=42)
    rf_multi.fit(X_train_tfidf, y_train)

    # Save the Random Forest models, TF-IDF vectorizer, and label encoder
    joblib.dump(rf_single, rf_model_path)
    joblib.dump(rf_multi, multi_rf_model_path)
    joblib.dump(vectorizer, vectorizer_path)
    joblib.dump(le, label_encoder_path)
    print("Random Forest model, vectorizer, and label encoder trained and saved to disk.")

# Single-label classification function for Random Forest model
def classify_text_rf(text):
    try:
        text_tfidf = vectorizer.transform([text])
        predicted_class_index = rf_single.predict(text_tfidf)[0]
        predicted_category = le.inverse_transform([predicted_class_index])[0]
        return predicted_category
    except Exception as e:
        print(f"Error in classify_text_rf: {e}")
        return None

# Multi-label classification function with top N predictions
def classify_text_rf_multi(text, top_n=3):
    try:
        text_tfidf = vectorizer.transform([text])
        probabilities = rf_multi.predict_proba(text_tfidf)[0]
        top_n_indices = np.argsort(probabilities)[::-1][:min(top_n, len(probabilities))]
        top_n_categories = le.inverse_transform(top_n_indices)
        return top_n_categories
    except Exception as e:
        print(f"Error in classify_text_rf_multi: {e}")
        return None