"""Sim-to-sim: an Isaac Lab-trained Franka Reach policy (rsl_rl, Newton/MJWarp) deployed on strands' MuJoCo
backend, driven by strands `run_policy` and recorded by strands `start_recording` (LeRobot v3 + camera).

Isaac-Reach-Franka contract (params/env.yaml of the run):
  obs (32) = joint_pos_rel(9) | joint_vel_rel(9) | pose_command(7: pos xyz + quat xyzw, base frame) | last_action(7)
  action (7) = panda_joint1..7 targets = default + 0.5 * a      control 30 Hz (dt 1/120, decimation 4)
  success = |p_hand - p_cmd| < 5 cm and orientation error < 0.2 rad (panda_hand body)
The strands side: MuJoCo menagerie `panda` (same menagerie model Isaac Lab's Newton preset uses),
joint names joint1..7 / finger_joint1..2, action keys = joint names.

Run (strands venv): MUJOCO_GL=egl python sim2sim_mujoco.py --jit RUN/exported/policy.pt --episodes 12 --root DS
"""

import argparse
import json
import math
import os
import shutil
import time

import numpy as np
import torch

ap = argparse.ArgumentParser()
ap.add_argument("--jit", required=True)
ap.add_argument("--episodes", type=int, default=12)
ap.add_argument("--seconds", type=float, default=4.0, help="per target (Isaac Lab resamples every 4 s)")
ap.add_argument("--root", required=True)
ap.add_argument("--repo_id", default="cagataydev/strands-isaaclab-reach-franka-mujoco-sim2sim")
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--no_record", action="store_true")
a = ap.parse_args()

from strands_robots.policies.base import Policy  # noqa: E402
from strands_robots.simulation.factory import create_simulation  # noqa: E402

DEFAULT = np.array([0.0, -0.569, 0.0, -2.81, 0.0, 3.037, 0.741, 0.04, 0.04])
JN = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7", "finger_joint1", "finger_joint2"]
ARM = JN[:7]


def quat_xyzw_from_euler(r, p, y):
    cr, sr, cp, sp, cy, sy = (math.cos(r / 2), math.sin(r / 2), math.cos(p / 2), math.sin(p / 2), math.cos(y / 2),
                              math.sin(y / 2))
    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    yy = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy
    q = np.array([x, yy, z, w])
    return q if q[3] >= 0 else -q


def mjid(m, kind, name):
    import mujoco

    for n in (name, "arm/" + name):
        i = mujoco.mj_name2id(m, kind, n)
        if i >= 0:
            return i
    raise KeyError(name)


def quat_err(q1, q2):  # xyzw, angle between
    d = abs(float(np.dot(q1, q2)))
    return 2 * math.acos(min(1.0, d))


