"""Keypoint layout, so a trimmed keypoint set can be trained without hardcoding 128.

The full DWPose layout is 128 keypoints: body 0:18, face 18:86, left hand 86:107,
right hand 107:128. Dropping keypoints (e.g. knees/ankles, which these clips crop
out) changes the model's input dimension, so the group boundaries can no longer be
constants.

Keypoints are always re-ordered into body / face / lhand / rhand blocks, so group
ranges stay contiguous slices -- only the boundaries move. A layout is persisted as
`layout.json` next to the packed memmaps.
"""
import json
import os

# --- the full DWPose layout, used when no layout.json is present -------------
FULL_GROUPS = {"body": (0, 18), "face": (18, 86), "lhand": (86, 107), "rhand": (107, 128)}
FULL_NKP = 128

# OpenPose-18 body indices, for reference when choosing what to drop:
# 0 nose  1 neck  2 Rshoulder 3 Relbow 4 Rwrist  5 Lshoulder 6 Lelbow 7 Lwrist
# 8 Rhip  9 Rknee 10 Rankle  11 Lhip  12 Lknee  13 Lankle
# 14 Reye 15 Leye 16 Rear    17 Lear
# Knees/ankles are detected in only ~19% / ~0.1% of Full_TriVis frames.
LOWER_BODY = (9, 10, 12, 13)

PRESETS = {
    "full": (),
    "upper": LOWER_BODY,            # drop knees + ankles, keep hips (100% valid)
    "upper_nohip": LOWER_BODY + (8, 11),
}


class Layout:
    """Which original DWPose keypoints are kept, and where the groups sit."""

    def __init__(self, keep, groups, name="full"):
        self.keep = list(keep)          # original indices, in new order
        self.groups = {k: tuple(v) for k, v in groups.items()}
        self.name = name
        self.n_kpts = len(self.keep)
        self.dim = self.n_kpts * 2
        # original index -> new index (None if dropped)
        self.old2new = {o: n for n, o in enumerate(self.keep)}

    # ------------------------------------------------------------------ build
    @classmethod
    def full(cls):
        return cls(list(range(FULL_NKP)), FULL_GROUPS, "full")

    @classmethod
    def from_drop(cls, drop, name="custom"):
        """Build a layout that drops `drop` (original indices), keeping group order."""
        drop = set(int(d) for d in drop)
        keep, groups, cur = [], {}, 0
        for g, (a, b) in FULL_GROUPS.items():
            idx = [i for i in range(a, b) if i not in drop]
            keep.extend(idx)
            groups[g] = (cur, cur + len(idx))
            cur += len(idx)
        return cls(keep, groups, name)

    @classmethod
    def preset(cls, name):
        if name not in PRESETS:
            raise ValueError(f"unknown layout preset {name!r}; choices {list(PRESETS)}")
        return cls.full() if name == "full" else cls.from_drop(PRESETS[name], name)

    # ------------------------------------------------------------------- i/o
    def save(self, data_dir):
        with open(os.path.join(data_dir, "layout.json"), "w") as f:
            json.dump({"name": self.name, "keep": self.keep,
                       "groups": {k: list(v) for k, v in self.groups.items()},
                       "n_kpts": self.n_kpts, "dim": self.dim}, f, indent=2)

    @classmethod
    def load(cls, data_dir):
        p = os.path.join(data_dir, "layout.json")
        if not os.path.exists(p):
            return cls.full()          # datasets packed before layouts existed
        with open(p) as f:
            d = json.load(f)
        return cls(d["keep"], d["groups"], d.get("name", "custom"))

    # ---------------------------------------------------------------- helpers
    def metric_groups(self):
        """Group ranges for reporting, with 'hands' merging both hands."""
        g = dict(self.groups)
        out = {"all": (0, self.n_kpts), "body": g["body"], "face": g["face"],
               "hands": (g["lhand"][0], g["rhand"][1])}
        return out

    def __repr__(self):
        return (f"Layout({self.name}, {self.n_kpts} kpts, dim {self.dim}, "
                f"groups {self.groups})")