import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
import cv2 
import os
from PIL import Image

# -----------------------------------------------------
# GLOBALS
# -----------------------------------------------------
cnn_model = None
gradcam_model = None
IMAGE_SIZE = (224, 224) 
LAST_CONV_LAYER_NAME = 'conv2' 

# -----------------------------------------------------
# LOAD MODEL & PREPARE GRAD-CAM
# -----------------------------------------------------
def load_cnn_model():
    """Loads the pre-trained CNN model and prepares Grad-CAM functional model."""
    global cnn_model, gradcam_model
    try:
        if not os.path.exists('cnn_model.h5'):
             raise FileNotFoundError("cnn_model.h5 is missing. Run create_dummy_cnn.py first.")
             
        cnn_model = tf.keras.models.load_model('cnn_model.h5')
        
        # CRITICAL FIX 1: Re-compile the model for reliable gradient calculation
        cnn_model.compile(optimizer='adam', loss='binary_crossentropy', run_eagerly=False)
        
        # Grad-CAM model outputs last conv layer + prediction
        gradcam_model = tf.keras.models.Model(
            inputs=cnn_model.input,
            outputs=[cnn_model.get_layer(LAST_CONV_LAYER_NAME).output, cnn_model.output]
        )
        
        print(f"✅ CNN model loaded. Last Conv Layer: {LAST_CONV_LAYER_NAME}")
    except Exception as e:
        print(f"FATAL ERROR: Could not load or compile CNN model: {e}")
        cnn_model = None

# Load immediately
load_cnn_model()

# -----------------------------------------------------
# IMAGE PROCESSING & PREDICTION
# -----------------------------------------------------
def get_img_array(img_path, size):
    """Utility function to load image and format it for the model."""
    img = image.load_img(img_path, target_size=size)
    array = image.img_to_array(img)
    array = np.expand_dims(array, axis=0)
    
    # CRITICAL FIX 2: Ensure array is float32 before normalization/prediction
    array = tf.cast(array, dtype=tf.float32)
    
    array /= 255.0 
    return array

def predict_xray_risk(img_path):
    if cnn_model is None: return 0.0
    try:
        img_array = get_img_array(img_path, IMAGE_SIZE)
        prediction = cnn_model.predict(img_array)[0]
        return float(prediction[0])
    except Exception as e:
        print(f"CNN Prediction Error: {e}")
        return 0.0

# -----------------------------------------------------
# GRAD-CAM FUNCTIONS (Ultimate Stability Fixes)
# -----------------------------------------------------

def make_gradcam_heatmap(img_path, pred_index=None):
    """Generates the Grad-CAM heatmap array with stability fixes."""
    if gradcam_model is None: return np.zeros(IMAGE_SIZE[:2])
    
    img_array = get_img_array(img_path, IMAGE_SIZE)
    
    with tf.GradientTape() as tape:
        last_conv_output, preds = gradcam_model(img_array)
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]
    
    grads = tape.gradient(class_channel, last_conv_output)
    pooled_grads = tf.reduce_mean(grads, axis=(0,1,2))
    
    # Use stable multiplication for heatmap generation
    last_conv_output_tensor = last_conv_output[0]
    
    # Final stable multiplication
    heatmap = last_conv_output_tensor * pooled_grads 
    heatmap = tf.reduce_sum(heatmap, axis=-1)
    
    # Safe normalization
    heatmap = tf.maximum(heatmap, 0)
    max_val = tf.reduce_max(heatmap)
    # CRITICAL FIX 3: Prevent division by zero if max_val is 0.0
    if max_val == 0: max_val = tf.constant(1e-10, dtype=tf.float32)
    heatmap /= max_val
    
    return heatmap.numpy()

def save_gradcam_overlay(img_path, heatmap, save_filename, alpha=0.4):
    """Overlays the heatmap onto the original image and saves it to the static folder."""
    
    static_dir = os.path.join(os.getcwd(), 'static')
    os.makedirs(static_dir, exist_ok=True)
    
    save_path = os.path.join(static_dir, save_filename)
    
    # Core OpenCV Image Processing Logic (The stable part)
    img = Image.open(img_path).convert('RGB')
    heatmap = cv2.resize(heatmap, (img.width, img.height))
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)

    img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    
    # Use cv2.addWeighted for cleaner overlay math
    superimposed_img = cv2.addWeighted(img_cv, 1-alpha, heatmap, alpha, 0)
    
    cv2.imwrite(save_path, superimposed_img)
    return save_filename

def generate_and_save_gradcam(original_xray_path):
    """Runs the full Grad-CAM pipeline and returns the FILENAME to the saved image."""
    if gradcam_model is None: return None
    base_name = os.path.basename(original_xray_path)
    gradcam_filename = f"gradcam_{base_name}"
    
    try:
        heatmap = make_gradcam_heatmap(original_xray_path, pred_index=0) 
        saved_file = save_gradcam_overlay(original_xray_path, heatmap, gradcam_filename)
        return saved_file
    except Exception as e:
        # Prints error to console if something failed during processing/saving
        print(f"GRAD-CAM Generation Failed: {e}") 
        return None