class IsaacReachOnStrands(Policy):
    """strands Policy: rebuilds the Isaac Lab Reach observation from a strands observation dict."""

    def __init__(self, jit, sim, seed):
        self.net = torch.jit.load(jit, map_location="cpu").eval()
        self.sim, self.rng = sim, np.random.default_rng(seed)
        self.last = np.zeros(7, np.float32)
        self.cmd = None
        self.log: list[dict] = []
        self.keys = JN

    def set_robot_state_keys(self, k):
        self.keys = list(k)

    @property
    def provider_name(self):
        return "isaaclab_reach_jit"

    def reset(self, seed=None, **kw):
        if seed is not None:
            self.rng = np.random.default_rng(seed)
        pos = np.array([self.rng.uniform(0.35, 0.65), self.rng.uniform(-0.2, 0.2), self.rng.uniform(0.15, 0.5)])
        q = quat_xyzw_from_euler(0.0, math.pi, self.rng.uniform(-3.14, 3.14))
        self.cmd = np.concatenate([pos, q]).astype(np.float32)
        self.last[:] = 0
        # Isaac Lab reset: joints at default * U(0.5,1.5) -> start from the default pose here
        d, m = self.sim.mj_data, self.sim.mj_model
        import mujoco

        for i, n in enumerate(JN):
            jid = mjid(m, mujoco.mjtObj.mjOBJ_JOINT, n)
            d.qpos[m.jnt_qposadr[jid]] = DEFAULT[i]
            d.qvel[m.jnt_dofadr[jid]] = 0
        d.ctrl[:7] = DEFAULT[:7]
        mujoco.mj_forward(m, d)
        self.log.append({"cmd": self.cmd.tolist(), "err": []})

    def hand_pose(self):
        import mujoco

        m, d = self.sim.mj_model, self.sim.mj_data
        bid = mjid(m, mujoco.mjtObj.mjOBJ_BODY, "hand")
        wxyz = d.xquat[bid]
        return d.xpos[bid].copy(), np.array([wxyz[1], wxyz[2], wxyz[3], wxyz[0]])

    async def get_actions(self, observation_dict, instruction, **kw):
        if self.cmd is None:
            self.reset()
        q = np.array([observation_dict[n] for n in JN])
        qd = np.array([observation_dict.get(n + ".vel", 0.0) for n in JN])
        obs = np.concatenate([q - DEFAULT, qd, self.cmd, self.last]).astype(np.float32)
        with torch.inference_mode():
            act = self.net(torch.from_numpy(obs)[None])[0].numpy()
        self.last = act.astype(np.float32)
        tgt = DEFAULT[:7] + 0.5 * act
        p, qq = self.hand_pose()
        self.log[-1]["err"].append([float(np.linalg.norm(p - self.cmd[:3])), quat_err(qq, self.cmd[3:])])
        # menagerie panda: actuator1..7 = position servos on joint1..7, actuator8 = gripper tendon (255 = open)
        out = {f"actuator{i + 1}": float(t) for i, t in enumerate(tgt)}
        out["actuator8"] = 255.0
        return [out]


sim = create_simulation("mujoco")
sim.create_world(timestep=1.0 / 120.0)
sim.add_robot("arm", data_config="panda")
cam = sim.add_camera("front", position=[1.6, 0.9, 0.9], target=[0.4, 0.0, 0.3], width=320, height=240)
pol = IsaacReachOnStrands(a.jit, sim, a.seed)
pol.set_robot_state_keys(JN)
print("action keys", sim.robot_action_keys("arm"))
meta = {"jit": a.jit, "episodes": a.episodes, "seconds": a.seconds, "seed": a.seed, "camera": str(cam)[:200]}
if not a.no_record:
    shutil.rmtree(a.root, ignore_errors=True)
    r = sim.start_recording(repo_id=a.repo_id, root=a.root, fps=30, task="reach the commanded end-effector pose",
                            overwrite=True, cameras=["front"])
    meta["start_recording"] = str(r)[:400]
t0 = time.time()
res = sim.run_policy("arm", policy_object=pol, instruction="reach the commanded end-effector pose",
                     control_frequency=30.0, control_substeps=4, n_steps=int(a.seconds * 30), fast_mode=True,
                     n_episodes=a.episodes, reset_between=True, seed=a.seed)
meta["run_policy"] = str(res)[:1500]
if not a.no_record:
    meta["stop_recording"] = str(sim.stop_recording())[:600]
meta["wall_s"] = round(time.time() - t0, 1)
eps = []
for e in pol.log:
    if not e["err"]:
        continue
    err = np.array(e["err"])
    fin = err[-10:].mean(0)
    eps.append({"final_pos_err_m": round(float(fin[0]), 4), "final_ori_err_rad": round(float(fin[1]), 3),
                "min_pos_err_m": round(float(err[:, 0].min()), 4),
                "success": bool(((err[:, 0] < 0.05) & (err[:, 1] < 0.2)).any()), "steps": len(err)})
meta["episodes_eval"] = eps
meta["success_rate"] = round(float(np.mean([e["success"] for e in eps])), 3) if eps else None
meta["mean_final_pos_err_m"] = round(float(np.mean([e["final_pos_err_m"] for e in eps])), 4) if eps else None
print("[sim2sim] RESULT " + json.dumps(meta, default=str))
os.makedirs(os.path.dirname(os.path.abspath(a.root)), exist_ok=True)
json.dump(meta, open(a.root.rstrip("/") + "_sim2sim.json", "w"), indent=1, default=str)