#!/usr/bin/env python
"""FelineBCS — Gradio inference app.
Upload a cat (or dog) photo -> Body Condition Score 1-9 + group + uncertainty flag.

NON-DIAGNOSTIC. Educational tool built on weak (vision-LLM) labels. Not a substitute
for veterinary examination.
"""
import gradio as gr

try:
    from pillow_heif import register_heif_opener

    register_heif_opener()
except Exception:
    pass

# The model (frozen CLIP backbone + trained heads) is loaded lazily on first
# prediction so this module can be imported without the model artifacts present
# (e.g. in CI import smoke tests). It is warmed at launch in __main__ below.
_MODEL = None


def get_model():
    global _MODEL
    if _MODEL is None:
        from felinebcs_predict import FelineBCS
        _MODEL = FelineBCS()
    return _MODEL


def model_unavailable_message(exc):
    return (
        "## Model artifacts unavailable\n\n"
        f"{exc}\n\n"
        "For deployment, provide the trained head files in `./models/` or set "
        "`FELINEBCS_MODEL_DIR` to the artifact directory. This app is still a "
        "non-diagnostic research tool and is not veterinary advice."
    )


BCS_DESC = {
    1: "Emaciated — ribs/spine/pelvis visible, no fat, severe waist tuck.",
    2: "Very thin — bones easily felt, minimal fat.",
    3: "Thin — ribs easily felt, obvious waist.",
    4: "Lean — ribs palpable, slight waist, minimal fat pad.",
    5: "Ideal — well-proportioned, ribs felt with light fat cover, visible waist.",
    6: "Slightly overweight — ribs felt with difficulty, waist less obvious.",
    7: "Overweight — ribs hard to feel, rounded abdomen, fat pad present.",
    8: "Obese — ribs not palpable, no waist, prominent abdominal fat pad.",
    9: "Severely obese — heavy fat deposits, distended abdomen.",
}

DISCLAIMER = """
### ⚠️ Non-diagnostic educational tool
- **Not veterinary advice.** Body condition scoring by a professional includes *palpation* (feeling ribs/fat), which no photo model can replicate.
- Labels are **weak** — generated by a vision language model, not clinicians. Expect bias toward "ideal" (BCS 5) and errors on extremes.
- Trained on a dataset that **mixes cats and dogs** and contains some **contaminated / off-distribution images** (puppies, non-cats, synthetic renders).
- Best performance on **clear, side-on, full-body** photos of short-haired animals. Fluffy coats, occlusion, odd angles, and close-up faces degrade accuracy.
- For any health concern, **consult a veterinarian.**
"""

def analyze(img):
    if img is None:
        return "Please upload an image.", {}, "", {}
    try:
        r = get_model().predict(img)
        bcs = float(r["bcs"])
        rd = int(r["bcs_rounded"])
        group = str(r["group"])
        group_clf = str(r["group_clf"])
        group_clf_conf = float(r["group_clf_conf"])
        group_agree = bool(r["group_agree"])
        classifier_expected = float(r["classifier_expected"])
        tta_std = float(r["tta_std"])
        disagreement = float(r["disagreement"])
        reject = bool(r["reject"])
        prob9 = {str(k): float(v) for k, v in r["prob9"].items()}

        # headline
        head = f"## Estimated BCS: **{bcs} / 9**  →  {group}\n\n*{BCS_DESC[rd]}*"
        # quality / reject warning
        warn = []
        if reject:
            warn.append(f"🚩 **Low-confidence prediction** (regressor–classifier disagreement "
                        f"{disagreement}, above reject threshold). The two model heads disagree "
                        f"on this image — treat the score as unreliable. Try a clearer, side-on, "
                        f"full-body photo.")
        if tta_std > 0.6:
            warn.append(f"⚠️ Prediction is unstable across image augmentations "
                        f"(TTA std {tta_std}). Interpret with caution.")
        warn_md = "\n\n".join(warn) if warn else "✅ Model heads agree; prediction is within normal confidence."
        # probability distribution over 1-9
        probs = {f"BCS {k}": v for k, v in prob9.items()}
        detail = (f"- **Regressor:** {bcs}\n"
                  f"- **Classifier expected value:** {classifier_expected}\n"
                  f"- **Group (from BCS):** {group}\n"
                  f"- **4-way classifier:** {group_clf} (confidence {group_clf_conf:.0%}, "
                  f"{'agrees' if group_agree else 'disagrees'})\n"
                  f"- **Disagreement:** {disagreement}  |  **TTA std:** {tta_std}")
        api_result = {
            "bcs": bcs,
            "bcs_rounded": rd,
            "group": group,
            "group_classifier": group_clf,
            "group_classifier_confidence": group_clf_conf,
            "group_agree": group_agree,
            "classifier_expected": classifier_expected,
            "tta_std": tta_std,
            "disagreement": disagreement,
            "reject": reject,
            "prob9": prob9,
        }
    except Exception as exc:
        return model_unavailable_message(exc), {}, "", {"error": str(exc)}
    return head + "\n\n" + warn_md, probs, detail, api_result

with gr.Blocks(title="FelineBCS") as demo:
    gr.Markdown("# 🐾 FelineBCS — Cat/Dog Body Condition Score Estimator")
    gr.Markdown("Upload a photo to estimate the 9-point Body Condition Score. "
                "Best with a clear, side-on, full-body shot.")
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="Upload photo")
            btn = gr.Button("Analyze", variant="primary")
        with gr.Column():
            out_head = gr.Markdown()
            out_prob = gr.Label(label="Score distribution (9-way classifier)", num_top_classes=5)
            out_detail = gr.Markdown()
            out_api = gr.JSON(visible=False)
    btn.click(analyze, inputs=inp, outputs=[out_head, out_prob, out_detail, out_api], api_name="analyze")
    gr.Markdown(DISCLAIMER)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)