from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
import numpy as np
from PIL import Image
import tensorflow as tf
import joblib
import json

app = FastAPI()

# Enable CORS (important for Laravel)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load class names
with open("class_indices.json", "r") as f:
    CLASS_NAMES = list(json.load(f).keys())

# ===== LOAD MODELS =====
print("Loading models...")

model_resnet = tf.keras.models.load_model("model_resnet.keras")
model_cnn = tf.keras.models.load_model("model_cnn.keras")
model_efficient = tf.keras.models.load_model("model_efficientnet.keras")

svm_model = joblib.load("svm_model.pkl")
rf_model = joblib.load("random_forest.pkl")

print("All models loaded!")


# ===== IMAGE PREPROCESSING =====
def preprocess_image(img: Image.Image):
    img = img.resize((224, 224))
    arr = np.array(img).astype("float32") / 255.0
    return np.expand_dims(arr, 0)


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    img = Image.open(file.file).convert("RGB")
    x = preprocess_image(img)

    # Deep Learning Predictions
    # pred_resnet = model_resnet.predict(x)
    # pred_cnn = model_cnn.predict(x)
    # pred_efficient = model_efficient.predict(x)

    res_resnet = CLASS_NAMES[int(np.argmax(pred_resnet))]
    res_cnn = CLASS_NAMES[int(np.argmax(pred_cnn))]
    res_efficient = CLASS_NAMES[int(np.argmax(pred_efficient))]

    # Machine Learning
    # flat = x.reshape(1, -1)  # flatten for ML models
    # pred_svm = svm_model.predict(flat)[0]
    # pred_rf = rf_model.predict(flat)[0]

    return {
        "resnet": {
            "prediction": res_resnet,
            "confidence": float(np.max(pred_resnet))
        },
        # "cnn": {
        #     "prediction": res_cnn,
        #     "confidence": float(np.max(pred_cnn))
        # },
        # "efficientnet": {
        #     "prediction": res_efficient,
        #     "confidence": float(np.max(pred_efficient))
        # },
        # "svm": str(pred_svm),
        # "random_forest": str(pred_rf)
    }