app.py

import base64
import io
import math
import os
import time
import uuid

import joblib
import numpy as np
import torch
from fastapi import FastAPI, File, Form, HTTPException, Query, Security, UploadFile
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from huggingface_hub import hf_hub_download
from PIL import Image

from gradcam import generate_heatmap_with_cam
from model_loader import get_model

app = FastAPI(title="Creative Intelligence Scorer")
security = HTTPBearer()

BEARER_TOKEN = os.environ["API_TOKEN"]

ERROR_MESSAGES = {
    "model_not_loaded": "Model is warming up. Please retry in 30 seconds.",
    "invalid_image": "Could not process image. Please upload a JPG or PNG.",
    "inference_failed": "Scoring failed. Please try again.",
}

BENCHMARKS = {
    "gaming":    {"median_ctr": 0.119, "median_halflife": 10.7, "sample_size": 7485},
    "ecommerce": {"median_ctr": 0.125, "median_halflife": 11.2, "sample_size": 7078},
    "finance":   {"median_ctr": 0.111, "median_halflife": 10.0, "sample_size": 6520},
    "other":     {"median_ctr": 0.118, "median_halflife": 10.6, "sample_size": 1000},
}

GRID_LABELS = [
    "top-left", "top-center", "top-right",
    "mid-left", "center",     "mid-right",
    "bot-left", "bot-center", "bot-right",
]

_clf = None


def get_classifier():
    global _clf
    if _clf is None:
        try:
            hf_repo  = os.environ["HF_MODEL_REPO"]
            hf_token = os.getenv("HF_TOKEN")
            clf_path = hf_hub_download(
                repo_id=hf_repo,
                filename="vertical_classifier.pkl",
                token=hf_token,
            )
            _clf = joblib.load(clf_path)
            print("[app] Vertical classifier loaded", flush=True)
        except Exception as e:
            print(f"[app] Classifier load failed (vertical detection disabled): {e}", flush=True)
    return _clf


def verify_token(
    credentials: HTTPAuthorizationCredentials = Security(security),
) -> None:
    if credentials.credentials != BEARER_TOKEN:
        raise HTTPException(status_code=401, detail="Invalid token")


def _open_image(data: bytes) -> Image.Image:
    try:
        return Image.open(io.BytesIO(data)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=422, detail=ERROR_MESSAGES["invalid_image"])


def _cam_region_labels(cam_16: np.ndarray) -> tuple[list[str], list[str]]:
    """3×3 grid aggregation over a 16×16 cam → top-3 high and top-3 low cells."""
    h, w = cam_16.shape
    cells = []
    for row in range(3):
        for col in range(3):
            r0, r1 = row * h // 3, (row + 1) * h // 3
            c0, c1 = col * w // 3, (col + 1) * w // 3
            cells.append(float(cam_16[r0:r1, c0:c1].mean()))

    ranked = sorted(range(9), key=lambda i: cells[i], reverse=True)
    return (
        [GRID_LABELS[i] for i in ranked[:3]],
        [GRID_LABELS[i] for i in ranked[-3:]],
    )


@app.on_event("startup")
async def startup() -> None:
    try:
        get_model()
    except Exception:
        pass  # health endpoint will report model_loaded=False; don't crash the server
    try:
        get_classifier()
    except Exception:
        pass  # vertical detection degrades gracefully; scoring still works


@app.get("/health")
async def health() -> dict:
    try:
        model, _ = get_model()
        loaded = model is not None
    except Exception:
        loaded = False
    return {"status": "ok", "model_loaded": loaded}


@app.get("/benchmark")
async def benchmark(vertical: str = Query(...)) -> dict:
    key = vertical.lower()
    if key not in BENCHMARKS:
        raise HTTPException(
            status_code=422,
            detail=f"Unknown vertical '{vertical}'. Valid: {list(BENCHMARKS)}",
        )
    return {"vertical": key, **BENCHMARKS[key]}


@app.post("/score")
async def score(
    image: UploadFile = File(...),
    vertical: str | None = Form(None),
    _: None = Security(verify_token),
) -> dict:
    try:
        model, processor = get_model()
    except Exception:
        raise HTTPException(status_code=503, detail=ERROR_MESSAGES["model_not_loaded"])

    try:
        pil_image = _open_image(await image.read())
    except HTTPException:
        raise
    except Exception:
        raise HTTPException(status_code=422, detail=ERROR_MESSAGES["invalid_image"])

    try:
        t0 = time.monotonic()

        inputs = processor(images=pil_image, return_tensors="pt")
        pixel_values = inputs["pixel_values"]  # (1, 3, 224, 224), CPU

        with torch.no_grad():
            clip_out = model.backbone(pixel_values=pixel_values)
            embedding = clip_out.pooler_output  # (1, 768)
            outputs = model(embedding=embedding)

        # Vertical classifier — non-fatal, falls back if pkl unavailable
        clf = get_classifier()
        if clf is not None:
            emb_np              = embedding.detach().numpy()   # (1, 768)
            predicted_vertical  = str(clf.predict(emb_np)[0])
            proba               = clf.predict_proba(emb_np)[0]
            print(f"[clf] proba={proba} predicted={predicted_vertical} conf={proba.max()}", flush=True)
            vertical_confidence = round(float(proba.max()), 2)
        else:
            predicted_vertical  = vertical or "other"
            vertical_confidence = 0.0

        ctr_score = float(outputs["ctr_score"].squeeze())

        log_scale = float(outputs["weibull_params"][0, 0].clamp(-10, 10))
        log_shape = float(outputs["weibull_params"][0, 1].clamp(-10, 10))
        scale = math.exp(log_scale)
        shape = math.exp(log_shape)
        # Weibull median = scale * ln(2)^(1/shape) — survival-theoretic halflife
        halflife_days = scale * (math.log(2) ** (1.0 / shape))

        # Confidence proxy: certainty peaks when prediction is near 0 or 1
        confidence = float(1.0 - 2.0 * abs(ctr_score - 0.5))

        inference_ms = int((time.monotonic() - t0) * 1000)

    except Exception:
        raise HTTPException(status_code=500, detail=ERROR_MESSAGES["inference_failed"])

    return {
        "ad_id": str(uuid.uuid4()),
        "ctr_score": round(ctr_score, 4),
        "halflife_days": round(halflife_days, 2),
        "confidence": round(confidence, 4),
        "inference_ms": inference_ms,
        "predicted_vertical": predicted_vertical,
        "vertical_confidence": vertical_confidence,
    }


@app.post("/heatmap")
async def heatmap(
    image: UploadFile = File(...),
    _: None = Security(verify_token),
) -> dict:
    try:
        model, processor = get_model()
    except Exception:
        raise HTTPException(status_code=503, detail=ERROR_MESSAGES["model_not_loaded"])

    try:
        pil_image = _open_image(await image.read())
    except HTTPException:
        raise
    except Exception:
        raise HTTPException(status_code=422, detail=ERROR_MESSAGES["invalid_image"])

    try:
        # Single pass: returns overlay (for base64) + raw cam_16 (for region labels)
        overlay, cam_16 = generate_heatmap_with_cam(model, processor, pil_image, device="cpu")

        high_attention, low_attention = _cam_region_labels(cam_16)

        buf = io.BytesIO()
        Image.fromarray(overlay).save(buf, format="PNG")
        heatmap_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

    except Exception:
        raise HTTPException(status_code=500, detail=ERROR_MESSAGES["inference_failed"])

    return {
        "heatmap_b64": heatmap_b64,
        "high_attention": high_attention,
        "low_attention": low_attention,
    }