import json
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

from sklearn.metrics import classification_report, roc_auc_score, average_precision_score, f1_score, roc_curve, confusion_matrix, ConfusionMatrixDisplay, precision_recall_curve
from sklearn.preprocessing import LabelBinarizer
from collections import defaultdict

# Load the prediction and ground truth JSON files
with open('pred.json', 'r') as pred_file:
    predictions = json.load(pred_file)

with open('GT.json', 'r') as gt_file:
    ground_truth = json.load(gt_file)

# Convert to DataFrame
pred_df = pd.DataFrame(predictions)
gt_df = pd.DataFrame(ground_truth)

# Create unique class name
pred_df['unique_class'] = pred_df['damage_class'] + '_' + pred_df['damage_level']
gt_df['unique_class'] = gt_df['damage_class'] + '_' + gt_df['damage_level']

# Merge predictions with ground truth based on image_path
merged_df = pd.merge(pred_df, gt_df, on='image_path', suffixes=('_pred', '_gt'))

# Extract true and predicted classes
y_true = merged_df['unique_class_gt']
y_pred = merged_df['unique_class_pred']

# Count the number of unique classes
unique_classes = set(y_true).union(set(y_pred))
num_unique_classes = len(unique_classes)
print(f"Total number of unique classes: {num_unique_classes}")

# Initialize the LabelBinarizer
lb = LabelBinarizer()
y_true_bin = lb.fit_transform(y_true)
y_pred_bin = lb.transform(y_pred)

# Classification Report
classification_rep = classification_report(y_true, y_pred, digits=6)
print("Classification Report:")
print(classification_rep)

# ROC-AUC Score (one-vs-rest)
roc_auc = roc_auc_score(y_true_bin, y_pred_bin, average='macro')
print(f"ROC-AUC Score: {roc_auc:.6f}")

# Mean Average Precision
mean_ap = average_precision_score(y_true_bin, y_pred_bin, average='macro')
print(f"Mean Average Precision: {mean_ap:.6f}")

# F1 Score
f1 = f1_score(y_true, y_pred, average='macro')
print(f"F1 Score: {f1:.6f}")

# Plot ROC-AUC curve
plt.figure(figsize=(12, 8))
roc_auc_values = {}
for i in range(num_unique_classes):
    fpr, tpr, _ = roc_curve(y_true_bin[:, i], y_pred_bin[:, i])
    roc_auc_value = roc_auc_score(y_true_bin[:, i], y_pred_bin[:, i])
    roc_auc_values[lb.classes_[i]] = roc_auc_value
    plt.plot(fpr, tpr, lw=2, label=f'Class {lb.classes_[i]} (area = {roc_auc_value:.6f})')

plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC-AUC Curve')
plt.legend(loc='lower right')
plt.savefig('roc_auc_curve.png')
plt.show()

# Plot Confusion Matrix
cm = confusion_matrix(y_true, y_pred, labels=lb.classes_)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=lb.classes_)

fig, ax = plt.subplots(figsize=(12, 12))
disp.plot(ax=ax, cmap=plt.cm.Blues)
plt.xticks(rotation=90)
plt.title('Confusion Matrix')
plt.savefig('confusion_matrix.png')
plt.show()

# Plot Precision-Recall Curve for each class
plt.figure(figsize=(12, 8))
average_precisions = {}
for i in range(num_unique_classes):
    precision, recall, _ = precision_recall_curve(y_true_bin[:, i], y_pred_bin[:, i])
    average_precision = average_precision_score(y_true_bin[:, i], y_pred_bin[:, i])
    average_precisions[lb.classes_[i]] = average_precision
    plt.plot(recall, precision, lw=2, label=f'Class {lb.classes_[i]} (AP = {average_precision:.6f})')

plt.xlabel('Recall')
plt.ylabel('Precision')
plt.title('Precision-Recall Curve')
plt.legend(loc='lower left')
plt.savefig('precision_recall_curve.png')
plt.show()

# Save combined performance metrics and class-wise metrics to a single text file
with open('performance_metrics.txt', 'w') as file:
    file.write("Classification Report:\n")
    file.write(classification_rep + "\n")
    file.write(f"ROC-AUC Score: {roc_auc:.6f}\n")
    file.write(f"Mean Average Precision: {mean_ap:.6f}\n")
    file.write(f"F1 Score: {f1:.6f}\n\n")
    
    file.write("Class-wise ROC-AUC values:\n")
    for cls, roc_auc_value in roc_auc_values.items():
        file.write(f"{cls}: {roc_auc_value:.6f}\n")
    file.write(f"Overall ROC-AUC: {roc_auc:.6f}\n\n")

    file.write("Class-wise Average Precision (AP) values:\n")
    for cls, average_precision in average_precisions.items():
        file.write(f"{cls}: {average_precision:.6f}\n")
    file.write(f"Overall Mean Average Precision (mAP): {mean_ap:.6f}\n")