import os
import io
import numpy as np
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import tensorflow as tf

# Patch Keras Layer initialization to pop quantization_config (avoiding deserialize issues in some Keras versions)
try:
    import keras
    original_layer_init = keras.layers.Layer.__init__
    def patched_layer_init(self, *args, **kwargs):
        kwargs.pop("quantization_config", None)
        original_layer_init(self, *args, **kwargs)
    keras.layers.Layer.__init__ = patched_layer_init
except Exception as e:
    pass

try:
    original_tf_layer_init = tf.keras.layers.Layer.__init__
    def patched_tf_layer_init(self, *args, **kwargs):
        kwargs.pop("quantization_config", None)
        original_tf_layer_init(self, *args, **kwargs)
    tf.keras.layers.Layer.__init__ = patched_tf_layer_init
except Exception as e:
    pass

app = FastAPI(title="Parkinson's Detection API", description="API to classify Parkinson's from images")

# Setup CORS to allow requests from the Cloudflare frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust this in production to match your Cloudflare Pages URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables to hold the model and status/error details
MODEL = None
MODEL_LOAD_STATUS = "Not started"

@app.on_event("startup")
async def load_model():
    global MODEL, MODEL_LOAD_STATUS
    model_path = "final_parkinson_model.keras"
    MODEL_LOAD_STATUS = f"Checking path: {os.path.abspath(model_path)}"
    if os.path.exists(model_path):
        try:
            MODEL_LOAD_STATUS = "Model file exists. Attempting to load..."
            MODEL = tf.keras.models.load_model(model_path)
            MODEL_LOAD_STATUS = "Loaded successfully"
            print("Model loaded successfully.")
        except Exception as e:
            import traceback
            error_msg = f"Error loading model: {str(e)}\n{traceback.format_exc()}"
            MODEL_LOAD_STATUS = error_msg
            print(error_msg)
    else:
        error_msg = f"Model file {model_path} not found. Files in current directory: {os.listdir('.')}"
        MODEL_LOAD_STATUS = error_msg
        print(error_msg)

def preprocess_image(image_bytes):
    try:
        # Load image, convert to grayscale (L) to match training's cv2.IMREAD_GRAYSCALE,
        # then convert to RGB to duplicate the grayscale channel to 3 channels.
        image = Image.open(io.BytesIO(image_bytes))
        image = image.convert("L")
        image = image.resize((128, 128))
        image = image.convert("RGB")
        
        # Convert to numpy array and apply MobileNetV2 preprocessing (scales to [-1, 1])
        # This matches the training pipeline which used tf.keras.applications.mobilenet_v2.preprocess_input
        img_array = np.array(image, dtype=np.float32)
        img_array = tf.keras.applications.mobilenet_v2.preprocess_input(img_array)
        # Expand dimensions to create a batch of 1: (1, 128, 128, 3)
        img_array = np.expand_dims(img_array, axis=0)
        return img_array
    except Exception as e:
        raise ValueError(f"Invalid image file: {e}")

@app.get("/")
def read_root():
    return {
        "status": "healthy",
        "model_loaded": MODEL is not None,
        "model_load_status": MODEL_LOAD_STATUS
    }

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    if MODEL is None:
        raise HTTPException(status_code=503, detail="Model is not loaded.")
    
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image.")
    
    try:
        image_bytes = await file.read()
        processed_image = preprocess_image(image_bytes)
        
        # Make prediction
        predictions = MODEL.predict(processed_image)
        # Based on the notebook, it's a binary classification with 2 dense output units.
        # Or a single unit. Let's return the raw predictions and class.
        
        # Determine class 0 = healthy, 1 = parkinson (based on standard alphabetical ordering of classes, but could vary)
        class_idx = int(np.argmax(predictions[0]))
        confidence = float(predictions[0][class_idx])
        
        # We assume output layer has 2 nodes: e.g., index 0 = healthy, index 1 = parkinson.
        classes = ["Healthy", "Parkinson"]
        result_class = classes[class_idx] if class_idx < len(classes) else "Unknown"
        
        return {
            "prediction": result_class,
            "confidence": confidence,
            "raw_scores": predictions[0].tolist()
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)