import pandas as pd
import numpy as np
import pickle
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import json
from sklearn.metrics import accuracy_score, confusion_matrix
import warnings

warnings.filterwarnings("ignore")


def test_model():
    # Load model and metadata
    with open("models/cve_severity_model.pkl", "rb") as f:
        model = pickle.load(f)
    with open("models/model_metadata.json", "r") as f:
        metadata = json.load(f)

    # Load test data
    test_df = pd.read_csv("data/cve_testing_dataset.csv")
    feature_columns = metadata["feature_columns"]
    severity_map = {"IMPORTANT": 0, "MODERATE": 1, "LOW": 2}
    severity_labels = ["IMPORTANT", "MODERATE", "LOW"]

    # Prepare features
    X_test = test_df[feature_columns].fillna(0).astype(int)
    y_test = test_df["severity_label"].map(severity_map)

    # Make predictions
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)

    print(f"🧪 Model: {metadata.get('model_type', 'Unknown')}")
    print(f"📊 Test Accuracy: {accuracy:.3f}")
    print(f"🎯 Test CVEs: {len(test_df)}")

    # Create visualization
    fig, axes = plt.subplots(2, 2, figsize=(14, 12))

    # 1. Confusion Matrix
    ax1 = axes[0, 0]
    cm = confusion_matrix(y_test, y_pred)

    cm_percent = np.zeros_like(cm, dtype=float)
    for i in range(cm.shape[0]):
        if cm[i].sum() > 0:
            cm_percent[i] = cm[i] / cm[i].sum() * 100

    sns.heatmap(
        cm,
        annot=False,
        cmap="Blues",
        ax=ax1,
        xticklabels=severity_labels,
        yticklabels=severity_labels,
    )

    # Add custom annotations with count and percentage
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            count = cm[i, j]
            percent = cm_percent[i, j]
            text_color = "white" if percent > 200 else "black"
            ax1.text(
                j + 0.5,
                i + 0.5,
                f"{count}\n({percent:.1f}%)",
                ha="center",
                va="center",
                color=text_color,
                fontweight="bold",
            )

    model_name = metadata.get("model_type", "XGBoost Balanced")
    ax1.set_title(
        f"{model_name.replace('_', ' ')}\nConfusion Matrix", fontsize=12, weight="bold"
    )
    ax1.set_xlabel("Predicted Label")
    ax1.set_ylabel("True Label")

    # 2. Class-wise Recall
    ax2 = axes[0, 1]
    recalls = [cm[i, i] / cm[i].sum() * 100 if cm[i].sum() > 0 else 0 for i in range(3)]
    bars = ax2.bar(
        severity_labels, recalls, color=["red", "orange", "green"], alpha=0.7
    )
    ax2.set_title("Class-wise Recall (%)", fontsize=12, weight="bold")
    ax2.set_ylabel("Recall Percentage")
    ax2.set_ylim(0, 100)

    for bar, recall in zip(bars, recalls):
        ax2.text(
            bar.get_x() + bar.get_width() / 2.0,
            bar.get_height() + 1,
            f"{recall:.1f}%",
            ha="center",
            va="bottom",
            fontweight="bold",
        )

    # 3. Prediction Distribution by True Class
    ax3 = axes[1, 0]

    # Create stacked bar chart
    bottom = np.zeros(len(severity_labels))
    colors_stack = ["green", "orange", "red"]

    for j, pred_class in enumerate(severity_labels):
        values = []
        for i in range(len(severity_labels)):
            total = cm[i].sum()
            pct = (cm[i, j] / total * 100) if total > 0 else 0
            values.append(pct)

        ax3.bar(
            severity_labels,
            values,
            bottom=bottom,
            label=f"Predicted as {pred_class}",
            color=colors_stack[j],
            alpha=0.7,
        )
        bottom += values

    ax3.set_title("Prediction Distribution by True Class", fontsize=12, weight="bold")
    ax3.set_ylabel("Percentage")
    ax3.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
    ax3.set_ylim(0, 100)

    # 4. Summary Stats
    ax4 = axes[1, 1]
    ax4.axis("off")

    from sklearn.metrics import precision_score, recall_score, f1_score

    precision = precision_score(y_test, y_pred, average="macro", zero_division=0)
    recall = recall_score(y_test, y_pred, average="macro", zero_division=0)
    f1 = f1_score(y_test, y_pred, average="macro", zero_division=0)

    # Count predictions
    pred_counts = {"IMPORTANT": 0, "MODERATE": 0, "LOW": 0}
    true_counts = {"IMPORTANT": 0, "MODERATE": 0, "LOW": 0}

    for true_label, pred_label in zip(y_test, y_pred):
        pred_counts[severity_labels[pred_label]] += 1
        true_counts[severity_labels[true_label]] += 1

    summary_text = f"BEST MODEL: {model_name.replace('_', ' ')}\n"
    summary_text += "=" * 40 + "\n\n"
    summary_text += f"Test Accuracy:  {accuracy:.3f}\n"
    summary_text += f"Precision:      {precision:.3f}\n"
    summary_text += f"Recall:         {recall:.3f}\n"
    summary_text += f"F1-Score:       {f1:.3f}\n\n"
    summary_text += "Training Time:  N/A\n"
    summary_text += "Testing Time:   N/A\n\n"

    summary_text += "CLASS BREAKDOWN:\n"
    summary_text += "-" * 20 + "\n"

    for i, severity in enumerate(severity_labels):
        summary_text += f"{severity}:\n"
        summary_text += f"  True:      {true_counts[severity]:3d}\n"
        summary_text += f"  Predicted: {pred_counts[severity]:3d}\n"
        summary_text += f"  Recall:    {recalls[i]:5.1f}%\n\n"

    ax4.text(
        0.1,
        0.9,
        summary_text,
        transform=ax4.transAxes,
        fontsize=10,
        verticalalignment="top",
        fontfamily="monospace",
    )

    plt.tight_layout()

    # Create test-results directory
    import os

    results_dir = "test-results"
    os.makedirs(results_dir, exist_ok=True)

    # Save visualization
    plt.savefig(f"{results_dir}/test_results.png", dpi=300, bbox_inches="tight")
    plt.show()

    # Save predictions
    predictions = []
    for cve_id, true_label, pred_label in zip(test_df["cve_id"], y_test, y_pred):
        predictions.append(
            f"{cve_id},{severity_labels[true_label]},{severity_labels[pred_label]}"
        )

    with open(f"{results_dir}/predictions.txt", "w") as f:
        f.write("CVE_ID,ACTUAL,PREDICTED\n")
        f.write("\n".join(predictions))

    # Save detailed results as JSON

    results_data = {
        "timestamp": datetime.now().isoformat(),
        "model_type": metadata.get("model_type", "Unknown"),
        "test_accuracy": float(accuracy),
        "test_precision": float(precision),
        "test_recall": float(recall),
        "test_f1_score": float(f1),
        "total_test_cves": len(test_df),
        "class_breakdown": {
            "IMPORTANT": {
                "true": int(true_counts["IMPORTANT"]),
                "predicted": int(pred_counts["IMPORTANT"]),
                "recall": float(recalls[0]),
            },
            "MODERATE": {
                "true": int(true_counts["MODERATE"]),
                "predicted": int(pred_counts["MODERATE"]),
                "recall": float(recalls[1]),
            },
            "LOW": {
                "true": int(true_counts["LOW"]),
                "predicted": int(pred_counts["LOW"]),
                "recall": float(recalls[2]),
            },
        },
    }

    with open(f"{results_dir}/test_summary.json", "w") as f:
        json.dump(results_data, f, indent=2)

    print(f"✅ Results saved in '{results_dir}/' folder:")
    print("   - test_results.png (visualization)")
    print("   - predictions.txt (CVE predictions)")
    print("   - test_summary.json (detailed metrics)")

    return accuracy


if __name__ == "__main__":
    test_model()
