#!/usr/bin/env python3
"""
deploy_wtw_pro_sdk.py

Walk-These-Ways (RMA) deployment on Unitree Go1 PRO via go1_pro_sdk.
Direct MCU control — no LCM, no official Unitree SDK.

Model: body_latest.jit (2102 → 12) + adaptation_module_latest.jit (2100 → 2)
Observation: 70-dim, 30-step history (2100 dims total)

Setup:
  cd /path/to/go1_pro_sdk && pip install -e .
  pip install torch onnxruntime numpy

Before running:
  ssh pi@192.168.123.161
  sudo pkill -9 -f keep_sport_alive; sudo pkill -9 -f Legged_sport; sudo pkill -9 -f appTransit

Usage:
  python deploy_wtw_pro_sdk.py --kill-sport --log-dir ../logs
"""

import argparse
import json
import signal
import subprocess
import time
from datetime import datetime
from enum import Enum
from pathlib import Path

import numpy as np
import torch

from go1_pro_sdk import (
    MCUClient, LowCmd, MotorCmd, MotorMode,
    apply_safety, JOINT_NAMES as SDK_JOINT_NAMES,
)

HERE = Path(__file__).parent.resolve()

# ─── WTW policy constants ───
NUM_OBS = 70
NUM_ACTIONS = 12
NUM_COMMANDS = 15
NUM_OBS_HISTORY = 30          # 30 steps of history
OBS_HISTORY_DIM = NUM_OBS * NUM_OBS_HISTORY  # 2100
BODY_INPUT_DIM = 2102          # 2100 (history) + 2 (latent)
ADAPT_INPUT_DIM = 2100
LATENT_DIM = 2

# WTW joint order: FL→FR→RL→RR (per leg: hip/thigh/calf)
# This is different from SDK order: FR→FL→RR→RL
WTW_JOINT_NAMES = [
    "FL_hip", "FL_thigh", "FL_calf",
    "FR_hip", "FR_thigh", "FR_calf",
    "RL_hip", "RL_thigh", "RL_calf",
    "RR_hip", "RR_thigh", "RR_calf",
]

# SDK joint order: FR→FL→RR→RL (per leg: hip/thigh/calf)
# Map: SDK index → WTW index
SDK_TO_WTW = np.array([3, 4, 5,  0, 1, 2,  9, 10, 11,  6, 7, 8], dtype=np.int64)
# Map: WTW index → SDK index
WTW_TO_SDK = np.array([3, 4, 5,  0, 1, 2,  9, 10, 11,  6, 7, 8], dtype=np.int64)

# Default joint angles in WTW order
DEFAULT_ANGLES_WTW = np.array([
     0.1,  0.8, -1.5,   # FL
    -0.1,  0.8, -1.5,   # FR
     0.1,  1.0, -1.5,   # RL
    -0.1,  1.0, -1.5,   # RR
], dtype=np.float32)

# Default joint angles in SDK order
DEFAULT_ANGLES_SDK = DEFAULT_ANGLES_WTW[WTW_TO_SDK]

# Observation scales (WTW standard)
OBS_SCALES = {
    "lin_vel": 2.0, "ang_vel": 0.25,
    "dof_pos": 1.0, "dof_vel": 0.05,
    "body_height_cmd": 2.0, "footswing_height_cmd": 0.15,
    "body_pitch_cmd": 0.3, "body_roll_cmd": 0.3,
    "stance_width_cmd": 1.0, "stance_length_cmd": 1.0,
    "aux_reward_cmd": 1.0,
}

COMMANDS_SCALE = np.array([
    OBS_SCALES["lin_vel"], OBS_SCALES["lin_vel"], OBS_SCALES["ang_vel"],
    OBS_SCALES["body_height_cmd"],  1.0, 1.0, 1.0, 1.0, 1.0,
    OBS_SCALES["footswing_height_cmd"],
    OBS_SCALES["body_pitch_cmd"], OBS_SCALES["body_roll_cmd"],
    OBS_SCALES["stance_width_cmd"], OBS_SCALES["stance_length_cmd"],
    OBS_SCALES["aux_reward_cmd"],
], dtype=np.float32)[:NUM_COMMANDS]

