#!/usr/bin/env python3
"""Build float32 BPU calibration inputs from recorded deployment logs."""

import argparse
import json
import random
from collections import deque
from pathlib import Path

import numpy as np


NUM_OBS = 45
TERM_DIMS = (3, 3, 3, 12, 12, 12)
HERE = Path(__file__).resolve().parent
DEFAULT_LOG_ROOT = HERE.parents[1] / "logs"


def build_policy_input(history, history_len):
    frames = list(history)
    while len(frames) < history_len:
        frames.insert(0, np.zeros(NUM_OBS, dtype=np.float32))

    chunks = []
    offset = 0
    for dim in TERM_DIMS:
        for frame in frames:
            chunks.append(frame[offset:offset + dim])
        offset += dim
    return np.concatenate(chunks, dtype=np.float32)


def reservoir_add(samples, value, seen, max_samples, rng):
    if len(samples) < max_samples:
        samples.append(value)
        return
    replace_index = rng.randrange(seen)
    if replace_index < max_samples:
        samples[replace_index] = value


def collect_samples(log_paths, max_samples, seed, history_len):
    rng = random.Random(seed)
    samples = []
    seen = 0
    usable_runs = []

    for steps_path in log_paths:
        history = deque(maxlen=history_len)
        run_seen = 0
        was_rl = False

        with steps_path.open("r", encoding="utf-8") as fp:
            for line in fp:
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue

                if record.get("mode") != "RL":
                    if was_rl:
                        history.clear()
                    was_rl = False
                    continue

                was_rl = True
                obs = np.asarray(record.get("obs_single"), dtype=np.float32)
                if obs.shape != (NUM_OBS,) or not np.all(np.isfinite(obs)):
                    continue

                history.append(obs)
                policy_input = build_policy_input(history, history_len)
                if not np.all(np.isfinite(policy_input)):
                    continue

                seen += 1
                run_seen += 1
                reservoir_add(samples, policy_input, seen, max_samples, rng)

        if run_seen:
            usable_runs.append({"steps": str(steps_path), "samples": run_seen})

    return samples, seen, usable_runs


def main():
    parser = argparse.ArgumentParser(
        description="Create BPU float32 calibration .bin files from JSONL deployment logs."
    )
    parser.add_argument(
        "--logs-root",
        type=Path,
        default=DEFAULT_LOG_ROOT,
        help="Directory containing <log-prefix>_*/steps.jsonl.",
    )
    parser.add_argument("--log-prefix", default="robotlab_go1_deploy",
                        help="Run directory prefix below --logs-root")
    parser.add_argument("--history-len", type=int, default=10,
                        help="Number of 45-dim observations to stack by term")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=HERE / "calibration_data",
        help="Output directory for raw float32 feature-map .bin files.",
    )
    parser.add_argument("--max-samples", type=int, default=512)
    parser.add_argument("--min-samples", type=int, default=32)
    parser.add_argument("--seed", type=int, default=20260727)
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    if args.max_samples < 1:
        raise ValueError("--max-samples must be positive")
    if args.history_len < 1:
        raise ValueError("--history-len must be positive")

    log_paths = sorted(args.logs_root.glob(f"{args.log_prefix}_*/steps.jsonl"))
    if not log_paths:
        raise FileNotFoundError(
            f"No step logs found for prefix {args.log_prefix!r} below {args.logs_root}"
        )

    samples, total_seen, usable_runs = collect_samples(
        log_paths, args.max_samples, args.seed, args.history_len
    )
    if len(samples) < args.min_samples:
        raise RuntimeError(
            f"Only {len(samples)} valid RL inputs found; need at least {args.min_samples} calibration samples."
        )

    output_dir = args.output_dir.resolve()
    existing = list(output_dir.glob("*.bin")) if output_dir.exists() else []
    if existing and not args.overwrite:
        raise FileExistsError(
            f"{output_dir} already contains calibration files; pass --overwrite to replace them."
        )
    output_dir.mkdir(parents=True, exist_ok=True)
    for path in existing:
        path.unlink()

    for index, sample in enumerate(samples):
        sample.astype(np.float32, copy=False).tofile(output_dir / f"{index:05d}.bin")

    policy_input_dim = NUM_OBS * args.history_len
    metadata = {
        "format": "raw float32 feature-map",
        "flat_shape": [1, policy_input_dim],
        "mapper_shape": [1, 1, 1, policy_input_dim],
        "history_len": args.history_len,
        "num_obs": NUM_OBS,
        "term_dims": list(TERM_DIMS),
        "log_prefix": args.log_prefix,
        "selected_samples": len(samples),
        "candidate_rl_inputs": total_seen,
        "seed": args.seed,
        "runs": usable_runs,
    }
    (output_dir.parent / f"{output_dir.name}_metadata.json").write_text(
        json.dumps(metadata, indent=2, ensure_ascii=True) + "\n",
        encoding="utf-8",
    )
    print(
        f"Wrote {len(samples)} calibration inputs from {total_seen} RL observations "
        f"to {output_dir}"
    )


if __name__ == "__main__":
    main()
