from fastapi import FastAPI, UploadFile, File
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import pickle
import logging

# ------------------- Logging Setup -------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ------------------- Load Model -------------------
try:
    model = tf.keras.models.load_model('MobilenetV2.keras')
    logger.info("Model loaded successfully")
except Exception as e:
    logger.error(f"Error loading model: {e}")
    raise

IMG_SIZE = (224, 224)

# ------------------- Load Classes -------------------
try:
    with open('classes.pkl', 'rb') as f:
        class_dict = pickle.load(f)  # e.g., {'Anthracnose': 0, 'Bacterial Canker': 1, ...}
    # Inverse mapping: index -> class name
    idx_to_class = {v: k for k, v in class_dict.items()}
    logger.info(f"Classes loaded: {idx_to_class}")
except Exception as e:
    logger.error(f"Error loading classes.pkl: {e}")
    raise

# ------------------- FastAPI App -------------------
app = FastAPI(title="Mango Diseases Detection API", version="0.1.0")

# ------------------- Image Preprocessing -------------------
def preprocess_image(image_bytes):
    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        image = image.resize(IMG_SIZE)
        image_array = np.array(image) / 255.0
        image_array = np.expand_dims(image_array, axis=0)
        logger.info(f"Image preprocessed: shape {image_array.shape}")
        return image_array
    except Exception as e:
        logger.error(f"Error preprocessing image: {e}")
        raise

# ------------------- Routes -------------------
@app.get("/")
def home():
    return {"message": "Welcome to the Mango Diseases Detection API"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        image_bytes = await file.read()
        logger.info(f"Received file: {file.filename}, size: {len(image_bytes)} bytes")

        input_data = preprocess_image(image_bytes)
        preds = model.predict(input_data)
        logger.info(f"Raw predictions: {preds}")

        predicted_index = np.argmax(preds[0])
        predicted_class = idx_to_class[predicted_index]
        confidence = float(np.max(preds[0]))
        logger.info(f"Predicted class: {predicted_class}, confidence: {confidence}")

        return {
            "predicted_class": predicted_class,
            "confidence": confidence
        }

    except Exception as e:
        logger.error(f"Error during prediction: {e}")
        return {"error": str(e)}