# eval_sample4d.py — 4维随机采样评估：两个 cube 的 (x,y) 同时随机，找失败组合
import argparse, sys, os
from isaaclab.app import AppLauncher
import cli_args  # isort: skip

parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, default=None)
parser.add_argument("--agent", type=str, default="rsl_rl_cfg_entry_point")
parser.add_argument("--num_samples", type=int, default=4096, help="4维空间随机采样点数 = 并行 env 数")
parser.add_argument("--sample_seed", type=int, default=0)
parser.add_argument("--out", type=str, default="eval_sample4d_out")
cli_args.add_rsl_rl_args(parser)
AppLauncher.add_app_launcher_args(parser)
args_cli, hydra_args = parser.parse_known_args()
sys.argv = [sys.argv[0]] + hydra_args

app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

import gymnasium as gym
import torch
import numpy as np
from rsl_rl.runners import OnPolicyRunner
from isaaclab.envs import ManagerBasedRLEnvCfg
from isaaclab.utils.assets import retrieve_file_path
from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper
import isaaclab_tasks  # noqa: F401
import uwlab_tasks     # noqa: F401
from uwlab_tasks.utils.hydra import hydra_task_config

# 两个 cube 在桌面上的采样范围（撑到训练范围之外，便于逼出失败）
X_MIN, X_MAX = 0.20, 0.65
Y_MIN, Y_MAX = -0.35, 0.45
Z_INS, Z_REC = 0.0065, 0.0070
UPRIGHT = (1.0, 0.0, 0.0, 0.0)


@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg, agent_cfg):
    N = args_cli.num_samples
    env_cfg.scene.num_envs = N
    agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
    agent_cfg = cli_args.sanitize_rsl_rl_cfg(agent_cfg)
    env_cfg.seed = agent_cfg.seed

    env = gym.make(args_cli.task, cfg=env_cfg)
    env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
    u = env.unwrapped
    device = u.device

    resume = retrieve_file_path(args_cli.checkpoint)
    runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
    runner.load(resume)
    policy = runner.get_inference_policy(device=device)
    try:
        policy_nn = runner.alg.policy
    except AttributeError:
        policy_nn = runner.alg.actor_critic

    # ---- 4维随机采样：每个 env 一个 (ins_x, ins_y, rec_x, rec_y) ----
    g = torch.Generator(device=device).manual_seed(args_cli.sample_seed)
    def u_rand(lo, hi):
        return lo + (hi - lo) * torch.rand(N, device=device, generator=g)
    ins_x, ins_y = u_rand(X_MIN, X_MAX), u_rand(Y_MIN, Y_MAX)
    rec_x, rec_y = u_rand(X_MIN, X_MAX), u_rand(Y_MIN, Y_MAX)

    origins = u.scene.env_origins
    ins = u.scene["insertive_object"]
    rec = u.scene["receptive_object"]
    pc = u.reward_manager.get_term_cfg("progress_context").func
    quat = torch.tensor(UPRIGHT, device=device).repeat(N, 1)

    def override():
        for asset, x, y, z in [(ins, ins_x, ins_y, Z_INS), (rec, rec_x, rec_y, Z_REC)]:
            pose = torch.zeros((N, 7), device=device)
            pose[:, 0] = origins[:, 0] + x
            pose[:, 1] = origins[:, 1] + y
            pose[:, 2] = origins[:, 2] + z
            pose[:, 3:7] = quat
            asset.write_root_pose_to_sim(pose)
            asset.write_root_velocity_to_sim(torch.zeros((N, 6), device=device))

    env.get_observations(); env.reset()
    override(); u.sim.forward()
    obs = env.get_observations()

    ever = torch.zeros(N, dtype=torch.bool, device=device)
    frozen = torch.zeros(N, dtype=torch.bool, device=device)
    steps = int(u.max_episode_length)
    print(f"[eval4d] N={N}, steps/episode={steps}")
    for t in range(steps):
        with torch.inference_mode():
            actions = policy(obs)
            obs, _, dones, _ = env.step(actions)
            policy_nn.reset(dones)
        ever |= (pc.success.clone() & ~frozen)
        frozen |= dones.bool()

    # ---- 存 4 列 CSV ----
    os.makedirs(args_cli.out, exist_ok=True)
    data = torch.stack([ins_x, ins_y, rec_x, rec_y, ever.float()], dim=1).cpu().numpy()
    np.save(os.path.join(args_cli.out, "samples.npy"), data)
    with open(os.path.join(args_cli.out, "samples.csv"), "w") as f:
        f.write("ins_x,ins_y,rec_x,rec_y,success\n")
        for r in data:
            f.write(f"{r[0]:.4f},{r[1]:.4f},{r[2]:.4f},{r[3]:.4f},{int(r[4])}\n")

    # ---- 降维可视化：把成败按不同 2D 平面 binning 求成功率 ----
    def binned_rate(ax, px, py, succ, xr, yr, title, xlabel, ylabel, nb=20):
        xb = np.linspace(xr[0], xr[1], nb + 1)
        yb = np.linspace(yr[0], yr[1], nb + 1)
        s, _, _ = np.histogram2d(px, py, bins=[xb, yb], weights=succ)
        c, _, _ = np.histogram2d(px, py, bins=[xb, yb])
        rate = np.divide(s, c, out=np.full_like(s, np.nan), where=c > 0)
        im = ax.imshow(rate.T, origin="lower", aspect="auto", vmin=0, vmax=1,
                       extent=[xr[0], xr[1], yr[0], yr[1]], cmap="RdYlGn")
        ax.set_title(title); ax.set_xlabel(xlabel); ax.set_ylabel(ylabel)
        return im

    try:
        import matplotlib; matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        ix, iy, rx, ry, sc = data[:, 0], data[:, 1], data[:, 2], data[:, 3], data[:, 4]
        fig, axes = plt.subplots(1, 3, figsize=(16, 4.5))
        # 1) 相对位置（最关键）：dx=ins-rec, dy=ins-rec
        binned_rate(axes[0], ix - rx, iy - ry, sc, (-0.6, 0.6), (-0.6, 0.6),
                    "Relative pos (ins - rec)", "dx", "dy")
        # 2) insertive 绝对位置
        binned_rate(axes[1], ix, iy, sc, (X_MIN, X_MAX), (Y_MIN, Y_MAX),
                    "Insertive abs pos", "ins_x", "ins_y")
        # 3) receptive 绝对位置
        im = binned_rate(axes[2], rx, ry, sc, (X_MIN, X_MAX), (Y_MIN, Y_MAX),
                         "Receptive abs pos", "rec_x", "rec_y")
        fig.colorbar(im, ax=axes, label="success rate", fraction=0.025)
        fig.savefig(os.path.join(args_cli.out, "failure_4d_views.png"), dpi=130, bbox_inches="tight")
        print(f"[eval4d] saved -> {args_cli.out}/failure_4d_views.png")
    except Exception as e:
        print(f"[eval4d] plot skipped: {e}")

    print(f"[eval4d] overall success rate = {ever.float().mean().item():.3f}")
    env.close()


if __name__ == "__main__":
    main()
    simulation_app.close()