# eval_grid.py — 网格评估：扫 insertive cube 的 (x,y)，固定 receptive，统计每格成功率
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("--grid", type=int, default=16, help="每个轴的网格点数 (G×G 格)")
parser.add_argument("--repeats", type=int, default=4, help="每格重复次数 (物理随机性求平均)")
parser.add_argument("--out", type=str, default="eval_grid_out", help="输出目录")
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
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

# ---- 网格范围（来自 reset_states_cfg 的 ObjectAnywhere 范围，留点余量）----
X_MIN, X_MAX = 0.30, 0.55
Y_MIN, Y_MAX = -0.10, 0.30
Z_REST = 0.0065                       # insertive cube 贴桌高度（探针实测）
RECEPTIVE_LOCAL = (0.45, 0.10, 0.0070)  # 固定目标 cube 的局部位置
UPRIGHT = (1.0, 0.0, 0.0, 0.0)        # wxyz，正立


@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg, agent_cfg):
    G, R = args_cli.grid, args_cli.repeats
    num_envs = G * G * R
    env_cfg.scene.num_envs = num_envs
    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

    # ---- 每个 env 对应的网格局部坐标 ----
    xs = torch.linspace(X_MIN, X_MAX, G, device=device)
    ys = torch.linspace(Y_MIN, Y_MAX, G, device=device)
    cell = torch.arange(num_envs, device=device) % (G * G)
    ix, iy = cell % G, cell // G
    gx, gy = xs[ix], ys[iy]                      # [num_envs]

    origins = u.scene.env_origins                # [num_envs,3]
    ins = u.scene["insertive_object"]
    rec = u.scene["receptive_object"]
    pc = u.reward_manager.get_term_cfg("progress_context").func

    def quat_t(q):
        return torch.tensor(q, device=device).repeat(num_envs, 1)

    def override_poses():
        # insertive → 网格位置，正立，贴桌
        ins_pose = torch.zeros((num_envs, 7), device=device)
        ins_pose[:, 0] = origins[:, 0] + gx
        ins_pose[:, 1] = origins[:, 1] + gy
        ins_pose[:, 2] = origins[:, 2] + Z_REST
        ins_pose[:, 3:7] = quat_t(UPRIGHT)
        ins.write_root_pose_to_sim(ins_pose)
        ins.write_root_velocity_to_sim(torch.zeros((num_envs, 6), device=device))
        # receptive → 固定标准位置
        rec_pose = torch.zeros((num_envs, 7), device=device)
        rec_pose[:, 0] = origins[:, 0] + RECEPTIVE_LOCAL[0]
        rec_pose[:, 1] = origins[:, 1] + RECEPTIVE_LOCAL[1]
        rec_pose[:, 2] = origins[:, 2] + RECEPTIVE_LOCAL[2]
        rec_pose[:, 3:7] = quat_t(UPRIGHT)
        rec.write_root_pose_to_sim(rec_pose)
        rec.write_root_velocity_to_sim(torch.zeros((num_envs, 6), device=device))

    # ---- 跑一个 episode ----
    obs = env.get_observations()
    env.reset()
    override_poses()
    u.sim.forward()
    obs = env.get_observations()

    ever_succ = torch.zeros(num_envs, dtype=torch.bool, device=device)
    frozen = torch.zeros(num_envs, dtype=torch.bool, device=device)  # 一旦该 env done 就冻结结果
    steps = int(u.max_episode_length)
    print(f"[eval] num_envs={num_envs} (G={G}, R={R}), steps/episode={steps}")

    for t in range(steps):
        with torch.inference_mode():
            actions = policy(obs)
            obs, _, dones, _ = env.step(actions)
            policy_nn.reset(dones)
        succ = pc.success.clone()
        ever_succ |= (succ & ~frozen)       # 只统计冻结前的成功
        frozen |= dones.bool()              # done 之后该 env 会自动重置、cube 位置会变，停止统计

    # ---- 聚合每格成功率 ----
    succ_f = ever_succ.float().view(R, G * G).mean(dim=0)   # [G*G]
    rate = succ_f.view(G, G).cpu().numpy()                  # [iy, ix]
    xs_np, ys_np = xs.cpu().numpy(), ys.cpu().numpy()

    os.makedirs(args_cli.out, exist_ok=True)
    import numpy as np
    np.save(os.path.join(args_cli.out, "success_rate.npy"), rate)
    with open(os.path.join(args_cli.out, "success_rate.csv"), "w") as f:
        f.write("x,y,success_rate\n")
        for j in range(G):
            for i in range(G):
                f.write(f"{xs_np[i]:.4f},{ys_np[j]:.4f},{rate[j,i]:.4f}\n")

    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        plt.figure(figsize=(6, 5))
        plt.imshow(rate, origin="lower", aspect="auto",
                   extent=[X_MIN, X_MAX, Y_MIN, Y_MAX], vmin=0, vmax=1, cmap="RdYlGn")
        plt.colorbar(label="success rate")
        plt.xlabel("insertive cube x (local)"); plt.ylabel("insertive cube y (local)")
        plt.title("Cube expert failure map")
        plt.scatter([RECEPTIVE_LOCAL[0]], [RECEPTIVE_LOCAL[1]], c="blue", marker="*", s=120, label="receptive")
        plt.legend()
        plt.savefig(os.path.join(args_cli.out, "failure_map.png"), dpi=130, bbox_inches="tight")
        print(f"[eval] saved heatmap -> {args_cli.out}/failure_map.png")
    except Exception as e:
        print(f"[eval] heatmap skipped: {e}")

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


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