from flask import Flask, request, jsonify, send_file
from tensorflow.keras.models import load_model, Model
from PIL import Image
import numpy as np
import os
import cv2
import tensorflow as tf
import firebase_admin
from firebase_admin import credentials, db
from datetime import datetime

app = Flask(__name__)

# ✅ 1. Initialize Firebase
cred = credentials.Certificate("glaucoma-4b682-firebase-adminsdk-fbsvc-cd31fbe99d.json")  # Path to your service account JSON
firebase_admin.initialize_app(cred, {
    'databaseURL': 'https://glaucoma-4b682-default-rtdb.firebaseio.com/'
})
results_ref = db.reference('results')  # Will save results here

# ✅ 2. Load the Model
model = load_model('mobilenet_glaucoma_model.h5', compile=False)

# ✅ 3. Preprocess Image
def preprocess_image(img):
    img = img.resize((224, 224))
    img = np.array(img) / 255.0
    img = np.expand_dims(img, axis=0)
    return img

# ✅ 4. Grad-CAM Generation
def make_gradcam(img_array, model, last_conv_layer_name='Conv2D_1'):
    """Generate Grad-CAM for the given image and model."""
    last_conv_layer = model.get_layer(last_conv_layer_name)
    grad_model = Model(inputs=model.inputs, outputs=[last_conv_layer.output, model.output])

    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(img_array)
        loss = predictions[:, 0]
    grads = tape.gradient(loss, conv_outputs)

    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    conv_outputs = conv_outputs[0]

    for i in range(conv_outputs.shape[-1]):
        conv_outputs[..., i] *= pooled_grads[i]

    heatmap = tf.reduce_mean(conv_outputs, axis=-1).numpy()
    heatmap = np.maximum(heatmap, 0)
    heatmap /= np.max(heatmap)

    return heatmap

# ✅ 5. Save Grad-CAM Overlay
def save_gradcam_image(original_img, heatmap, filename='gradcam.png', output_dir='results'):
    """Save the Grad-CAM overlay image and return its path."""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    img = np.array(original_img.resize((224, 224)))
    heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
    heatmap = np.uint8(255 * heatmap)

    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    overlay = cv2.addWeighted(img, 0.6, heatmap, 0.4, 0)

    filepath = os.path.join(output_dir, filename)
    cv2.imwrite(filepath, overlay)

    return filepath

@app.route('/')
def home():
    return "Glaucoma Detection Flask API is running!"

@app.route("/test_file")
def test_file():
    """Check if the Firebase service account JSON is present and readable."""
    filepath = "glaucoma-4b682-firebase-adminsdk-fbsvc-cd31fbe99d.json"
    if os.path.exists(filepath):
        return f"✅ Service account file found at: {filepath}"
    else:
        return "❌ Service account JSON NOT found."

@app.route('/predict', methods=['POST'])
def predict():
    """Perform prediction and save results to Firebase."""
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400

    try:
        img = Image.open(file.stream).convert('RGB')
        img_array = preprocess_image(img)

        prediction = model.predict(img_array)[0]
        glaucoma_prob = 1 - prediction[0]
        normal_prob = prediction[0]
        result = 'Glaucoma' if glaucoma_prob > normal_prob else 'Normal'
        confidence = float(glaucoma_prob) if result == 'Glaucoma' else float(normal_prob)

        # Grad-CAM
        heatmap = make_gradcam(img_array, model, last_conv_layer_name='Conv2D_1')
        gradcam_filename = f"gradcam_{int(datetime.now().timestamp())}.png"
        save_gradcam_image(img, heatmap, filename=gradcam_filename)

        # Save to Firebase
        results_ref.push({
            'image_filename': file.filename,
            'prediction': result,
            'confidence': confidence,
            'gradcam_filename': gradcam_filename,
            'timestamp': datetime.now().isoformat()
        })

        return jsonify({
            'prediction': result,
            'confidence': confidence,
            'normal_probability': float(normal_prob),
            'glaucoma_probability': float(glaucoma_prob),
            'gradcam_image': gradcam_filename
        })

    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/results', methods=['GET'])
def results():
    """List all results from the Firebase database."""
    results_data = results_ref.get()
    if not results_data:
        results_data = []
    return jsonify(results_data)

@app.route('/gradcam/<filename>')
def get_gradcam(filename):
    """Serve the Grad-CAM overlay image."""
    filepath = os.path.join('results', filename)
    if os.path.exists(filepath):
        return send_file(filepath, mimetype='image/png')
    else:
        return jsonify({'error': 'File not found'}), 404

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)