import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import joblib
import pandas as pd
import time
import re
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Paths
MINILM_PATH = "/home/office-7/Downloads/minilm_v2/models/minilm"
XLM_ROBERTA_PATH = "/home/office-7/Downloads/xlm_roberta_v2/models/xlm_roberta"
ML_MODEL_PATH = "/home/office-7/support-intelligence-backend/core/services/models/spam_detection_model.pkl"

def preprocess_ml(text):
    text = str(text).lower()
    text = re.sub(r'[^\w\s]', '', text)
    text = re.sub(r'\d+', '', text)
    return text.strip()

def evaluate():
    logger.info("Loading models...")
    minilm_tokenizer = AutoTokenizer.from_pretrained(MINILM_PATH)
    minilm_model = AutoModelForSequenceClassification.from_pretrained(MINILM_PATH)
    
    xlm_tokenizer = AutoTokenizer.from_pretrained(XLM_ROBERTA_PATH)
    xlm_model = AutoModelForSequenceClassification.from_pretrained(XLM_ROBERTA_PATH)
    
    ml_model = joblib.load(ML_MODEL_PATH)
    
    df = pd.read_csv("eval_dataset.csv")
    
    results = []

    models = [
        ("MiniLM v2", minilm_model, minilm_tokenizer, "transformer"),
        ("XLM-Roberta v2", xlm_model, xlm_tokenizer, "transformer"),
        ("ML Model (English)", ml_model, None, "ml")
    ]

    for model_name, model, tokenizer, model_type in models:
        logger.info(f"Evaluating {model_name}...")
        
        for category in df['category'].unique():
            cat_df = df[df['category'] == category]
            texts = cat_df['text'].tolist()
            labels = cat_df['label'].tolist()
            
            predictions = []
            start_time = time.time()
            
            if model_type == "transformer":
                for text in texts:
                    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
                    with torch.no_grad():
                        outputs = model(**inputs)
                    pred = torch.argmax(outputs.logits, dim=-1).item()
                    # Mapping based on typical spam detection labels (1=spam, 0=ham)
                    # Check id2label if available
                    if hasattr(model.config, 'id2label'):
                        label_str = model.config.id2label[pred].lower()
                        pred = 1 if 'spam' in label_str else 0
                    predictions.append(pred)
            else:
                clean_texts = [preprocess_ml(t) for t in texts]
                predictions = model.predict(clean_texts).tolist()
            
            end_time = time.time()
            total_time = end_time - start_time
            avg_speed = total_time / len(texts)
            
            acc = accuracy_score(labels, predictions)
            prec = precision_score(labels, predictions, zero_division=0)
            rec = recall_score(labels, predictions, zero_division=0)
            f1 = f1_score(labels, predictions, zero_division=0)
            
            results.append({
                "Model": model_name,
                "Category": category,
                "Accuracy": f"{acc:.4f}",
                "Precision": f"{prec:.4f}",
                "Recall": f"{rec:.4f}",
                "F1 Score": f"{f1:.4f}",
                "Avg Speed (s/req)": f"{avg_speed:.4f}"
            })

    results_df = pd.DataFrame(results)
    results_df.to_csv("model_evaluation.csv", index=False)
    print("Evaluation completed: model_evaluation.csv")

if __name__ == "__main__":
    evaluate()