#!/usr/bin/env python
"""FelineBCS inference: frozen CLIP ViT-B/32 backbone + trained heads.
Predicts Body Condition Score (1-9) with an uncertainty-based reject option.

NON-DIAGNOSTIC: weak-label model trained on vision-LLM scores. Not veterinary advice.
"""
import os

import numpy as np
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
import torch, joblib, open_clip
import torchvision.transforms as T
from PIL import Image

MODEL_DIR_ENV = "FELINEBCS_MODEL_DIR"
DEFAULT_MODEL_DIR = "models"
REQUIRED_MODEL_FILES = (
    "clip_vitb32_regressor.joblib",
    "clip_vitb32_clf9.joblib",
    "clip_vitb32_clf4.joblib",
)
_MEAN = (0.48145466, 0.4578275, 0.40821073)
_STD  = (0.26862954, 0.26130258, 0.27577711)
GRP_NAMES = {0: "Underweight (BCS 1-3)", 1: "Ideal (BCS 4-5)",
             2: "Overweight (BCS 6-7)", 3: "Obese (BCS 8-9)"}
# Reject threshold on regressor-classifier disagreement (calibrated on test set;
# ~70% coverage operating point where MAE drops 0.83->0.74).
DISAGREE_REJECT = 1.5


class MissingModelArtifactsError(RuntimeError):
    """Raised when the trained head artifacts are unavailable."""


def resolve_model_dir(model_dir=None):
    return os.path.expanduser(model_dir or os.environ.get(MODEL_DIR_ENV, DEFAULT_MODEL_DIR))


def validate_model_dir(model_dir):
    missing = [
        os.path.join(model_dir, name)
        for name in REQUIRED_MODEL_FILES
        if not os.path.exists(os.path.join(model_dir, name))
    ]
    if missing:
        missing_list = "\n".join(f"- {path}" for path in missing)
        raise MissingModelArtifactsError(
            "FelineBCS model artifacts were not found.\n"
            f"Set {MODEL_DIR_ENV} to a directory containing the trained heads, "
            f"or place them in ./{DEFAULT_MODEL_DIR}/.\n"
            "Required files:\n"
            f"{missing_list}"
        )

def _bcs_to_grp(b):
    b = int(b)
    if b <= 3: return 0
    if b <= 5: return 1
    if b <= 7: return 2
    return 3

class FelineBCS:
    def __init__(self, model_dir=None, device=None):
        model_dir = resolve_model_dir(model_dir)
        validate_model_dir(model_dir)
        self.model_dir = model_dir
        if device is None:
            device = os.environ.get("FELINEBCS_DEVICE", "cpu")
        self.device = device
        model, _, self.preprocess = open_clip.create_model_and_transforms(
            "ViT-B-32", pretrained="openai")
        self.visual = model.visual.eval().to(device)
        self.reg  = joblib.load(os.path.join(model_dir, "clip_vitb32_regressor.joblib"))
        self.clf9 = joblib.load(os.path.join(model_dir, "clip_vitb32_clf9.joblib"))
        self.clf4 = joblib.load(os.path.join(model_dir, "clip_vitb32_clf4.joblib"))
        self._tta = [self.preprocess,
                     self._mk(True, 1.0), self._mk(False, 1.1),
                     self._mk(True, 1.1), self._mk(False, 1.2)]

    @staticmethod
    def _mk(hflip, scale):
        ops = ([T.RandomHorizontalFlip(1.0)] if hflip else []) + \
              [T.Resize(int(224*scale)), T.CenterCrop(224),
               T.ToTensor(), T.Normalize(_MEAN, _STD)]
        return T.Compose(ops)

    def _feat(self, img, tf):
        with torch.no_grad():
            return self.visual(tf(img).unsqueeze(0).to(self.device)).cpu().numpy()[0]

    def predict(self, img):
        if isinstance(img, str):
            img = Image.open(img)
        img = img.convert("RGB")
        f = self._feat(img, self.preprocess).reshape(1, -1)
        bcs = float(np.clip(self.reg.predict(f)[0], 1, 9))
        p9 = self.clf9.predict_proba(f)[0]
        ev = float((p9 * self.clf9.classes_).sum())
        p4 = self.clf4.predict_proba(f)[0]
        grp_clf = int(self.clf4.classes_[p4.argmax()])
        # Primary group is derived from the production regressor so the band is
        # always coherent with the reported BCS; the 4-way classifier is a
        # secondary opinion (its agreement is an extra confidence signal).
        grp = _bcs_to_grp(round(bcs))
        # TTA spread for a stability read
        tta = np.array([np.clip(self.reg.predict(self._feat(img, tf).reshape(1, -1))[0], 1, 9)
                        for tf in self._tta])
        disagree = abs(bcs - ev)
        return {
            "bcs": round(bcs, 1),
            "bcs_rounded": int(round(bcs)),
            "group": GRP_NAMES[grp],
            "group_clf": GRP_NAMES[grp_clf],
            "group_clf_conf": round(float(p4.max()), 3),
            "group_agree": bool(grp == grp_clf),
            "classifier_expected": round(ev, 1),
            "tta_std": round(float(tta.std()), 2),
            "disagreement": round(disagree, 2),
            "reject": bool(disagree > DISAGREE_REJECT),
            "prob9": {int(c): round(float(p), 3) for c, p in zip(self.clf9.classes_, p9)},
        }

if __name__ == "__main__":
    import sys, json
    m = FelineBCS()
    print(json.dumps(m.predict(sys.argv[1]), indent=2))