ACTION_SCALE = 0.25
HIP_SCALE_REDUCTION = 0.5  # hip joints get half action
CLIP_ACTIONS = 10.0
CLIP_OBS = 100.0

EXIT = False


def _sig_handler(signum, frame):
    global EXIT
    EXIT = True


signal.signal(signal.SIGINT, _sig_handler)
signal.signal(signal.SIGTERM, _sig_handler)


# ─── State machine ───
class State(Enum):
    IDLE = "IDLE"
    CALIBRATE = "CALIBRATE"
    HOLD = "HOLD"
    RL = "RL"


# ─── Quaternion math ───
def quat_to_rot_matrix(q):
    w, x, y, z = q
    return np.array([
        [1 - 2*y*y - 2*z*z,     2*x*y - 2*w*z,     2*x*z + 2*w*y],
        [    2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z,     2*y*z - 2*w*x],
        [    2*x*z - 2*w*y,     2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y],
    ], dtype=np.float32)


def get_projected_gravity(quaternion):
    R = quat_to_rot_matrix(quaternion)
    return (R.T @ np.array([0., 0., -1.], dtype=np.float32)).astype(np.float32)


# ─── Observation ───
def build_commands_default():
    """Build default command vector for trotting gait."""
    cmd = np.zeros(NUM_COMMANDS, dtype=np.float32)
    cmd[0:3] = [0.0, 0.0, 0.0]    # vx, vy, wz
    cmd[3] = 0.0                     # height command (zero = nominal)
    cmd[4] = 3.0                     # frequency (Hz)
    cmd[5] = 0.5                     # phase (trot = 0.5 offset)
    cmd[6] = 0.0                     # offset
    cmd[7] = 0.0                     # bound
    cmd[8] = 0.5                     # duration (stance ratio)
    cmd[9] = 0.15                    # swing_height (matches reference footswing_height_cmd)
    cmd[10] = 0.0                    # body_pitch
    cmd[11] = 0.0                    # body_roll
    cmd[12] = 0.25                   # stance_width
    cmd[13] = 0.4                    # stance_length
    cmd[14] = 0.0                    # aux_reward
    return cmd


class ClockState:
    """Track gait indices and compute clock_inputs (4-dim sin per foot)."""
    def __init__(self):
        self.gait_indices = 0.0
        self.dt = 0.01  # 100Hz

    def step(self, commands, dt=None):
        if dt is not None:
            self.dt = dt
        freq = commands[4]
        phase = commands[5]
        offset = commands[6]
        bound = commands[7] if NUM_COMMANDS > 8 else 0.0

        self.gait_indices = (self.gait_indices + self.dt * freq) % 1.0

        foot_indices = [
            self.gait_indices + phase + offset + bound,   # FL
            self.gait_indices + offset,                    # FR
            self.gait_indices + bound,                     # RL
            self.gait_indices + phase,                     # RR
        ]
        clock = np.array([np.sin(2 * np.pi * fi) for fi in foot_indices], dtype=np.float32)
        return clock

    def reset(self):
        self.gait_indices = 0.0


def compute_obs_wtw(imu, motor_states, commands, actions, last_actions, clock_inputs):
    """Build 70-dim observation matching WTW LCM agent layout."""
    obs = np.zeros(NUM_OBS, dtype=np.float32)

    # 1. projected_gravity (3)
    obs[0:3] = get_projected_gravity(imu.quaternion)

    # 2. commands * scale (15)
    offset = 3
    obs[offset:offset+NUM_COMMANDS] = commands * COMMANDS_SCALE
    offset += NUM_COMMANDS

    # 3. dof_pos_rel in WTW order (12)
    dof_pos_sdk = np.array([motor_states[i].q for i in range(12)], dtype=np.float32)
    dof_pos_wtw = dof_pos_sdk[SDK_TO_WTW]
    obs[offset:offset+12] = (dof_pos_wtw - DEFAULT_ANGLES_WTW) * OBS_SCALES["dof_pos"]
    offset += 12

    # 4. dof_vel in WTW order (12)
    dof_vel_sdk = np.array([motor_states[i].dq for i in range(12)], dtype=np.float32)
    dof_vel_wtw = dof_vel_sdk[SDK_TO_WTW]
    obs[offset:offset+12] = dof_vel_wtw * OBS_SCALES["dof_vel"]
    offset += 12

    # 5. actions clipped (12)
    obs[offset:offset+12] = np.clip(actions, -CLIP_ACTIONS, CLIP_ACTIONS)
    offset += 12

    # 6. last_actions (12)
    obs[offset:offset+12] = last_actions
    offset += 12

    # 7. clock_inputs (4)
    obs[offset:offset+4] = clock_inputs

    obs = np.clip(obs, -CLIP_OBS, CLIP_OBS)
    obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)
    return obs


