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

# =========================
# APP
# =========================
app = FastAPI(title="Alzheimer DATSCAN Ensemble API")

# =========================
# DEVICE
# =========================
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", DEVICE)

# =========================
# CLASSES (VERY IMPORTANT - YOUR ORIGINAL SETUP)
# =========================
CLASSES = [
    "Mild Demented",
    "Moderate Demented",
    "Non Demented",
    "Very Mild Demented"
]

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

# =========================
# MODEL BUILDER (MATCH TRAINING)
# =========================
def build_model(version="121"):
    if version == "121":
        model = models.densenet121(weights=None)
        in_features = 1024
    elif version == "169":
        model = models.densenet169(weights=None)
        in_features = 1664
    else:
        model = models.densenet201(weights=None)
        in_features = 1920

    model.classifier = nn.Sequential(
        nn.Dropout(0.4),
        nn.Linear(in_features, 4)
    )

    return model

# =========================
# SAFE MODEL LOADER (FIX ALL YOUR ERRORS)
# =========================
def load_model(path, version):
    model = build_model(version)

    try:
        state = torch.load(path, map_location=DEVICE)

        # handle different checkpoint formats
        if isinstance(state, dict):
            if "state_dict" in state:
                state = state["state_dict"]
            elif "model_state_dict" in state:
                state = state["model_state_dict"]

        model.load_state_dict(state, strict=False)

        print(f"Loaded {path}")

    except Exception as e:
        print(f"Failed loading {path}: {e}")

    model.to(DEVICE)
    model.eval()
    return model

# =========================
# LOAD MODELS
# =========================
model_121 = load_model("alzheimers_densenet121.pth", "121")
model_169 = load_model("alzheimers_densenet169.pth", "169")
model_201 = load_model("alzheimers_densenet201.pth", "201")

# =========================
# IMAGE PROCESSING
# =========================
def process_image(image_bytes):
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    img = transform(img).unsqueeze(0).to(DEVICE)
    return img

# =========================
# PREDICT FUNCTION (FIXED)
# =========================
def predict(model, x):
    with torch.no_grad():
        out = model(x)
        probs = torch.softmax(out, dim=1)[0]
        conf, cls = torch.max(probs, 0)

        return {
            "class_id": int(cls.item()),
            "class_name": CLASSES[int(cls.item())],
            "confidence": float(conf.item()),
            "probabilities": {
                CLASSES[i]: float(probs[i].item())
                for i in range(len(CLASSES))
            }
        }

# =========================
# ROOT
# =========================
@app.get("/")
def home():
    return {
        "status": "running",
        "models": ["densenet121", "densenet169", "densenet201"],
        "classes": CLASSES,
        "endpoints": ["/predict/121", "/predict/169", "/predict/201", "/predict/ensemble"]
    }

# =========================
# SINGLE MODEL ROUTES
# =========================
@app.post("/predict/121")
async def predict_121(file: UploadFile = File(...)):
    img = process_image(await file.read())
    return {"model": "121", **predict(model_121, img)}

@app.post("/predict/169")
async def predict_169(file: UploadFile = File(...)):
    img = process_image(await file.read())
    return {"model": "169", **predict(model_169, img)}

@app.post("/predict/201")
async def predict_201(file: UploadFile = File(...)):
    img = process_image(await file.read())
    return {"model": "201", **predict(model_201, img)}

# =========================
# ENSEMBLE (FINAL FIXED VERSION)
# =========================
@app.post("/predict/ensemble")
async def ensemble(file: UploadFile = File(...)):

    img = process_image(await file.read())

    r1 = predict(model_121, img)
    r2 = predict(model_169, img)
    r3 = predict(model_201, img)

    avg = {}

    for c in CLASSES:
        avg[c] = (r1["probabilities"][c] +
                  r2["probabilities"][c] +
                  r3["probabilities"][c]) / 3

    final_class = max(avg, key=avg.get)

    return {
        "final_prediction": final_class,
        "final_confidence": avg[final_class],
        "individual_models": {
            "121": r1,
            "169": r2,
            "201": r3
        },
        "ensemble_probabilities": avg
    }