#!/usr/bin/env python3
"""
Grid-test motion planning for VerbObjectColor-v1 (requires working mplib; use CPU if GPU segfaults).

  python scripts/benchmark_verb_object_mp.py -n 5 --obj-xy 0 0.02 -b cpu --render-backend cpu
  python scripts/benchmark_verb_object_mp.py -v -n 8   # per-seed OK/FAIL lines

Some machines still hit mplib segfaults; run this locally to get real failure tables.
Prints any (verb, shape) with at least one failed trial. Exit 1 if any failure.
"""
from __future__ import annotations

import argparse
import os
import sys

script_dir = os.path.dirname(os.path.abspath(__file__))
root = os.path.dirname(script_dir)
if root not in sys.path:
    sys.path.insert(0, root)

import gymnasium as gym

import mani_skill.envs  # noqa: F401
from mani_skill.examples.motionplanning.panda.run import MP_SOLUTIONS


def main():
    p = argparse.ArgumentParser()
    p.add_argument(
        "--verbs",
        nargs="*",
        default=[
            "lift",
            "grasp",
            "push",
            "pull",
            "rotate",
            "slide",
            "nudge",
            "touch",
        ],
    )
    p.add_argument(
        "--shapes",
        nargs="*",
        default=[
            "cube",
            "sphere",
            "cup",
            "ring",
            "car",
            "pyramid",
            "star",
        ],
    )
    p.add_argument("-n", type=int, default=5, help="seeds 0..n-1 per pair")
    p.add_argument("--obj-xy", type=float, nargs=2, default=[0.0, 0.02])
    p.add_argument("-b", "--sim-backend", default="cpu", choices=["cpu", "gpu"])
    p.add_argument("--render-backend", default="cpu", choices=["cpu", "gpu"])
    p.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Print OK/FAIL per trial (verb, shape, seed, reason).",
    )
    args = p.parse_args()

    solve = MP_SOLUTIONS["VerbObjectColor-v1"]
    bad: list[tuple[str, str, int, int, int]] = []

    for verb in args.verbs:
        for shape in args.shapes:
            fails = exc = 0
            for seed in range(args.n):
                env = gym.make(
                    "VerbObjectColor-v1",
                    obs_mode="none",
                    control_mode="pd_joint_pos",
                    render_mode=None,
                    sim_backend=args.sim_backend,
                    render_backend=args.render_backend,
                    verb=verb,
                    object_shape=shape,
                    object_color_id=0,
                    distractor_max=0,
                    object_size_jiggle=0.0,
                    task_difficulty=1.0,
                )
                try:
                    res = solve(
                        env,
                        seed=seed,
                        debug=False,
                        vis=False,
                        options={"obj_xy": list(args.obj_xy)},
                    )
                except Exception as e:
                    exc += 1
                    fails += 1
                    if args.verbose:
                        print(f"FAIL {verb:7} {shape:8} seed={seed} exception={e!r}")
                else:
                    if res == -1 or res is None:
                        fails += 1
                        if args.verbose:
                            print(
                                f"FAIL {verb:7} {shape:8} seed={seed} reason=plan_or_exec_-1"
                            )
                    elif not bool(res[-1]["success"].item()):
                        fails += 1
                        if args.verbose:
                            print(
                                f"FAIL {verb:7} {shape:8} seed={seed} reason=eval_success_false"
                            )
                    elif args.verbose:
                        print(f"OK   {verb:7} {shape:8} seed={seed}")
                finally:
                    env.close()
            if fails:
                bad.append((verb, shape, fails, args.n, exc))

    print(
        f"VerbObjectColor-v1 MP grid: obj_xy={list(args.obj_xy)}, {args.n} seeds/pair, backends={args.sim_backend}/{args.render_backend}"
    )
    if not bad:
        print("All pairs succeeded.")
        return 0
    print("Failures (verb, shape, fail_count, trials, exceptions):")
    for row in bad:
        print(" ", row)
    return 1


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