# ─── Model ───
class WTWPolicy:
    def __init__(self, body_path, adapt_path):
        self.body = torch.jit.load(str(body_path), map_location='cpu')
        self.adapt = torch.jit.load(str(adapt_path), map_location='cpu')
        self.body.eval()
        self.adapt.eval()

        self.obs_history = torch.zeros(1, OBS_HISTORY_DIM, dtype=torch.float)
        self.latent = torch.zeros(1, LATENT_DIM, dtype=torch.float)

        print(f"[INFO] WTW body: {body_path}")
        print(f"[INFO] WTW adapt: {adapt_path}")
        print(f"[INFO] History: {NUM_OBS} obs × {NUM_OBS_HISTORY} steps = {OBS_HISTORY_DIM}")

    def reset(self):
        self.obs_history.zero_()
        self.latent.zero_()

    def __call__(self, obs):
        obs_t = torch.from_numpy(obs.reshape(1, -1)).float()

        # Update history: shift left, append new obs
        self.obs_history = torch.cat(
            (self.obs_history[:, NUM_OBS:], obs_t), dim=-1)

        # Adaptation module: history → latent
        with torch.no_grad():
            self.latent = self.adapt(self.obs_history)

        # Body: [history, latent] → action
        body_input = torch.cat((self.obs_history, self.latent), dim=-1)
        with torch.no_grad():
            action = self.body(body_input)

        return action.numpy().flatten().astype(np.float32)


# ─── Remote controller ───
def get_rc_commands(state, base_cmd, args):
    r = state.remote
    base_cmd[0] = r.ly * args.rc_vx_scale
    base_cmd[1] = -r.lx * args.rc_vy_scale
    base_cmd[2] = -r.rx * args.rc_wz_scale
    return base_cmd


class RCEdgeDetector:
    def __init__(self):
        self._prev = set()

    def update(self, state):
        current = set(state.remote.pressed)
        rising = current - self._prev
        falling = self._prev - current
        self._prev = current
        return rising, falling


# ─── Safety wrappers ───
def send_hold_cmd(client, state, args):
    cmd = LowCmd()
    for j in range(12):
        cmd.set_motor(j, MotorCmd(
            mode=MotorMode.Servo,
            q=float(DEFAULT_ANGLES_SDK[j]), dq=0.0, tau=0.0,
            Kp=args.kp, Kd=args.kd,
        ))
    apply_safety(cmd, state, power_factor=args.power_factor,
                 position_limit_on=True, position_protect_limit=None)
    client.send(cmd)


def send_rl_cmd(client, state, action_wtw, args):
    """Convert WTW action to SDK targets and send."""
    # Scale action → position offset in WTW order
    offset_wtw = action_wtw * ACTION_SCALE
    # Apply hip scale reduction
    for i in [0, 3, 6, 9]:  # hip indices in WTW order
        offset_wtw[i] *= HIP_SCALE_REDUCTION
    # Target in WTW order
    targets_wtw = DEFAULT_ANGLES_WTW + offset_wtw
    # Convert to SDK order
    targets_sdk = targets_wtw[WTW_TO_SDK]

    cmd = LowCmd()
    for j in range(12):
        cmd.set_motor(j, MotorCmd(
            mode=MotorMode.Servo,
            q=float(targets_sdk[j]), dq=0.0, tau=0.0,
            Kp=args.kp, Kd=args.kd,
        ))
    pp_limit = args.position_protect_limit if args.position_protect_limit > 0 else None
    apply_safety(cmd, state, power_factor=args.power_factor,
                 position_limit_on=True, position_protect_limit=pp_limit)
    client.send(cmd)


