#!/usr/bin/env python3
"""Build the dual-rate numeric cache directly from packed val tarballs.

Streams scene JSONL files out of data_scenarios_packed/val/<type>.tar.gz
(no raw extraction — disk is tight) and writes one compressed .npz per scene:

    states   (T, N, 6) float32   [px, py, vx, vy, angle, angvel]
    static   (N, 8)    float32   [is_circle, is_box, log_mass, friction,
                                  elasticity, size0/100, size1/100, 0]
    ids      (N,)      int32     object ids (for obj_{id} text rendering)
    gravity  (2,)      float32
    dt       ()        float32
    header_text        str       pre-rendered scene header (baseline format)

Scenes are split train/val by seed hash (seed % 10 == 0 -> val).
OOD scenario types are never processed here.
"""

import argparse
import json
import math
import sys
import tarfile
from multiprocessing import Pool
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent.parent))

from src.training.data_loader import _format_header_text  # noqa: E402

OOD_TYPES = {"pong", "bowling", "ramp_roll", "angry_birds", "hourglass", "newtons_cradle"}


def scene_to_arrays(lines):
    header = json.loads(lines[0])
    if header.get("type") != "scene_header":
        raise ValueError("first line is not scene_header")
    frames = [json.loads(ln) for ln in lines[1:]]

    ids = [obj["id"] for obj in header["objects"]]
    order = {oid: i for i, oid in enumerate(ids)}
    n = len(ids)
    t = len(frames)

    states = np.zeros((t, n, 6), dtype=np.float32)
    for fi, frame in enumerate(frames):
        for obj in frame["objects"]:
            i = order[obj["id"]]
            states[fi, i, 0] = obj["position"]["x"]
            states[fi, i, 1] = obj["position"]["y"]
            states[fi, i, 2] = obj["velocity"]["x"]
            states[fi, i, 3] = obj["velocity"]["y"]
            states[fi, i, 4] = obj.get("angle", 0.0)
            states[fi, i, 5] = obj.get("angular_velocity", 0.0)

    static = np.zeros((n, 8), dtype=np.float32)
    for obj in header["objects"]:
        i = order[obj["id"]]
        is_circle = 1.0 if obj.get("type") == "circle" else 0.0
        mat = obj.get("material", {})
        if "radius" in obj:
            s0 = s1 = float(obj["radius"])
        else:
            s0 = float(obj.get("width", obj.get("size", 0.0)) or 0.0)
            s1 = float(obj.get("height", obj.get("size", 0.0)) or 0.0)
        static[i] = [
            is_circle,
            1.0 - is_circle,
            math.log(max(mat.get("mass", 1.0), 1e-3)),
            mat.get("friction", 0.0),
            mat.get("elasticity", 0.0),
            s0 / 100.0,
            s1 / 100.0,
            0.0,
        ]

    gravity = np.array([header["gravity"]["x"], header["gravity"]["y"]], dtype=np.float32)
    dt = np.float32(header["timestep"])
    header_text = _format_header_text(header)
    return header, states, static, np.array(ids, dtype=np.int32), gravity, dt, header_text


def process_tarball(args):
    tar_path, out_dir, scenes_per_type = args
    scenario = Path(tar_path).name.replace(".tar.gz", "")
    done = skipped = 0
    with tarfile.open(tar_path, "r:gz") as tf:
        for member in tf:
            if done >= scenes_per_type:
                break
            if not member.name.endswith(".jsonl"):
                continue
            try:
                lines = tf.extractfile(member).read().decode().splitlines()
                header, states, static, ids, gravity, dt, header_text = scene_to_arrays(lines)
                if states.shape[0] < 50:
                    skipped += 1
                    continue
                seed = header["seed"]
                split = "val" if seed % 10 == 0 else "train"
                out_path = Path(out_dir) / split / f"{scenario}_{seed}.npz"
                np.savez_compressed(
                    out_path,
                    states=states,
                    static=static,
                    ids=ids,
                    gravity=gravity,
                    dt=dt,
                    header_text=header_text,
                )
                done += 1
            except Exception as e:  # noqa: BLE001 — skip malformed scenes, count them
                skipped += 1
                if skipped <= 3:
                    print(f"[{scenario}] skip {member.name}: {e}", flush=True)
    print(f"[{scenario}] cached {done}, skipped {skipped}", flush=True)
    return scenario, done, skipped


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--packed-dir", default="data_scenarios_packed/val")
    ap.add_argument("--out", default="data_dualrate/cache")
    ap.add_argument("--scenes-per-type", type=int, default=2500)
    ap.add_argument("--workers", type=int, default=12)
    args = ap.parse_args()

    out_dir = Path(args.out)
    (out_dir / "train").mkdir(parents=True, exist_ok=True)
    (out_dir / "val").mkdir(parents=True, exist_ok=True)

    tarballs = sorted(
        p for p in Path(args.packed_dir).glob("*.tar.gz")
        if p.name.replace(".tar.gz", "") not in OOD_TYPES
    )
    print(f"{len(tarballs)} seen-type tarballs, {args.scenes_per_type}/type -> {out_dir}")

    jobs = [(str(p), str(out_dir), args.scenes_per_type) for p in tarballs]
    with Pool(args.workers) as pool:
        results = pool.map(process_tarball, jobs)

    total = sum(d for _, d, _ in results)
    print(f"TOTAL cached {total} scenes "
          f"(train {len(list((out_dir / 'train').glob('*.npz')))}, "
          f"val {len(list((out_dir / 'val').glob('*.npz')))})")


if __name__ == "__main__":
    main()