#!/usr/bin/env python3
"""
Test workspace XY sampling region: reset color envs many times and measure
the empirical X/Y range of object positions.

Output: X range [min, max], extent; Y range [min, max], extent.

Run from ManiSkill repo root with: pip install -e .  (or PYTHONPATH)
"""
import argparse

import gymnasium as gym
import numpy as np

import mani_skill.envs  # noqa: F401


def test_xy_sampling(env_id: str, num_samples: int = 1000):
    env = gym.make(
        env_id,
        obs_mode="state",
        control_mode="pd_joint_pos",
        num_envs=1,
        sim_backend="cpu",
    )
    xs, ys = [], []
    for seed in range(num_samples):
        obs, _ = env.reset(seed=seed)
        # obj position: env.unwrapped.obj.pose.p or from obs
        obj = env.unwrapped.obj
        p = obj.pose.p[0].cpu().numpy()  # (3,) xyz
        xs.append(p[0])
        ys.append(p[1])
    env.close()

    xs = np.array(xs)
    ys = np.array(ys)
    x_min, x_max = xs.min(), xs.max()
    y_min, y_max = ys.min(), ys.max()
    x_extent = x_max - x_min
    y_extent = y_max - y_min

    print(f"Env: {env_id}")
    print(f"Samples: {num_samples}")
    print(f"X: [{x_min:.4f}, {x_max:.4f}], extent (width): {x_extent:.4f}")
    print(f"Y: [{y_min:.4f}, {y_max:.4f}], extent (length): {y_extent:.4f}")
    print(f"XY area: {x_extent:.4f} x {y_extent:.4f}")
    return x_min, x_max, y_min, y_max


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-e", "--env-id",
        default="PushCubeColor-v1",
        help="Env to test (PushCubeColor-v1, PullCubeColor-v1, PickCubeColor-v1, etc.)",
    )
    parser.add_argument(
        "-n", "--num-samples",
        type=int,
        default=1000,
        help="Number of reset samples",
    )
    args = parser.parse_args()

    test_xy_sampling(args.env_id, args.num_samples)


if __name__ == "__main__":
    main()