import os
import pandas as pd
import joblib
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import seaborn as sns
import matplotlib.pyplot as plt

def evaluate_best_model(data_dir, model_path):
    print("Loading test dataset...")
    X_test = pd.read_csv(os.path.join(data_dir, "X_test.csv"))
    y_test = pd.read_csv(os.path.join(data_dir, "y_test.csv")).values.ravel()
    
    print(f"Loading best model from {model_path}...")
    if not os.path.exists(model_path):
        print(f"Model not found at {model_path}. Please train the model first.")
        return
        
    model = joblib.load(model_path)
    
    # Check if label encoder exists
    le_path = os.path.join(data_dir, "label_encoder.pkl")
    target_names = None
    if os.path.exists(le_path):
        le = joblib.load(le_path)
        target_names = [str(c) for c in le.classes_]
    
    print("Generating predictions...")
    y_pred = model.predict(X_test)
    
    print("\n--- Model Evaluation ---")
    acc = accuracy_score(y_test, y_pred)
    print(f"Test Accuracy: {acc:.4f}\n")
    
    print("Classification Report:")
    print(classification_report(y_test, y_pred, target_names=target_names if target_names else None, zero_division=0))
    
    # Generates a confusion matrix (optional visual)
    cm = confusion_matrix(y_test, y_pred)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=target_names if target_names else "auto", 
                yticklabels=target_names if target_names else "auto")
    plt.title('Confusion Matrix')
    plt.ylabel('Actual Class')
    plt.xlabel('Predicted Class')
    cm_path = os.path.join(os.path.dirname(model_path), "confusion_matrix.png")
    plt.savefig(cm_path)
    print(f"Saved confusion matrix matrix chart to {cm_path}")
    print("Evaluation complete.")

if __name__ == "__main__":
    datasets_directory = r"C:\Users\KAUSTAV\OneDrive\Desktop\NetWokie\NetWokie-AI\datasets"
    model_filepath = r"C:\Users\KAUSTAV\OneDrive\Desktop\NetWokie\NetWokie-AI\ml_models\device_classifier.pkl"
    
    evaluate_best_model(datasets_directory, model_filepath)