"""Portable inference API for the OMI 34-attribute trait-regression model.

Designed to be the single file you import from another project. It:
- Loads one or more DINOv2 backbones (via torch.hub, weights cached on first use).
- Loads one or more MLP head checkpoints per backbone.
- Averages predictions within backbone, then across backbones.
- Returns a (N, 34) numpy array (or pandas DataFrame) on the 0-100 scale.

Bundle layout (see scripts/export_bundle.py):
    bundle/
        manifest.json          {attr_names, backbones: {name: {head_files: [...]}}}
        head_vitl14/            *.pt head checkpoints
        head_vitg14/            *.pt head checkpoints

Minimal usage in another codebase:
    from trait_predictor import TraitPredictor
    predictor = TraitPredictor.from_bundle("path/to/bundle")
    df = predictor.predict(["path/to/face.jpg", ...])   # pandas DataFrame, 34 cols
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable, Sequence

import numpy as np
import torch
from PIL import Image
from torch import nn


# ---------- Backbone + transform ----------

_IMAGENET_MEAN = (0.485, 0.456, 0.406)
_IMAGENET_STD = (0.229, 0.224, 0.225)


def _build_transform(image_size: int = 224):
    from torchvision import transforms  # local import keeps optional deps localized
    return transforms.Compose([
        transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize(_IMAGENET_MEAN, _IMAGENET_STD),
    ])


def _pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


# ---------- MLP head (mirror of src/model.py) ----------

class _TraitHead(nn.Module):
    def __init__(self, in_dim: int, out_dim: int = 34,
                 hidden: int | None = 512, dropout: float = 0.2):
        super().__init__()
        if hidden is None:
            self.net = nn.Linear(in_dim, out_dim)
        else:
            self.net = nn.Sequential(
                nn.Linear(in_dim, hidden), nn.GELU(), nn.Dropout(dropout),
                nn.Linear(hidden, out_dim),
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# ---------- Predictor ----------

class TraitPredictor:
    def __init__(
        self,
        attr_names: list[str],
        backbones: dict[str, dict] | None = None,
        finetunes: list[Path] | None = None,
        device: torch.device | None = None,
    ):
        """
        Parameters
        ----------
        attr_names : list of 34 attribute column names (in the order heads output).
        backbones  : dict {group_name: {"base_model": hub_name, "image_size": int,
                     "head_files": [paths]}}. group_name is a user-defined key
                     (e.g. "vitg_518"); `base_model` is the DINOv2 hub name shared
                     across groups that reuse the same weights.
        finetunes  : list of paths to finetune checkpoints (each carries its own
                     backbone weights + head, loaded as a single module).
        device     : torch device; auto-selected if None.
        """
        self.attr_names = list(attr_names)
        self.device = device or _pick_device()

        self._backbone_groups = backbones or {}
        # Map base_model → lazy-loaded backbone (shared across groups)
        self._base_models: dict[str, tuple[nn.Module, int]] = {}
        # Map group_name → (transform, heads) loaded lazily
        self._group_state: dict[str, tuple[object, list[_TraitHead]]] = {}
        self._finetune_paths = [Path(p) for p in (finetunes or [])]
        self._finetunes: list[tuple[nn.Module, object]] = []

    # ----- constructors -----

    @classmethod
    def from_bundle(cls, bundle_dir: str | Path,
                    device: torch.device | None = None) -> "TraitPredictor":
        bundle_dir = Path(bundle_dir)
        manifest = json.loads((bundle_dir / "manifest.json").read_text())
        attr_names = manifest["attr_names"]
        backbones: dict[str, dict] = {}
        for gname, spec in manifest.get("backbones", {}).items():
            # Back-compat: old manifests used the DINOv2 hub name as the group key
            # and put only head_files in the spec.
            base_model = spec.get("base_model", gname)
            image_size = int(spec.get("image_size", 224))
            backbones[gname] = {
                "base_model": base_model,
                "image_size": image_size,
                "head_files": [bundle_dir / f for f in spec["head_files"]],
            }
        finetunes = [bundle_dir / f for f in manifest.get("finetune_files", [])]
        return cls(attr_names=attr_names, backbones=backbones,
                   finetunes=finetunes, device=device)

    # ----- lazy loading -----

    def _ensure_group_loaded(self, group_name: str):
        if group_name in self._group_state:
            return
        spec = self._backbone_groups[group_name]
        base = spec["base_model"]; image_size = int(spec["image_size"])
        if base not in self._base_models:
            print(f"[TraitPredictor] loading base model {base} on {self.device}...")
            model = torch.hub.load("facebookresearch/dinov2", base)
            model.eval().to(self.device)
            self._base_models[base] = (model, int(model.norm.weight.shape[0]))
        _, in_dim = self._base_models[base]
        transform = _build_transform(image_size=image_size)
        heads: list[_TraitHead] = []
        for p in spec["head_files"]:
            ck = torch.load(p, map_location=self.device, weights_only=False)
            cfg = ck["config"]
            if cfg["in_dim"] != in_dim:
                raise ValueError(
                    f"head {p.name} expects in_dim={cfg['in_dim']} but backbone "
                    f"{base} produces {in_dim}"
                )
            if list(cfg["attr_names"]) != self.attr_names:
                raise ValueError(f"head {p.name} attr_names disagree with manifest")
            head = _TraitHead(in_dim=cfg["in_dim"], out_dim=cfg["out_dim"],
                              hidden=cfg["hidden"], dropout=cfg["dropout"])
            head.load_state_dict(ck["state_dict"])
            head.eval().to(self.device)
            heads.append(head)
        self._group_state[group_name] = (transform, heads)

    # ----- prediction -----

    @torch.inference_mode()
    def _group_forward(self, group_name: str, pils: list[Image.Image]) -> torch.Tensor:
        self._ensure_group_loaded(group_name)
        transform, heads = self._group_state[group_name]
        base = self._backbone_groups[group_name]["base_model"]
        model, _ = self._base_models[base]
        x = torch.stack([transform(p) for p in pils]).to(self.device)
        feats = model(x)
        return torch.stack([h(feats) for h in heads], dim=0).mean(0)

    @torch.inference_mode()
    def _group_forward_flip(self, group_name: str, pils: list[Image.Image]) -> torch.Tensor:
        """TTA: run on horizontal flip."""
        flipped = [p.transpose(Image.FLIP_LEFT_RIGHT) for p in pils]
        return self._group_forward(group_name, flipped)

    def _ensure_finetunes_loaded(self):
        if self._finetunes or not self._finetune_paths:
            return
        for p in self._finetune_paths:
            ck = torch.load(p, map_location=self.device, weights_only=False)
            cfg = ck["config"]
            if list(cfg["attr_names"]) != self.attr_names:
                raise ValueError(f"finetune {p.name} attr_names disagree with manifest")
            print(f"[TraitPredictor] loading finetune {p.name} on {self.device}...")
            backbone = torch.hub.load("facebookresearch/dinov2", cfg["backbone"])
            head = _TraitHead(in_dim=cfg["in_dim"], out_dim=cfg["out_dim"],
                              hidden=cfg["hidden"], dropout=cfg["dropout"])

            class _Wrap(nn.Module):
                def __init__(self, b, h):
                    super().__init__(); self.backbone = b; self.head = h
                def forward(self, x):
                    return self.head(self.backbone(x))

            mod = _Wrap(backbone, head).to(self.device)
            mod.load_state_dict(ck["state_dict"])
            mod.eval()
            image_size = int(cfg.get("image_size", 224))
            self._finetunes.append((mod, _build_transform(image_size=image_size)))

    def predict(
        self,
        images: str | Path | Image.Image | Iterable[str | Path | Image.Image],
        batch_size: int = 16,
        return_dataframe: bool = True,
        tta: bool = True,
    ):
        """Predict the 34-d trait vector for one or many images.

        `images` can be a path, a PIL.Image, or an iterable of those.
        Returns pandas.DataFrame (default) or numpy.ndarray of shape (N, 34),
        with values on the 0-100 scale (clamped).
        """
        if isinstance(images, (str, Path, Image.Image)):
            inputs: Sequence = [images]
            single = True
        else:
            inputs = list(images)
            single = False

        pils: list[Image.Image] = []
        filenames = []
        for item in inputs:
            if isinstance(item, Image.Image):
                img = item.convert("RGB"); name = "<PIL image>"
            else:
                img = Image.open(item).convert("RGB"); name = str(item)
            pils.append(img); filenames.append(name)

        def _chunks(pil_list: list[Image.Image], fwd):
            outs = []
            for s in range(0, len(pil_list), batch_size):
                outs.append(fwd(pil_list[s : s + batch_size]))
            return torch.cat(outs, dim=0)

        group_preds = []
        for gname in self._backbone_groups:
            preds = _chunks(pils, lambda b: self._group_forward(gname, b))
            if tta:
                flipped = _chunks(pils, lambda b: self._group_forward_flip(gname, b))
                preds = (preds + flipped) / 2
            group_preds.append(preds)

        self._ensure_finetunes_loaded()
        for mod, tfm in self._finetunes:
            def _run(pl, tfm=tfm, mod=mod, flip=False):
                if flip:
                    pl = [p.transpose(Image.FLIP_LEFT_RIGHT) for p in pl]
                x = torch.stack([tfm(p) for p in pl]).to(self.device)
                with torch.inference_mode():
                    return mod(x)
            preds = _chunks(pils, _run)
            if tta:
                flipped = _chunks(pils, lambda b: _run(b, flip=True))
                preds = (preds + flipped) / 2
            group_preds.append(preds)

        if not group_preds:
            raise SystemExit("no models to run — bundle is empty")
        y = torch.stack(group_preds, dim=0).mean(0).detach().cpu().numpy()
        y100 = np.clip(y * 100.0, 0.0, 100.0)

        if return_dataframe:
            import pandas as pd
            df = pd.DataFrame(y100, columns=self.attr_names)
            df.insert(0, "filename", filenames)
            return df.iloc[0] if single else df
        return y100[0] if single else y100

    # ----- diagnostic figure for a single image -----

    def predict_with_figure(
        self,
        image: str | Path | Image.Image,
        out_path: str | Path | None = None,
        show: bool = False,
    ):
        """Predict traits for a single image and render the diagnostic panel.

        Returns (pandas.Series of 34 predictions, matplotlib.figure.Figure).
        If `out_path` is given, the figure is saved there.
        """
        import matplotlib.pyplot as plt
        from matplotlib.patches import FancyBboxPatch  # noqa: F401

        row = self.predict(image, return_dataframe=True)
        values = row.drop(labels=["filename"]).to_numpy(dtype=float)
        attr = self.attr_names

        if isinstance(image, Image.Image):
            pil = image.convert("RGB")
        else:
            pil = Image.open(image).convert("RGB")

        fig = plt.figure(figsize=(18, 11))
        gs = fig.add_gridspec(1, 3, width_ratios=[1.0, 1.8, 1.6], wspace=0.35)

        ax_img = fig.add_subplot(gs[0, 0])
        ax_img.imshow(pil); ax_img.axis("off")
        title = str(image) if not isinstance(image, Image.Image) else "<PIL image>"
        ax_img.set_title(title, fontsize=11)

        # --- radar ---
        ax_r = fig.add_subplot(gs[0, 1], projection="polar")
        n = len(attr)
        angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
        vals = np.clip(values, 0, 100)
        vc = np.concatenate([vals, vals[:1]])
        ac = np.concatenate([angles, angles[:1]])
        ax_r.set_theta_offset(np.pi / 2)
        ax_r.set_theta_direction(-1)
        ax_r.set_ylim(0, 100)
        ax_r.set_yticks([50, 100])
        ax_r.set_yticklabels([])
        ax_r.set_xticks(angles); ax_r.set_xticklabels([])
        ax_r.fill(ac, vc, color="#d62728", alpha=0.2, linewidth=0)
        ax_r.plot(ac, vc, "-", color="#d62728", lw=1.5)
        ax_r.plot(angles, vals, "o", color="#d62728", markersize=3)
        for ang, name in zip(angles, attr):
            deg = np.degrees(ang); rot = 90 - deg
            if rot > 90: rot -= 180
            elif rot < -90: rot += 180
            ax_r.text(ang, 112, name, rotation=rot, rotation_mode="anchor",
                      ha="center", va="center", fontsize=7)
        ax_r.set_title("predicted-trait radar (0-100)", fontsize=10, pad=18)

        # --- horizontal bars sorted by value ---
        order = np.argsort(-vals)
        sorted_attrs = [attr[i] for i in order]
        sorted_vals = vals[order]
        ax_b = fig.add_subplot(gs[0, 2])
        y = np.arange(len(attr))
        ax_b.barh(y, sorted_vals, 0.7, color="#d62728", edgecolor="#000", linewidth=0.5)
        for yi, v in enumerate(sorted_vals):
            ax_b.text(v + 1, yi, f"{v:.0f}", va="center", fontsize=7, color="#222")
        ax_b.set_yticks(y); ax_b.set_yticklabels(sorted_attrs, fontsize=7)
        ax_b.invert_yaxis()
        ax_b.set_xlim(0, 110)
        ax_b.set_xticks([0, 25, 50, 75, 100])
        ax_b.set_xlabel("predicted rating (0-100)", fontsize=8)
        ax_b.tick_params(axis="x", labelsize=7)
        ax_b.grid(axis="x", alpha=0.2)
        ax_b.axvline(0, color="#000", lw=0.6)
        ax_b.set_title("predictions, sorted high→low", fontsize=10)

        if out_path is not None:
            from pathlib import Path as _P
            out = _P(out_path)
            out.parent.mkdir(parents=True, exist_ok=True)
            fig.savefig(out, dpi=150, bbox_inches="tight", facecolor="white")
        if show:
            plt.show()
        return row.drop(labels=["filename"]), fig