import numpy as np
import pandas as pd
import cv2
import pydicom
from flask import Flask, request, jsonify
from tensorflow.keras.models import load_model
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force CPU mode for TensorFlow

# Initialize Flask app
Pneumonia_predictor_api = Flask("Pneumonia predictor")

# Load trained model (must match TF/Keras version used during training)
model = load_model("final_model.keras")


@Pneumonia_predictor_api.get('/')
def home():
    return "Welcome to the Pneumonia Prediction API!"


@Pneumonia_predictor_api.route('/v1/Pneumonia_prediction', methods=['POST'])
def predict_pneumonia():

    # --------- STEP 1: CHECK FILE ---------
    if 'file' not in request.files:
        return jsonify({'error': 'No DICOM file uploaded'}), 400

    dcm_file = request.files['file']
    dcm_file.seek(0)  # ensure pointer at beginning

    # --------- STEP 2: READ DICOM SAFELY ---------
    try:
        dicom_data = pydicom.dcmread(dcm_file, force=True)  # force=True fixes header issues
        image = dicom_data.pixel_array.astype(np.float32)
    except Exception as e:
        return jsonify({"error": f"Invalid DICOM file: {str(e)}"}), 400

    # --------- STEP 3: PREPROCESS IMAGE ---------
    image = cv2.resize(image, (224, 224))

    # Convert grayscale → RGB 3-channel
    if len(image.shape) == 2:
        image = np.stack([image] * 3, axis=-1)

    # Normalize 0–1
    max_val = np.max(image)
    if max_val > 0:
        image = image / max_val

    image = np.expand_dims(image, axis=0)  # (1,224,224,3)

    # --------- STEP 4: PREDICT ---------
    try:
        prob = float(model.predict(image)[0][0])  # probability
        pred_class = int(prob >= 0.5)             # predicted class
    except Exception as e:
        return jsonify({"error": f"Model prediction failed: {str(e)}"}), 500
    
    return jsonify({
        "Predicted_Class": pred_class,
        "Probability": prob,
        "Message": "Prediction completed"
    })


# Run Flask locally
if __name__ == '__main__':
    Pneumonia_predictor_api.run(debug=True)