import json
import re
import numpy as np
from pathlib import Path
from graph.smplx_wholebody import SMPLX_KEEP

ROOT = Path("/home/s56fasla/MOTIONX/data")
OUT = Path("/home/s56fasla/MOTIONX/cache")
SUBSETS = ["animation", "haa500", "humman", "idea400", "kungfu", "music", "perform"]
T = 64
MIN_FRAMES = 16
VAL_FRACTION = 0.05
SEED = 0
CLIP_SUFFIX = re.compile(r"_clip\d+$")

OUT.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(SEED)

records = []
for subset in SUBSETS:
    jdir = ROOT / "joint" / subset
    tdir = ROOT / "semantic_label" / subset
    if not jdir.exists():
        continue
    for npy in sorted(jdir.glob("*.npy")):
        txt = tdir / (npy.stem + ".txt")
        if not txt.exists():
            continue
        records.append((subset, npy, txt))
print("paired clips", len(records), flush=True)

groups = {}
for subset, npy, _ in records:
    g = subset + "/" + CLIP_SUFFIX.sub("", npy.stem)
    groups.setdefault(g, []).append(npy.stem)
keys = sorted(groups)
rng.shuffle(keys)
n_val = int(len(keys) * VAL_FRACTION)
val_groups = set(keys[:n_val])
print("groups", len(keys), "val groups", len(val_groups), flush=True)

kept = []
for subset, npy, txt in records:
    g = subset + "/" + CLIP_SUFFIX.sub("", npy.stem)
    kept.append((subset, npy, txt, "val" if g in val_groups else "train"))

arr = np.lib.format.open_memmap(OUT / "pose.npy", mode="w+", dtype=np.float16,
                                shape=(len(kept), T, len(SMPLX_KEEP), 3))
meta = []
idx = np.linspace(0, 1, T)
w = 0
skipped = 0
for i, (subset, npy, txt, split) in enumerate(kept):
    try:
        a = np.load(npy)
    except Exception:
        skipped += 1
        continue
    if a.ndim != 3 or a.shape[0] < MIN_FRAMES or a.shape[1] < 55:
        skipped += 1
        continue
    a = a[:, SMPLX_KEEP, :].astype(np.float32)
    pos = idx * (a.shape[0] - 1)
    lo = np.floor(pos).astype(np.int64)
    hi = np.minimum(lo + 1, a.shape[0] - 1)
    frac = (pos - lo)[:, None, None]
    out = a[lo] * (1.0 - frac) + a[hi] * frac
    out = out - out[0, 0]
    arr[w] = out.astype(np.float16)
    meta.append({"i": w, "subset": subset, "stem": npy.stem, "split": split,
                 "frames": int(a.shape[0]), "text": txt.read_text(encoding="utf-8").strip()})
    w += 1
    if w % 2000 == 0:
        print("processed", w, flush=True)

arr.flush()
del arr
full = np.load(OUT / "pose.npy", mmap_mode="r")
np.save(OUT / "pose_trimmed.npy", np.asarray(full[:w]))
(OUT / "pose.npy").unlink()
(OUT / "pose_trimmed.npy").rename(OUT / "pose.npy")
json.dump(meta, open(OUT / "meta.json", "w"))
print("written", w, "skipped", skipped, flush=True)
print("train", sum(1 for m in meta if m["split"] == "train"), "val", sum(1 for m in meta if m["split"] == "val"), flush=True)