import os

# Force TensorFlow Keras compatibility
os.environ["TF_USE_LEGACY_KERAS"] = "1"

import numpy as np
import tensorflow as tf
from flask import Flask, request, jsonify
from model import build_model

pneumonia_detector_api = Flask(__name__)

IMG_SIZE = 224
MODEL_PATH = "pneumonia_resnet.weights.h5"

print("TensorFlow Version:", tf.__version__)

# -----------------------------
# SAFE MODEL LOADING
# -----------------------------
def load_trained_model():

    try:

        print("Building model architecture...")

        model = build_model()

        print("Loading weights...")

        model.load_weights(
            MODEL_PATH,
            skip_mismatch=True
        )

        print("Model loaded successfully")

        return model

    except Exception as e:

        print("Model loading failed:", e)
        raise RuntimeError("Failed to load model")


model = load_trained_model()


# -----------------------------
# HOME ROUTE
# -----------------------------
@pneumonia_detector_api.route("/")
def home():
    return "Pneumonia Detection API Running"


# -----------------------------
# IMAGE PREPROCESSING
# -----------------------------
def preprocess_image(img_array):

    img = np.array(img_array)

    # Ensure image has channel dimension
    if len(img.shape) == 2:
        img = np.expand_dims(img, axis=-1)

    # Resize image
    img = tf.image.resize(img, (IMG_SIZE, IMG_SIZE)).numpy()

    # Normalize
    img = img.astype("float32") / 255.0

    # Convert grayscale → RGB
    if img.shape[-1] == 1:
        img = np.repeat(img, 3, axis=-1)

    # Add batch dimension
    img = np.expand_dims(img, axis=0)

    return img


# -----------------------------
# PREDICTION API
# -----------------------------
@pneumonia_detector_api.route("/v1/predict", methods=["POST"])
def predict():

    try:

        if "file" not in request.files:
            return jsonify({"error": "No file uploaded"}), 400

        file = request.files["file"]

        # Load .npy file
        img_array = np.load(file)

        processed_img = preprocess_image(img_array)

        print("Input shape to model:", processed_img.shape)

        prediction = model.predict(processed_img)

        probability = float(prediction[0][0])

        label = "Pneumonia" if probability >= 0.5 else "Normal"

        return jsonify({
            "Predicted_Class": label,
            "Probability_Pneumonia": round(probability, 4)
        })

    except Exception as e:

        print("Prediction error:", str(e))

        return jsonify({
            "error": str(e)
        }), 500


# -----------------------------
# RUN SERVER
# -----------------------------
if __name__ == "__main__":

    pneumonia_detector_api.run(
        host="0.0.0.0",
        port=7860,
        debug=False
    )