import numpy as np
import time
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    confusion_matrix, roc_curve, precision_recall_curve, auc
)
from sklearn.preprocessing import StandardScaler, LabelEncoder, label_binarize
from sklearn.pipeline import Pipeline

try:
    import xgboost as xgb
    XGB_AVAILABLE = True
except:
    XGB_AVAILABLE = False
    print("XGBoost not installed")

PSEUDO_PATH = r"path to saved pseudolabeled .npz"
VAL_PATH = r"path to extracted test data as .npz"
SAVE_RESULTS_PATH = r"save results as .csv"

pseudo_data = np.load(PSEUDO_PATH)
X_train = pseudo_data["combined_features"]
y_train = pseudo_data["combined_labels"]

val_data = np.load(VAL_PATH)
X_val = val_data["features"]
y_val = val_data["labels"]

print("Train:", X_train.shape)
print("Validation:", X_val.shape)

le = LabelEncoder()
y_train = le.fit_transform(y_train)
y_val = le.transform(y_val)

def get_models():
    models = {}

    models["Linear SVM"] = Pipeline([
        ("scaler", StandardScaler()),
        ("svm", SVC(kernel="rbf", probability=True, random_state=42))
    ])

    models["Random Forest"] = RandomForestClassifier(
        n_estimators=100,
        random_state=42,
        n_jobs=-1
    )

    if XGB_AVAILABLE:
        models["XGBoost"] = xgb.XGBClassifier(
            n_estimators=100,
            max_depth=6,
            learning_rate=0.1,
            eval_metric="mlogloss",
            random_state=42
        )

    return models

roc_storage = []
pr_storage = []
results = []


def store_multiclass_curves(model, X_val, y_val, model_name):

    classes = np.unique(y_val)
    y_score = model.predict_proba(X_val)

    if len(classes) == 2:
        y_bin = label_binarize(y_val, classes=classes).ravel()
        y_score_bin = y_score[:, 1]

        fpr, tpr, _ = roc_curve(y_bin, y_score_bin)
        roc_auc = auc(fpr, tpr)

        precision, recall, _ = precision_recall_curve(y_bin, y_score_bin)
        pr_auc = auc(recall, precision)

    else:
        y_bin = label_binarize(y_val, classes=classes)

        fpr, tpr, _ = roc_curve(y_bin.ravel(), y_score.ravel())
        roc_auc = auc(fpr, tpr)

        precision, recall, _ = precision_recall_curve(
            y_bin.ravel(),
            y_score.ravel()
        )
        pr_auc = auc(recall, precision)

    roc_storage.append((model_name, fpr, tpr, roc_auc))
    pr_storage.append((model_name, recall, precision, pr_auc))

def evaluate_model(model, name):

    print(f"\nTraining {name}")

    start = time.time()
    model.fit(X_train, y_train)
    train_time = time.time() - start

    start = time.time()
    y_pred = model.predict(X_val)
    test_time = time.time() - start

    store_multiclass_curves(model, X_val, y_val, name)

    acc = accuracy_score(y_val, y_pred)
    prec = precision_score(y_val, y_pred, average="weighted", zero_division=0)
    rec = recall_score(y_val, y_pred, average="weighted", zero_division=0)
    f1 = f1_score(y_val, y_pred, average="weighted", zero_division=0)

    print("Accuracy  :", format(acc, '.4f'))
    print("Precision :", format(prec, '.4f'))
    print("Recall    :", format(rec, '.4f'))
    print("F1 Score  :", format(f1, '.4f'))

    # Confusion Matrix
    cm = confusion_matrix(y_val, y_pred)
    plt.figure(figsize=(6,5))
    sns.heatmap(cm, annot=True, fmt='d', cbar=False)
    plt.title(f"Confusion Matrix: {name}")
    plt.xlabel("Predicted")
    plt.ylabel("True")
    plt.tight_layout()
    plt.show()

    return {
        "Model": name,
        "Accuracy": acc,
        "Precision": prec,
        "Recall": rec,
        "F1": f1,
        "TrainTime": train_time,
        "TestTime": test_time
    }
for name, model in get_models().items():
    results.append(evaluate_model(model, name))

results_df = pd.DataFrame(results)
print("\nFinal Results:\n", results_df)

results_df.to_csv(SAVE_RESULTS_PATH, index=False)

plt.figure(figsize=(8,6))
for name, fpr, tpr, roc_auc in roc_storage:
    plt.plot(fpr, tpr, label=f"{name} (AUC={roc_auc:.3f})")

plt.plot([0,1],[0,1],'--')
plt.title("ROC Curve")
plt.xlabel("False Positive Rate")
plt.ylabel("True Positive Rate")
plt.legend()
plt.grid(True)
plt.show()


plt.figure(figsize=(8,6))
for name, recall, precision, pr_auc in pr_storage:
    plt.plot(recall, precision, label=f"{name} (AUC={pr_auc:.3f})")

plt.title("Precision-Recall Curve")
plt.xlabel("Recall")
plt.ylabel("Precision")
plt.legend()
plt.grid(True)
plt.show()



