import os
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Headless backend
import matplotlib.pyplot as plt
from PIL import Image
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.convnext import preprocess_input
from sklearn.metrics import confusion_matrix, accuracy_score, ConfusionMatrixDisplay

# ------------------------------
# This is the verification for convnext_20260228_90.h5.keras
# We are currently at 99% accuracy in the predictions with the data
# ------------------------------


# ----------------------------
# Paths
# ----------------------------
# MODEL_PATH = "/home/pi/CNN/Weights/convnext_20260228_90.h5.keras"
# TEST_DIR = "/home/pi/DeepLearning/Data"  # test images structured by class
# OUTPUT_DIR = "/home/pi/CNN/Models/evaluation_outputs"

MODEL_PATH = "c:\boneyard\DeepLearning\CNN\Weights\convnext_20260228_90.h5.keras"
TEST_DIR = "c:\boneyard\DeepLearning\CNN\Data"  # test images structured by class
OUTPUT_DIR = "c:\boneyard\DeepLearning\CNN\Models\evaluation_outputs"

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ----------------------------
# Suppress TensorFlow warnings
# ----------------------------
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # only errors

# ----------------------------
# Test Data Generator
# ----------------------------
test_datagen = ImageDataGenerator(preprocessing_function=preprocess_input)

test_generator = test_datagen.flow_from_directory(
    TEST_DIR,
    target_size=(224, 224),
    batch_size=32,
    class_mode='binary',
    shuffle=False
)

# ----------------------------
# Load Model
# ----------------------------
model = load_model(MODEL_PATH)
print("Model loaded successfully.")

# ----------------------------
# Predictions
# ----------------------------
y_true = test_generator.classes
y_pred_probs = model.predict(test_generator, verbose=1)
y_pred = (y_pred_probs > 0.5).astype(int).reshape(-1)

# ----------------------------
# Accuracy
# ----------------------------
accuracy = accuracy_score(y_true, y_pred)
print(f"Test Accuracy: {accuracy*100:.2f}%")

# ----------------------------
# Confusion Matrix
# ----------------------------
cm = confusion_matrix(y_true, y_pred)
disp = ConfusionMatrixDisplay(
    confusion_matrix=cm,
    display_labels=list(test_generator.class_indices.keys())
)
disp.plot(cmap=plt.cm.Blues)
plt.title("Confusion Matrix")
conf_matrix_path = os.path.join(OUTPUT_DIR, "confusion_matrix.png")
plt.savefig(conf_matrix_path, dpi=300)
plt.close()
print(f"Confusion matrix saved to: {conf_matrix_path}")

# ----------------------------
# Misclassified Images
# ----------------------------
misclassified_idx = np.where(y_true != y_pred)[0]
print(f"Number of misclassified images: {len(misclassified_idx)}")

MAX_IMAGES = 20  # Number of misclassified images to save
for i, idx in enumerate(misclassified_idx[:MAX_IMAGES]):
    img_path = test_generator.filepaths[idx]

    # Open image with Pillow
    img = Image.open(img_path).convert('RGB')

    plt.imshow(np.array(img))
    true_label = list(test_generator.class_indices.keys())[y_true[idx]]
    pred_label = list(test_generator.class_indices.keys())[y_pred[idx]]
    plt.title(f"True: {true_label}, Pred: {pred_label}")
    plt.axis('off')

    save_path = os.path.join(OUTPUT_DIR, f"misclassified_{i+1}.png")
    plt.savefig(save_path, dpi=200)
    plt.close()
    print(f"Saved misclassified image: {save_path}")