#!/usr/bin/env python3
"""Run native LeRobot training/resume and offline inference on real action data.

This is a mechanics check, not a robot rollout or task success benchmark.
"""
from __future__ import annotations

import argparse
import importlib.metadata
import json
import math
from pathlib import Path
import re
import subprocess
import sys


def run(command, directory, name):
    with (directory / f"{name}.log").open("w") as log:
        result = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT)
    record = {"command": command, "exit_code": result.returncode}
    (directory / f"{name}_command.json").write_text(json.dumps(record, indent=2) + "\n")
    if result.returncode:
        raise RuntimeError(f"{name} failed; see {directory / (name + '.log')}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--dataset-root", type=Path, required=True)
    parser.add_argument("--dataset-repo-id", required=True)
    parser.add_argument("--mode", choices=["binary", "continuous"], required=True)
    parser.add_argument("--camera-key", default="observation.images.image")
    parser.add_argument("--episodes", type=int, nargs="+", default=[0, 1, 2])
    parser.add_argument("--steps", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--verify-only", action="store_true")
    args = parser.parse_args()
    directory = args.output.resolve()
    directory.mkdir(parents=True, exist_ok=True)
    info = json.loads((args.dataset_root / "meta/info.json").read_text())
    image = info["features"][args.camera_key]
    shape = image["shape"]
    names = image.get("names", [])
    if names == ["height", "width", "channels"]:
        shape = [shape[2], shape[0], shape[1]]
    features = {"observation.images.image": {"type": "VISUAL", "shape": shape}}
    training_dir = directory / "train"
    first_model = training_dir / f"checkpoints/{args.steps:06d}/pretrained_model"
    if not args.verify_only:
        run([
            sys.executable, "-m", "lerobot.scripts.lerobot_train",
            f"--policy.path={args.checkpoint.resolve()}", f"--policy.gripper_mode={args.mode}",
            "--policy.device=cuda", "--policy.dtype=bfloat16", "--policy.push_to_hub=false",
            "--policy.optimizer_lr=0.00001", "--policy.input_features=" + json.dumps(features),
            "--policy.output_features=null", "--rename_map=" + json.dumps({args.camera_key: "observation.images.image"}),
            f"--dataset.repo_id={args.dataset_repo_id}", f"--dataset.root={args.dataset_root.resolve()}",
            "--dataset.video_backend=pyav", "--dataset.episodes=" + json.dumps(args.episodes),
            f"--output_dir={training_dir}", f"--steps={args.steps}", f"--save_freq={args.steps}",
            f"--batch_size={args.batch_size}", "--num_workers=0", "--log_freq=1",
            "--env_eval_freq=0", "--wandb.enable=false", "--seed=123",
        ], directory, "training")
        run([
            sys.executable, "-m", "lerobot.scripts.lerobot_train", "--resume=true",
            f"--config_path={first_model / 'train_config.json'}", f"--steps={args.steps + 1}",
        ], directory, "resume")

    import torch
    from safetensors import safe_open
    import lerobot_policy_dita
    from lerobot.configs.policies import PreTrainedConfig
    from lerobot.datasets.factory import resolve_delta_timestamps
    from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
    from lerobot.policies.factory import make_pre_post_processors
    from lerobot_policy_dita import DitaPolicy
    from lerobot_policy_dita.policy.processor_dita import DitaMaskedNormalizerProcessorStep, DitaMaskedUnnormalizerProcessorStep

    torch.set_num_threads(4)
    model_dir = training_dir / f"checkpoints/{args.steps + 1:06d}/pretrained_model"
    cfg = PreTrainedConfig.from_pretrained(model_dir)
    assert cfg.gripper_mode == args.mode
    pre, post = make_pre_post_processors(cfg, pretrained_path=model_dir)
    encoder = next(s for s in pre.steps if isinstance(s, DitaMaskedNormalizerProcessorStep))
    decoder = next(s for s in post.steps if isinstance(s, DitaMaskedUnnormalizerProcessorStep))
    assert encoder.gripper_mode == decoder.gripper_mode == args.mode
    raw_stats = json.loads((args.dataset_root / "meta/stats.json").read_text())["action"]
    # Episode subsets can produce different statistics. Compare the initial and
    # resumed saved processors; these must agree exactly.
    initial_pre, initial_post = make_pre_post_processors(cfg, pretrained_path=first_model,
        preprocessor_overrides={"device_processor": {"device": "cpu"}})
    initial_encoder = next(s for s in initial_pre.steps if isinstance(s, DitaMaskedNormalizerProcessorStep))
    for key, value in initial_encoder._tensor_stats["action"].items():
        torch.testing.assert_close(encoder._tensor_stats["action"][key].cpu().reshape(-1), value.cpu().reshape(-1), rtol=0, atol=0)
        torch.testing.assert_close(decoder._tensor_stats["action"][key].cpu().reshape(-1), value.cpu().reshape(-1), rtol=0, atol=0)
    del initial_pre, initial_post, initial_encoder
    meta = LeRobotDatasetMetadata(args.dataset_repo_id, root=args.dataset_root)
    dataset = LeRobotDataset(args.dataset_repo_id, root=args.dataset_root, episodes=args.episodes,
        delta_timestamps=resolve_delta_timestamps(cfg, meta), video_backend="pyav")
    indices = [0, len(dataset)//4, len(dataset)//2, 3*len(dataset)//4]
    raw = torch.utils.data.default_collate([dataset[i] for i in indices])
    batch = pre(raw)
    expected = raw["action"][..., -1].to(batch["action"])
    if args.mode == "binary":
        torch.testing.assert_close(batch["action"][..., -1], (1 - expected) / 2)
        reconstructed = batch["action"].clone()
        reconstructed[..., -1] = torch.where(reconstructed[..., -1] > cfg.gripper_threshold, 1., -1.)
    else:
        lo = encoder._tensor_stats["action"]["min"][-1].to(expected)
        hi = encoder._tensor_stats["action"]["max"][-1].to(expected)
        torch.testing.assert_close(batch["action"][..., -1], 2 * (expected - lo) / (hi - lo) - 1)
        reconstructed = batch["action"]
    torch.testing.assert_close(post(reconstructed).cpu(), raw["action"], atol=1e-5, rtol=1e-5)

    policy = DitaPolicy.from_pretrained(model_dir, config=cfg, strict=True)
    policy.eval()
    with torch.no_grad():
        loss, _ = policy(batch)
        assert torch.isfinite(loss)
        torch.manual_seed(7123)
        model_actions = policy.predict_action_chunk(batch)
        actions = post(model_actions)
    assert actions.shape == (4, cfg.chunk_size, 7) and torch.isfinite(actions).all()
    gripper = actions[..., -1]
    if args.mode == "binary":
        assert set(gripper.flatten().tolist()) <= {-1., 1.}
    else:
        assert ((gripper >= lo.cpu() - 1e-5) & (gripper <= hi.cpu() + 1e-5)).all()
        assert ((gripper > lo.cpu() + 1e-4) & (gripper < hi.cpu() - 1e-4)).any(), "All gripper predictions saturated"
    # Also exercise the online single-frame history and select_action path.
    policy.reset()
    online = {k: v for k, v in raw.items() if k != "action" and k != "action_is_pad"}
    for key in list(online):
        if key.startswith("observation.images."):
            online[key] = online[key][:, -1]
    selected = post(policy.select_action(pre(online)))
    assert selected.shape == (4, 7) and torch.isfinite(selected).all()

    losses, gradients = [], []
    for name, count in [("training", args.steps), ("resume", 1)]:
        log = (directory / f"{name}.log").read_text()
        found = re.findall(r"step:(\d+).*?loss:(\S+) grdn:(\S+)", log)
        assert len(found) == count, (name, found)
        losses.extend(float(x[1]) for x in found)
        gradients.extend(float(x[2]) for x in found)
        assert "End of training" in log and "Traceback" not in log
    assert all(math.isfinite(x) for x in losses + gradients)
    changed = 0
    with safe_open(args.checkpoint / "model.safetensors", framework="pt") as before, safe_open(model_dir / "model.safetensors", framework="pt") as after:
        for key in after.keys():
            value = after.get_tensor(key)
            assert torch.isfinite(value).all(), key
            changed += not torch.equal(before.get_tensor(key), value)
    assert changed > 0
    optimizer_steps = set()
    with safe_open(model_dir.parent / "training_state/optimizer_state.safetensors", framework="pt") as state:
        for key in state.keys():
            value = state.get_tensor(key)
            assert torch.isfinite(value).all(), key
            if key.endswith("/step"):
                optimizer_steps.add(value.item())
    assert optimizer_steps == {float(args.steps + 1)}, optimizer_steps
    report = {
        "status": "passed", "scope": "Real-data training, resume and offline inference; no task success or hardware rollout claim",
        "mode": args.mode, "source_checkpoint": str(args.checkpoint.resolve()),
        "dataset_repo_id": args.dataset_repo_id, "dataset_root": str(args.dataset_root.resolve()),
        "episodes": args.episodes, "frames": len(dataset), "sample_indices": indices,
        "training_steps": args.steps, "resume_steps": 1, "batch_size": args.batch_size,
        "losses": losses, "gradient_norms": gradients, "reloaded_loss": loss.item(),
        "changed_weight_tensors": changed, "all_model_and_optimizer_tensors_finite": True,
        "optimizer_steps": sorted(optimizer_steps), "processor_stats_preserved_on_resume": True,
        "real_action_encode_decode_roundtrip": True, "gripper_bounds_source": "action.min/max" if args.mode == "continuous" else "LIBERO signed commands",
        "dataset_gripper_min_max": [raw_stats['min'][-1], raw_stats['max'][-1]],
        "bound_gripper_min_max": [encoder._tensor_stats['action'][k][-1].item() for k in ['min','max']],
        "raw_gripper_targets": sorted(set(raw['action'][..., -1].flatten().tolist())),
        "inference_shape": list(actions.shape), "selected_action_shape": list(selected.shape),
        "inference_gripper_values": gripper.flatten().tolist(), "checkpoint": str(model_dir),
        "runtime_source": lerobot_policy_dita.__file__, "runtime_version": lerobot_policy_dita.__version__,
        "versions": {n: importlib.metadata.version(n) for n in ['lerobot','torch','accelerate','diffusers']},
    }
    (directory / "verification.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2), flush=True)


if __name__ == "__main__":
    main()