# ------------------------------------------------------------
# FastAPI service exposing BinhQuocNguyen/food-recognition-model
# ------------------------------------------------------------
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import base64, io
from PIL import Image
import torch
import numpy as np

# Transformers imports
from transformers import AutoModel, AutoImageProcessor

# -------------------------------------------------------------------
# 1️⃣ Load the model & processor (once, at import time)
# -------------------------------------------------------------------
MODEL_NAME = "BinhQuocNguyen/food-recognition-model"

# AutoModel knows the custom architecture (food_recognition) because
# the repository ships a proper `config.json`.
model = AutoModel.from_pretrained(MODEL_NAME)
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

# Put the model on CPU – the Space has no GPU.
device = torch.device("cpu")
model.to(device)
model.eval()

# Mapping from class index → readable label (comes from the config)
id2label = model.config.id2label  # dict[int, str]

# -------------------------------------------------------------------
# 2️⃣ Minimal nutrient lookup table (extend as you like)
# -------------------------------------------------------------------
nutrient_db = {
    "Apple":  {"calories_per_100g": 52,  "portion_g": 182},
    "Banana": {"calories_per_100g": 89,  "portion_g": 118},
    "Orange": {"calories_per_100g": 43,  "portion_g": 131},
    "Pizza":  {"calories_per_100g": 266, "portion_g": 200},
    "Bread":  {"calories_per_100g": 265, "portion_g": 30},
    # Add the rest of the 101 categories if you need them
}

# -------------------------------------------------------------------
# 3️⃣ Pydantic model for the incoming JSON payload
# -------------------------------------------------------------------
class ImageRequest(BaseModel):
    image: str   # base64‑encoded JPEG/PNG

app = FastAPI()


# ------------------------------------------------------------
# Health‑check endpoint (optional)
# ------------------------------------------------------------
@app.get("/")
def health():
    return {"message": "Food‑Recognition API is up"}


# ------------------------------------------------------------
# 4️⃣ Main inference endpoint
# ------------------------------------------------------------
@app.post("/analyze")
def analyze(request: ImageRequest):
    # ---- 4.1 decode the base64 image ---------------------------------
    try:
        raw = base64.b64decode(request.image)
        pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid base64 image")

    # ---- 4.2 preprocess ------------------------------------------------
    inputs = processor(images=pil_img, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

    # ---- 4.3 forward pass ---------------------------------------------
    with torch.no_grad():
        outputs = model(**inputs)

    # The model returns logits (shape [1, num_classes])
    logits = outputs.logits.squeeze(0)          # [num_classes]
    probs = torch.nn.functional.softmax(logits, dim=-1)

    # ---- 4.4 get top‑1 prediction --------------------------------------
    top_idx = int(probs.argmax().item())
    confidence = float(probs[top_idx].item())
    label = id2label.get(top_idx, "unknown")

    # ---- 4.5 lookup nutrition -----------------------------------------
    nutrition = nutrient_db.get(label, {"calories_per_100g": 0, "portion_g": 100})
    calories_per_100g = nutrition["calories_per_100g"]
    portion_g = nutrition["portion_g"]
    estimated_calories = calories_per_100g * (portion_g / 100.0)

    # ---- 4.6 build JSON response ---------------------------------------
    return {
        "label": label,
        "confidence": confidence,
        "estimated_portion_g": portion_g,
        "calories_per_100g": calories_per_100g,
        "estimated_calories": round(estimated_calories, 2)
    }