import os
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import os
import numpy as np
import tensorflow as tf

from tensorflow.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array      
from tensorflow.keras.models import load_model

from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import cv2

from io import BytesIO
from PIL import Image
import requests

MODEL_PATH = r'D:\Studies\Sem-7\Deep Learning Concepts and Architectures\Exercise\Exercise-5\Output\Xception_Transfer_Learning.hdf5'
EVAL_DIR   = r'D:\Studies\Sem-7\Deep Learning Concepts and Architectures\Exercise\OCT_Dataset\OCT_Dataset\val'
IMG_SIZE   = (299, 299)
BATCH_SIZE = 32
LAST_CONV  = 'block14_sepconv2_act'

model = load_model(MODEL_PATH)

eval_gen = ImageDataGenerator(rescale=1./255).flow_from_directory(
    EVAL_DIR, target_size=IMG_SIZE, batch_size=BATCH_SIZE,
    class_mode='categorical', shuffle=False)

classes = list(eval_gen.class_indices.keys())
print("Classes:", classes)

loss, acc = model.evaluate(eval_gen, verbose=1)
print(f"Eval Loss: {loss:.4f} | Eval Acc: {acc:.4f}")

probs = model.predict(eval_gen, verbose=1)
y_pred = np.argmax(probs, axis=1)
y_true = eval_gen.classes

cm = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(7,6))
sns.heatmap(cm, annot=True, fmt='d', cmap='YlGnBu',
            xticklabels=classes, yticklabels=classes)
plt.xlabel('Predicted'); plt.ylabel('Actual')
plt.title('Confusion Matrix (Eval)')
plt.tight_layout()
plt.savefig('confusion_matrix_eval.png', dpi=150)
plt.show()

print("Classification report (Eval):")
print(classification_report(y_true, y_pred, target_names=classes, digits=4))


def preprocess_input_pil(im):
    x = img_to_array(im) / 255.0
    return np.expand_dims(x, axis=0)


def predict_from_image_path(image_path):
    im = load_img(image_path, target_size=IMG_SIZE)
    x = preprocess_input_pil(im)
    probs = model.predict(x, verbose=0)
    idx = int(np.argmax(probs, axis=1)[0])
    conf = float(np.max(probs))
    return idx, classes[idx], conf


print(predict_from_image_path(
    r'D:\Studies\Sem-7\Deep Learning Concepts and Architectures\Exercise\OCT_Dataset\OCT_Dataset\val\DME\dme_val_1001.jpg'
))

# def predict_from_image_url(image_url):
#     res = requests.get(image_url, timeout=10)
#     im = Image.open(BytesIO(res.content)).convert('RGB').resize(IMG_SIZE)
#     x = preprocess_input_pil(im)
#     probs = model.predict(x, verbose=0)
#     idx = int(np.argmax(probs, axis=1)[0])
#     conf = float(np.max(probs))
#     return idx, classes[idx], conf


# def grad_cam(image_path, last_conv_name=LAST_CONV, intensity=0.5):
#     try:
#         last_conv_layer = model.get_layer(last_conv_name)
#     except ValueError:
#         conv_candidates = [l for l in model.layers if hasattr(l, 'output_shape') and len(l.output_shape) == 4]
#         if not conv_candidates:
#             raise ValueError("No 4D conv layer found for Grad-CAM.")
#         last_conv_layer = conv_candidates[-1]

#     cam_model = tf.keras.Model(model.input, [last_conv_layer.output, model.output])

#     im = load_img(image_path, target_size=IMG_SIZE)
#     x = preprocess_input_pil(im)

#     with tf.GradientTape() as tape:
#         conv_out, preds = cam_model(x, training=False)
#         if isinstance(preds, (list, tuple)):
#             preds = preds[0]
#         preds = tf.convert_to_tensor(preds)
#         idx = tf.argmax(preds[0])
#         class_score = tf.gather(preds[0], idx)

#     grads = tape.gradient(class_score, conv_out)
#     weights = tf.reduce_mean(grads, axis=(0, 1, 2))
#     cam = tf.reduce_sum(weights * conv_out[0], axis=-1)
#     cam = tf.maximum(cam, 0) / (tf.reduce_max(cam) + 1e-8)
#     cam = cam.numpy()

#     img_cv = cv2.imread(image_path)
#     if img_cv is None:
#         img_cv = cv2.cvtColor(np.array(im), cv2.COLOR_RGB2BGR)

#     cam = cv2.resize(cam, (img_cv.shape[1], img_cv.shape[0]))
#     heatmap = np.uint8(255 * cam)
#     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
#     overlay = cv2.addWeighted(img_cv, 1.0, heatmap, intensity, 0)
#     out_path = './gradcam_tmp.jpg'
#     cv2.imwrite(out_path, overlay)
#     plt.figure(figsize=(12, 6))
#     plt.imshow(plt.imread(out_path))
#     plt.axis('off')
#     plt.show()


# grad_cam(
#     r'D:\Studies\Sem-7\Deep Learning Concepts and Architectures\Exercise\OCT_Dataset\OCT_Dataset\val\DME\dme_val_1001.jpg',
#     last_conv_name=LAST_CONV,
#     intensity=0.5
# )