def kill_sport_processes(host, user):
    cmds = [
        "sudo pkill -9 -f keep_sport_alive",
        "sudo pkill -9 -f Legged_sport",
        "sudo pkill -9 -f appTransit",
    ]
    ssh_target = f"{user}@{host}"
    print(f"[INFO] Killing sport processes on {ssh_target}...")
    try:
        result = subprocess.run(
            ["ssh", ssh_target, " && ".join(cmds)],
            capture_output=True, text=True, timeout=15)
        if result.returncode == 0 or "no process" in result.stderr.lower():
            print("[INFO] Sport processes killed.")
            return True
        print(f"[WARN] SSH returned {result.returncode}: {result.stderr.strip()}")
        return False
    except Exception as e:
        print(f"[WARN] Failed: {e}")
        return False


# ─── JSONL logger ───
class JsonlLogger:
    def __init__(self, log_dir, args):
        self.enabled = bool(log_dir)
        self.fp = None
        self.run_dir = None
        self.flush_every = 50
        if not self.enabled:
            return
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.run_dir = Path(log_dir).expanduser().resolve() / f"wtw_deploy_{ts}"
        self.run_dir.mkdir(parents=True, exist_ok=True)
        meta = {
            "created_at": ts, "num_obs": NUM_OBS, "num_actions": NUM_ACTIONS,
            "num_commands": NUM_COMMANDS, "num_obs_history": NUM_OBS_HISTORY,
            "action_scale": ACTION_SCALE,
            "default_angles_wtw": DEFAULT_ANGLES_WTW.tolist(),
            "default_angles_sdk": DEFAULT_ANGLES_SDK.tolist(),
            "joint_names_wtw": WTW_JOINT_NAMES,
            "joint_names_sdk": list(SDK_JOINT_NAMES),
        }
        for k, v in vars(args).items():
            if isinstance(v, (str, int, float, bool, type(None))):
                meta[k] = v
        (self.run_dir / "metadata.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False))
        self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8", buffering=1)
        print(f"[INFO] Log dir: {self.run_dir}")

    def log(self, step, **kw):
        if not self.enabled:
            return
        rec = {"step": int(step), "time_wall": time.time()}
        for k, v in kw.items():
            if isinstance(v, np.ndarray):
                rec[k] = np.asarray(v, dtype=np.float32).reshape(-1).tolist()
            elif isinstance(v, (np.float32, np.float64)):
                rec[k] = float(v)
            elif isinstance(v, (np.int32, np.int64)):
                rec[k] = int(v)
            else:
                rec[k] = v
        self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n")
        if step % self.flush_every == 0:
            self.fp.flush()

    def close(self):
        if self.fp:
            self.fp.flush(); self.fp.close()
            print(f"[INFO] Log saved: {self.run_dir}")


# ─── Ramp to default ───
def ramp_to_default(client, args, state):
    print("[INFO] Ramping to default pose (~2s)...")
    current_sdk = np.array([state.motorState[i].q for i in range(12)], dtype=np.float32)
    error = current_sdk - DEFAULT_ANGLES_SDK

    if np.max(np.abs(error)) < 0.05:
        print("[INFO] Already near default pose.")
        return state

    ramp_steps = 200
    step_err = error / ramp_steps

    for i in range(ramp_steps):
        if EXIT: return state
        new_state = client.recv_latest()
        if new_state is not None:
            state = new_state
        targets = DEFAULT_ANGLES_SDK + (error - step_err * min(i + 1, ramp_steps))
        cmd = LowCmd()
        for j in range(12):
            cmd.set_motor(j, MotorCmd(
                mode=MotorMode.Servo,
                q=float(targets[j]), dq=0.0, tau=0.0,
                Kp=args.kp_cal, Kd=args.kd_cal,
            ))
        apply_safety(cmd, state, power_factor=args.power_factor,
                     position_limit_on=True, position_protect_limit=None)
        client.send(cmd)
        time.sleep(0.01)
        if i % 50 == 0:
            actual = np.array([state.motorState[j].q for j in range(12)])
            print(f"  ramp {i}/{ramp_steps} target_err={np.max(np.abs(targets-DEFAULT_ANGLES_SDK)):.3f} actual_err={np.max(np.abs(actual-DEFAULT_ANGLES_SDK)):.3f}")

    for _ in range(50):
        if EXIT: return state
        new_state = client.recv_latest()
        if new_state is not None:
            state = new_state
        send_hold_cmd(client, state, args)
        time.sleep(0.01)

    print("[INFO] Default pose reached.")
    return state


# ─── Main ───
def main():
    parser = argparse.ArgumentParser(description="WTW RMA deployment on Go1 PRO via go1_pro_sdk")
    parser.add_argument("--body", default=str(HERE / "body_latest.jit"))
    parser.add_argument("--adapt", default=str(HERE / "adaptation_module_latest.jit"))

    parser.add_argument("--kill-sport", action="store_true")
    parser.add_argument("--pi-host", default="192.168.123.161")
    parser.add_argument("--pi-user", default="pi")

    parser.add_argument("--kp", type=float, default=20.0)
    parser.add_argument("--kd", type=float, default=0.5)
    parser.add_argument("--kp-cal", type=float, default=15.0)
    parser.add_argument("--kd-cal", type=float, default=0.5)
    parser.add_argument("--power-factor", type=int, default=7)
    parser.add_argument("--position-protect-limit", type=float, default=1.0)

    parser.add_argument("--rc-vx-scale", type=float, default=1.0)
    parser.add_argument("--rc-vy-scale", type=float, default=1.0)
    parser.add_argument("--rc-wz-scale", type=float, default=1.0)
    parser.add_argument("--cmd-x", type=float, default=0.0)
    parser.add_argument("--cmd-y", type=float, default=0.0)
    parser.add_argument("--cmd-yaw", type=float, default=0.0)

    parser.add_argument("--rate-hz", type=float, default=50.0)
    parser.add_argument("--warmup-steps", type=int, default=50)
    parser.add_argument("--max-steps", type=int, default=0)
    parser.add_argument("--print-every", type=int, default=50)

    parser.add_argument("--log-dir", default="")
    parser.add_argument("--log-flush-every", type=int, default=50)

    args = parser.parse_args()

    print("""
╔══════════════════════════════════════════════════════════════╗
║  WTW RMA Deployment (go1_pro_sdk direct MCU)                ║
║  1. Robot SUSPENDED                                         ║
║  2. Use --kill-sport to auto-kill Pi processes               ║
║  3. R2=go/stop, L2=estop, Left-stick=move, Right-stick=turn ║
╚══════════════════════════════════════════════════════════════╝
""")
    input("Press Enter when ready...")

    if args.kill_sport:
        kill_sport_processes(args.pi_host, args.pi_user)

    print("[INFO] Loading WTW models...")
    policy = WTWPolicy(args.body, args.adapt)

    print("[INFO] Connecting to MCU...")
    client = MCUClient()
    logger = JsonlLogger(args.log_dir, args)

    try:
        print("[INFO] Waking MCU...")
        client.wake_mcu(n_frames=50, dt=0.01)

        state = client.recv_state(timeout=2.0)
        if state is None:
            print("[ERROR] No state received.")
            return 1

        print(f"[INFO] Connected. Battery={state.bms.SOC}%")
        print(f"[INFO] RPY: {np.round(np.degrees(state.imu.rpy), 1)} deg")

        sm_state = State.IDLE
        edge = RCEdgeDetector()
        edge.update(state)
        clock = ClockState()
        base_cmd = build_commands_default()
        actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
        last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
        step = 0
        dt = 1.0 / args.rate_hz
        next_t = time.perf_counter()

        print(f"[INFO] Rate: {args.rate_hz}Hz. Ctrl+C to exit.")
        print(f"[INFO] State machine: IDLE → (R2) → CALIBRATE → HOLD → (R2) → RL")

        while not EXIT:
            new_state = client.recv_latest()
            if new_state is not None:
                state = new_state
            if state is None:
                time.sleep(0.001)
                continue

            rising, falling = edge.update(state)
            r2_rose = "R2" in rising
            l2_rose = "L2" in rising

            # Emergency stop
            if l2_rose and sm_state != State.IDLE:
                print(f"\n[L2 EMERGENCY] {sm_state.value} → IDLE")
                sm_state = State.IDLE
                actions[:] = 0.0
                last_actions[:] = 0.0
                policy.reset()
                clock.reset()
                client.send(LowCmd())

            # State machine
            if sm_state == State.IDLE:
                if step % 10 == 0:
                    client.send(LowCmd())

                if r2_rose:
                    print("\n[R2] IDLE → CALIBRATE")
                    sm_state = State.CALIBRATE
                    state = ramp_to_default(client, args, state)
                    if EXIT: break
                    sm_state = State.HOLD
                    print("[STATE] → HOLD")

            elif sm_state == State.HOLD:
                send_hold_cmd(client, state, args)

                if r2_rose:
                    print("\n[R2] HOLD → RL")
                    sm_state = State.RL
                    actions[:] = 0.0
                    last_actions[:] = 0.0
                    policy.reset()
                    clock.reset()

            elif sm_state == State.RL:
                if r2_rose:
                    print("\n[R2] RL → HOLD")
                    sm_state = State.HOLD
                    actions[:] = 0.0
                    last_actions[:] = 0.0
                    state = ramp_to_default(client, args, state)
                    if EXIT: break
                    continue

                # Update commands from RC
                base_cmd = get_rc_commands(state, base_cmd, args)

                # Clock inputs
                clock_inputs = clock.step(base_cmd, dt)

                # Observation
                obs = compute_obs_wtw(state.imu, state.motorState,
                                      base_cmd, actions, last_actions, clock_inputs)

                # Inference
                action_raw = policy(obs)
                actions = np.clip(action_raw, -CLIP_ACTIONS, CLIP_ACTIONS).astype(np.float32)
                last_actions = actions.copy()

                if step >= args.warmup_steps:
                    send_rl_cmd(client, state, action_raw, args)

                logger.log(
                    step, mode="RL",
                    commands=base_cmd,
                    obs_wtw=obs,
                    action_raw=action_raw,
                    action_safe=actions,
                    dof_pos_sdk=np.array([state.motorState[i].q for i in range(12)]),
                    dof_vel_sdk=np.array([state.motorState[i].dq for i in range(12)]),
                    clock_inputs=clock_inputs,
                    imu_rpy_deg=np.degrees(state.imu.rpy),
                    rc_buttons=state.remote.pressed,
                )
            else:
                logger.log(
                    step, mode=sm_state.value,
                    dof_pos_sdk=np.array([state.motorState[i].q for i in range(12)]),
                    imu_rpy_deg=np.degrees(state.imu.rpy),
                    rc_buttons=state.remote.pressed,
                )

            # Status print
            if step % args.print_every == 0:
                dof_pos = np.array([state.motorState[i].q for i in range(12)])
                print(f"\n[STEP {step}] state={sm_state.value} bat={state.bms.SOC}% "
                      f"rpy={np.round(np.degrees(state.imu.rpy), 1)}")
                print(f"  RC: lx={state.remote.lx:+.2f} ly={state.remote.ly:+.2f} "
                      f"btns={state.remote.pressed}")
                print(f"  joint: {np.round(dof_pos, 2)}")
                if sm_state == State.RL:
                    print(f"  action max: {np.max(np.abs(actions)):.2f}")

            step += 1
            if args.max_steps > 0 and step >= args.max_steps:
                print("[INFO] max_steps reached.")
                break

            next_t += dt
            sleep = next_t - time.perf_counter()
            if sleep > 0:
                time.sleep(sleep)
            else:
                next_t = time.perf_counter()

    finally:
        if logger is not None:
            logger.close()
        print("[INFO] Safe stopping...")
        client.safe_stop(n_frames=50, dt=0.002)
        client.close()
        print("[INFO] Done.")


if __name__ == "__main__":
    main()