#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.generation.maniskill_lattice import (  # noqa: E402
    ManiSkillLatticeConfig,
    generate_maniskill_lattice,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Generate measured same-state CIL branches from ManiSkill demonstrations."
    )
    parser.add_argument("--demo", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--num-groups", type=int, default=128)
    parser.add_argument(
        "--group-offset",
        type=int,
        default=0,
        help="Start index in the deterministic global episode-step plan.",
    )
    parser.add_argument("--k", type=int, default=8)
    parser.add_argument("--horizon", type=int, default=4)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--shard-size", type=int, default=1024)
    parser.add_argument("--env-id", default="PickCube-v1")
    parser.add_argument("--obs-mode", default="state")
    parser.add_argument(
        "--image-quality",
        type=int,
        default=90,
        help="JPEG quality for observations.h5 when obs-mode includes RGB.",
    )
    parser.add_argument("--control-mode", default="pd_ee_delta_pose")
    parser.add_argument("--sim-backend", default="physx_cuda")
    parser.add_argument(
        "--render-backend",
        default="gpu",
        help="ManiSkill render backend. Use 'gpu' for standard task assets and RGB rendering.",
    )
    parser.add_argument(
        "--parallel-branches",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Execute K same-state interventions in K vectorized ManiSkill environments.",
    )
    parser.add_argument(
        "--state-storage",
        choices=("archive", "files", "none"),
        default="archive",
        help="Persist replay states in one archive, separate files, or only by source reference.",
    )
    parser.add_argument(
        "--state-batch-size",
        type=int,
        default=1,
        help="Number of distinct simulator states to execute together; total envs are G*K.",
    )
    parser.add_argument(
        "--candidate-mode",
        choices=("structured", "random"),
        default="structured",
        help="Use the proposed intervention lattice or a matched random-negative baseline.",
    )
    args = parser.parse_args(argv)
    summary = generate_maniskill_lattice(
        ManiSkillLatticeConfig(
            demo_path=args.demo,
            output_dir=args.out,
            num_groups=args.num_groups,
            group_offset=args.group_offset,
            k=args.k,
            horizon=args.horizon,
            seed=args.seed,
            shard_size=args.shard_size,
            env_id=args.env_id,
            obs_mode=args.obs_mode,
            image_quality=args.image_quality,
            control_mode=args.control_mode,
            sim_backend=args.sim_backend,
            render_backend=args.render_backend,
            parallel_branches=args.parallel_branches,
            state_storage=args.state_storage,
            state_batch_size=args.state_batch_size,
            candidate_mode=args.candidate_mode,
        )
    )
    print(json.dumps(summary, indent=2, default=str))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())