import zipfile
import shutil
from pathlib import Path

RAW = Path("/home/s56fasla/MOTIONX/raw")
OUT = Path("/home/s56fasla/MOTIONX/data")
PREFIX = "comp_robot/zhangyuhong1/data/Motion-X++/v7/"
KINDS = {
    "joint": RAW / "motion/motion_generation/joint",
    "semantic_label": RAW / "text/semantic_label",
    "pose_description": RAW / "text/pose_description",
}

for kind, src in KINDS.items():
    for zp in sorted(src.glob("*.zip")):
        subset = zp.stem
        dst = OUT / kind / subset
        if dst.exists() and any(dst.iterdir()):
            print("skip", kind, subset, flush=True)
            continue
        dst.mkdir(parents=True, exist_ok=True)
        n = 0
        with zipfile.ZipFile(zp) as z:
            for info in z.infolist():
                if info.is_dir():
                    continue
                name = info.filename
                if name.startswith(PREFIX):
                    name = name[len(PREFIX):]
                leaf = Path(name).name
                with z.open(info) as fsrc, open(dst / leaf, "wb") as fdst:
                    shutil.copyfileobj(fsrc, fdst)
                n += 1
        print(kind, subset, n, flush=True)