from fastapi import FastAPI, Form
from fastapi.middleware.cors import CORSMiddleware
import torch
from torchvision import models, transforms
from PIL import Image
import io
import base64

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

device = torch.device("cpu")
model = models.efficientnet_b0(pretrained=False)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2) 

# Memuat checkpoint dengan ekstraksi state_dict yang aman
checkpoint = torch.load("best_model_EfficientNet-B0.pth", map_location=device, weights_only=False)

if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
    state_dict = checkpoint["model_state_dict"]
else:
    state_dict = checkpoint

model.load_state_dict(state_dict)
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

@app.post("/analyze")
async def analyze_image(imageBase64: str = Form(...), target: str = Form(...)):
    try:
        # Menangani input base64
        img_str = imageBase64.split(",")[1] if "," in imageBase64 else imageBase64
        img_data = base64.b64decode(img_str)
        img = Image.open(io.BytesIO(img_data)).convert("RGB")
        
        # Inferensi model
        tensor = transform(img).unsqueeze(0).to(device)
        with torch.no_grad():
            outputs = model(tensor)
            probs = torch.softmax(outputs, dim=1)[0]
            anemia_prob = probs[1].item() * 100 

        prediction = "anemic" if anemia_prob >= 50 else "non_anemic"
        risk_level = "high" if anemia_prob >= 70 else "moderate" if anemia_prob >= 50 else "optimal"

        return {
            "id": "real-inference",
            "timestamp": "2026-06-30T13:40:00Z",
            "target": target,
            "prediction": prediction,
            "probability": int(anemia_prob),
            "riskLevel": risk_level,
            "xaiFactors": [
                {"label": "Reflektansi mucosal (Visual Model)", "contribution": 65},
                {"label": "Saturasi eritrosit (Red channel)", "contribution": 25},
                {"label": "Distribusi vaskularisasi", "contribution": 10}
            ]
        }
    except Exception as e:
        return {"error": str(e)}

@app.get("/")
def read_root():
    return {"status": "FeMora API is running"}