import torch
import torchvision.models as models
import torch.nn as nn
from fastapi import FastAPI, UploadFile, File, HTTPException
from PIL import Image
import io
import torchvision.transforms as transforms

app = FastAPI(title="Alzheimer MRI API")

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# -----------------------------
# MODEL LOADER
# -----------------------------
def load_model(path, num_classes=3):
    model = models.densenet121(weights=None)
    model.classifier = nn.Linear(model.classifier.in_features, num_classes)
    model.load_state_dict(torch.load(path, map_location=DEVICE))
    model.to(DEVICE)
    model.eval()
    return model


# -----------------------------
# LOAD ALL MODELS
# -----------------------------
MODELS = {
    "ad_dn121": load_model("alzheimers_densenet121.pth"),
    "ad_dn169": load_model("alzheimers_densenet169.pth"),
    "ad_dn201": load_model("alzheimers_densenet201.pth"),
}


# -----------------------------
# IMAGE TRANSFORM
# -----------------------------
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])


# -----------------------------
# PREDICTION FUNCTION
# -----------------------------
def predict(model, image_bytes):
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    image = transform(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        outputs = model(image)
        probs = torch.softmax(outputs, dim=1)

    return probs[0].tolist()


# -----------------------------
# ROUTES
# -----------------------------
@app.get("/")
def home():
    return {"message": "Alzheimer MRI API is running"}


@app.post("/predict")
async def predict_image(model_key: str, file: UploadFile = File(...)):

    if model_key not in MODELS:
        raise HTTPException(status_code=400, detail="Invalid model key")

    image_bytes = await file.read()

    result = predict(MODELS[model_key], image_bytes)

    return {
        "model": model_key,
        "prediction": result
    }