"""Metadata prediction from SkinMap embeddings via bundled linear probes.

Default uses the global probes, which reproduce the demographic estimates reported
in the paper (e.g. Fitzpatrick V--VI ~ 11%). Passing ``modality`` switches the
Fitzpatrick prediction to a modality-specific probe (clinical vs. dermoscopy/TBP),
which is more accurate when the imaging modality is known.
"""
from __future__ import annotations

import json
import pickle
from pathlib import Path
from typing import Dict, List, Optional, Union

import numpy as np

# imaging modality -> probe group used for the modality-specific Fitzpatrick model
_MODALITY_GROUP = {"clinical": "clinical", "dermoscopy": "derm", "tbp": "derm"}
_FST_GROUP = {1.0: "1-2", 2.0: "1-2", 3.0: "3-4", 4.0: "3-4", 5.0: "5-6", 6.0: "5-6"}


class MetaPredictor:
    """Loads a bundled probe directory and predicts metadata from embeddings."""

    def __init__(self, probes_dir: Union[str, Path]):
        self.dir = Path(probes_dir)
        manifest = self.dir / "manifest.json"
        if not manifest.exists():
            raise FileNotFoundError(f"probe manifest not found at {manifest}")
        self.manifest = json.loads(manifest.read_text())
        self._cache: Dict[str, object] = {}

    def _load(self, name: str):
        if name not in self._cache:
            with open(self.dir / f"{name}.pkl", "rb") as f:
                self._cache[name] = pickle.load(f)
        return self._cache[name]

    @property
    def attributes(self) -> List[str]:
        return list(self.manifest["attributes"].keys())

    def predict(
        self,
        embedding: np.ndarray,
        modality: Optional[str] = None,
        attributes: Optional[List[str]] = None,
    ) -> Dict[str, object]:
        """Predict metadata for one embedding (1024,) or a batch (N, 1024).

        Args:
            embedding: L2-normalized SkinMap embedding(s).
            modality: optional {'clinical','dermoscopy','TBP'}; switches Fitzpatrick
                to the modality-specific probe. Other attributes stay global.
            attributes: subset to predict (default: all available).
        Returns:
            dict attribute -> scalar (single input) or np.ndarray (batch).
        """
        emb = np.asarray(embedding, dtype=np.float32)
        single = emb.ndim == 1
        if single:
            emb = emb[None, :]
        grp = _MODALITY_GROUP.get(str(modality).lower()) if modality else None
        attrs = attributes or self.attributes
        out: Dict[str, object] = {}
        for a in attrs:
            spec = self.manifest["attributes"].get(a)
            if spec is None:
                continue
            if spec["task"] == "regression":
                pred = self._load(f"model_{a}").predict(emb).astype(float)
                clip = spec.get("clip")
                if clip:
                    pred = np.clip(pred, clip[0], clip[1])
                vals = pred
            else:
                if a == "fitzpatrick" and grp and spec.get("modality_specific"):
                    clf = self._load(f"model_fitzpatrick__{grp}")
                    le = self._load("label_encoder_fitzpatrick_modspecific")
                else:
                    clf = self._load(f"model_{a}")
                    le = self._load(f"label_encoder_{a}")
                vals = le.inverse_transform(clf.predict(emb))
                try:
                    vals = vals.astype(float)
                except (ValueError, TypeError):
                    pass
            out[a] = vals

        # derived: grouped Fitzpatrick (I-II / III-IV / V-VI)
        if "fitzpatrick" in out:
            fst = np.asarray(out["fitzpatrick"], dtype=float)
            out["fitzpatrick_grouped"] = np.array(
                [_FST_GROUP.get(round(v), "unknown") for v in np.atleast_1d(fst)]
            )

        if single:
            out = {
                k: (v[0].item() if hasattr(v[0], "item") else v[0])
                for k, v in out.items()
            }
        return out