#!/usr/bin/env python3
"""
Deploy the RoboGauge Go1 45-dim RL-Gym BPU policy on Unitree Go1 PRO.

This uses go1_pro_sdk direct MCU control, not LCM or the official Unitree SDK.

Policy:
  - policy_35k_int16_gemm.bin by default
  - single-frame obs: 45 dims
  - BPU input: 5-frame history, 1x1x1x225 featuremap, stacked by observation terms
  - joint order: FR, FL, RR, RL, matching go1_pro_sdk motor order

Safety-first workflow:
  1. MONITOR:     --monitor, no motor command
  2. OBS-CHECK:   --obs-check, no motor command
  3. INFER-CHECK: --infer-check, BPU only, no motor command
  4. STATE MACHINE:
       IDLE -> CALIBRATE -> HOLD -> OBS_TEST -> INFER_TEST -> RL
       R2 advances one layer, L2 emergency-stops to IDLE.
       RL output is only enabled when --enable-rl is passed.

Before low-level control, kill sport processes on the Pi:
  ssh pi@192.168.123.161
  sudo pkill -9 -f keep_sport_alive
  sudo pkill -9 -f Legged_sport
  sudo pkill -9 -f appTransit

Initial tests should be done with the robot suspended. For extra-conservative
low-level checks, override the default with --power-factor 1.
"""

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

import numpy as np

from go1_pro_sdk import (
    MotorMode, PowerProtectViolation, JOINT_NAMES,
)

HERE = Path(__file__).parent.resolve()
DEPLOY_ROOT = HERE.parent
WORKSPACE_ROOT = HERE.parents[1]
sys.path.insert(0, str(HERE / "bpu_deploy_x5"))
from bpu_policy import BpuInferLibPolicy


SDK_FAST_LOW_CMD = WORKSPACE_ROOT / "go1_pro_sdk" / "fast_lowcmd_cpp"
if SDK_FAST_LOW_CMD.exists():
    sys.path.insert(0, str(SDK_FAST_LOW_CMD))
try:
    from fast_lowcmd import FastLowCmdBuilder
    from fast_mcu import FastMCUClient
except ImportError as exc:
    raise RuntimeError(
        "C++ fast LowCmd/LowState backend is required for this entrypoint. "
        "Build it first: cd /root/go1_pro_sdk/fast_lowcmd_cpp && "
        "PYTHONPATH=/root/go1_pro_sdk python3 setup.py build_ext --inplace"
    ) from exc

BPU_MODEL_REGISTRY = {
    "15k": HERE / "bpu_quantization" / "mapper_output_15k_gemm" / "policy_15k_int16_gemm.bin",
    "25k": HERE / "bpu_quantization" / "mapper_output_25k_gemm" / "policy_25k_int16_gemm.bin",
    "30k": HERE / "bpu_quantization" / "mapper_output_30k_gemm" / "policy_30k_int16_gemm.bin",
    "35k": HERE / "bpu_quantization" / "mapper_output_35k_gemm" / "policy_35k_int16_gemm.bin",
}
DEFAULT_BPU_ROUND = "35k"
DEFAULT_BPU_MODEL = BPU_MODEL_REGISTRY[DEFAULT_BPU_ROUND]
LOWCMD_BACKEND = "cpp_lowcmd_cpp_lowstate"
SPORT_KILL_CMD = (
    'ssh pi@192.168.123.161 "sudo pkill -9 -f keep_sport_alive; '
    'sudo pkill -9 -f Legged_sport; sudo pkill -9 -f appTransit"'
)

NUM_OBS = 45
NUM_ACTIONS = 12
HISTORY_LEN = 5
POLICY_INPUT_DIM = NUM_OBS * HISTORY_LEN
BPU_INPUT_SHAPE = [1, 1, 1, POLICY_INPUT_DIM]
BPU_OUTPUT_SHAPE = [1, NUM_ACTIONS, 1, 1]

ACTION_SCALE = 0.25
CLIP_OBS = 100.0
ANG_VEL_SCALE = 0.25
DOF_VEL_SCALE = 0.05
CMD_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32)

MAX_LIN_VEL_X = 1.0
MAX_LIN_VEL_Y = 0.5
MAX_ANG_VEL_YAW = 1.0

DEFAULT_DOF_POS = np.array([
    -0.1, 0.8, -1.5,   # FR_hip, FR_thigh, FR_calf
     0.1, 0.8, -1.5,   # FL_hip, FL_thigh, FL_calf
    -0.1, 1.0, -1.5,   # RR_hip, RR_thigh, RR_calf
     0.1, 1.0, -1.5,   # RL_hip, RL_thigh, RL_calf
], dtype=np.float32)
EXPECTED_SDK_JOINT_NAMES = [
    "FR_0", "FR_1", "FR_2",
    "FL_0", "FL_1", "FL_2",
    "RR_0", "RR_1", "RR_2",
    "RL_0", "RL_1", "RL_2",
]
_FAST_LOW_CMD_BUILDER = None
POLICY_JOINT_NAMES = [
    "FR_hip", "FR_thigh", "FR_calf",
    "FL_hip", "FL_thigh", "FL_calf",
    "RR_hip", "RR_thigh", "RR_calf",
    "RL_hip", "RL_thigh", "RL_calf",
]

EXIT = False


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


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


class State(Enum):
    IDLE = "IDLE"
    CALIBRATE = "CALIBRATE"
    HOLD = "HOLD"
    OBS_TEST = "OBS_TEST"
    INFER_TEST = "INFER_TEST"
    RL = "RL"
    FAULT = "FAULT"


def get_projected_gravity(quat_wxyz):
    qw, qx, qy, qz = quat_wxyz
    g = np.zeros(3, dtype=np.float32)
    g[0] = 2.0 * (-qz * qx + qw * qy)
    g[1] = -2.0 * (qz * qy + qw * qx)
    g[2] = 1.0 - 2.0 * (qw * qw + qz * qz)
    return g


def as_np(values, dtype=np.float32):
    return np.asarray(values, dtype=dtype)


def motor_pos(state):
    return np.array([state.motorState[i].q for i in range(NUM_ACTIONS)], dtype=np.float32)


def motor_vel(state):
    return np.array([state.motorState[i].dq for i in range(NUM_ACTIONS)], dtype=np.float32)


def motor_tau(state):
    return np.array([state.motorState[i].tauEst for i in range(NUM_ACTIONS)], dtype=np.float32)


def motor_mode(state):
    return np.array([state.motorState[i].mode for i in range(NUM_ACTIONS)], dtype=np.int32)


def motor_temperature(state):
    return np.array([state.motorState[i].temperature for i in range(NUM_ACTIONS)], dtype=np.int32)


def motor_reserve(state):
    return np.array([state.motorState[i].reserve for i in range(NUM_ACTIONS)], dtype=np.int32)


def motor_servo_fault(state):
    modes = motor_mode(state)
    bad = [
        f"{JOINT_NAMES[i]}={int(modes[i])}"
        for i in range(NUM_ACTIONS)
        if int(modes[i]) != int(MotorMode.Servo)
    ]
    if bad:
        return "motor feedback not servo: " + ", ".join(bad)
    return None


class MotorServoGuard:
    def __init__(self, fault_frames):
        self.fault_frames = max(1, int(fault_frames))
        self.consecutive_bad = 0

    def reset(self):
        self.consecutive_bad = 0

    def update(self, state):
        reason = motor_servo_fault(state)
        if reason is None:
            self.reset()
            return None
        self.consecutive_bad += 1
        if self.consecutive_bad >= self.fault_frames:
            return f"{reason} ({self.consecutive_bad} consecutive frames)"
        return None


def validate_joint_order():
    sdk_names = list(JOINT_NAMES)
    if sdk_names != EXPECTED_SDK_JOINT_NAMES:
        raise RuntimeError(
            "go1_pro_sdk JOINT_NAMES order mismatch.\n"
            f"  expected: {EXPECTED_SDK_JOINT_NAMES}\n"
            f"  actual  : {sdk_names}"
        )
    print("[INFO] Joint order check passed: SDK and policy both use FR, FL, RR, RL.")
    for i, (sdk_name, policy_name, q0) in enumerate(
            zip(sdk_names, POLICY_JOINT_NAMES, DEFAULT_DOF_POS)):
        print(f"  [{i:02d}] {sdk_name:4s} -> {policy_name:8s} default={q0:+.3f}")


def resolve_bpu_model(args):
    if args.bpu_model:
        return Path(args.bpu_model).expanduser().resolve()
    key = str(args.bpu_round).lower()
    if key not in BPU_MODEL_REGISTRY:
        choices = ", ".join(sorted(BPU_MODEL_REGISTRY))
        raise ValueError(f"Unknown --bpu-round {args.bpu_round!r}; choices: {choices}")
    return BPU_MODEL_REGISTRY[key].expanduser().resolve()


def make_policy(args):
    model_path = resolve_bpu_model(args)
    policy = BpuInferLibPolicy(
        model_path,
        priority=args.bpu_priority,
        bpu_cores=args.bpu_cores,
    )
    if policy.input_size != POLICY_INPUT_DIM:
        policy.close()
        raise ValueError(
            f"BPU model input has {policy.input_size} values, expected {POLICY_INPUT_DIM}. "
            "Use a Gym 5-frame/225-dim .bin, not a RobotLab 10-frame/450-dim .bin."
        )
    if policy.output_size != NUM_ACTIONS:
        policy.close()
        raise ValueError(
            f"BPU model output has {policy.output_size} values, expected exactly {NUM_ACTIONS}. "
            "Compile an actions-only ONNX for deployment."
        )
    return policy


def apply_deadzone(value, deadzone):
    if deadzone <= 0.0:
        return float(value)
    mag = abs(float(value))
    if mag <= deadzone:
        return 0.0
    return float(np.sign(value) * (mag - deadzone) / max(1e-6, 1.0 - deadzone))


def get_command(state, args):
    if args.no_rc:
        cmd = np.array([args.cmd_x, args.cmd_y, args.cmd_yaw], dtype=np.float32)
    else:
        r = state.remote
        ly = apply_deadzone(r.ly, args.rc_deadzone)
        lx = apply_deadzone(r.lx, args.rc_deadzone)
        rx = apply_deadzone(r.rx, args.rc_deadzone)
        if args.swap_vy_yaw:
            cmd = np.array([
                ly * args.rc_vx_scale,
                -rx * args.rc_vy_scale,
                -lx * args.rc_wz_scale,
            ], dtype=np.float32)
        else:
            cmd = np.array([
                ly * args.rc_vx_scale,
                -lx * args.rc_vy_scale,
                -rx * args.rc_wz_scale,
            ], dtype=np.float32)

    limits = np.array([MAX_LIN_VEL_X, MAX_LIN_VEL_Y, MAX_ANG_VEL_YAW], dtype=np.float32)
    return np.clip(cmd, -limits, limits)


class CommandFilter:
    def __init__(self, args):
        self.alpha = float(args.cmd_ema_alpha)
        self.max_step = np.array([
            args.max_cmd_step_x,
            args.max_cmd_step_y,
            args.max_cmd_step_yaw,
        ], dtype=np.float32)
        self.prev = np.zeros(3, dtype=np.float32)

    def reset(self):
        self.prev[:] = 0.0

    def update(self, raw_cmd):
        cmd = np.asarray(raw_cmd, dtype=np.float32)
        if 0.0 < self.alpha < 1.0:
            cmd = self.alpha * cmd + (1.0 - self.alpha) * self.prev
        if np.any(self.max_step > 0.0):
            limit = np.where(self.max_step > 0.0, self.max_step, np.inf)
            cmd = self.prev + np.clip(cmd - self.prev, -limit, limit)
        self.prev = cmd.astype(np.float32)
        return self.prev.copy()


class ObsHistoryBuilder:
    """Build RoboGauge 45-dim obs and 225-dim term-stacked policy history."""

    def __init__(self):
        self.history = deque(maxlen=HISTORY_LEN)

    def reset(self):
        self.history.clear()

    def build_single(self, state, cmd, last_action):
        quat = as_np(state.imu.quaternion)
        gyro = as_np(state.imu.gyroscope)
        q = motor_pos(state)
        dq = motor_vel(state)

        obs = np.zeros(NUM_OBS, dtype=np.float32)
        obs[0:3] = gyro * ANG_VEL_SCALE
        obs[3:6] = get_projected_gravity(quat)
        obs[6:9] = cmd * CMD_SCALE
        obs[9:21] = q - DEFAULT_DOF_POS
        obs[21:33] = dq * DOF_VEL_SCALE
        obs[33:45] = last_action
        obs = np.clip(obs, -CLIP_OBS, CLIP_OBS)
        return np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)

    def build_policy_input(self, obs_single):
        self.history.append(obs_single.copy())
        frames = list(self.history)
        while len(frames) < HISTORY_LEN:
            frames.insert(0, np.zeros(NUM_OBS, dtype=np.float32))

        term_dims = [3, 3, 3, 12, 12, 12]
        chunks = []
        offset = 0
        for dim in term_dims:
            for frame in frames:
                chunks.append(frame[offset:offset + dim])
            offset += dim
        obs = np.concatenate(chunks, dtype=np.float32).reshape(1, POLICY_INPUT_DIM)
        return np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)


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


class JsonlLogger:
    def __init__(self, log_dir, args):
        self.enabled = bool(log_dir)
        self.fp = None
        self.run_dir = None
        self.flush_every = max(1, int(args.log_flush_every))
        if not self.enabled:
            return

        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.run_dir = Path(log_dir).expanduser().resolve() / f"rlgym_go1_deploy_{ts}"
        self.run_dir.mkdir(parents=True, exist_ok=True)
        meta = {
            "created_at": ts,
            "num_obs": NUM_OBS,
            "history_len": HISTORY_LEN,
            "policy_input_dim": POLICY_INPUT_DIM,
            "bpu_input_shape": BPU_INPUT_SHAPE,
            "bpu_output_shape": BPU_OUTPUT_SHAPE,
            "action_scale": ACTION_SCALE,
            "default_dof_pos": DEFAULT_DOF_POS.tolist(),
            "joint_names_sdk": list(JOINT_NAMES),
            "joint_names_sdk_expected": EXPECTED_SDK_JOINT_NAMES,
            "joint_names_policy": POLICY_JOINT_NAMES,
            "joint_order_policy": ["FR", "FL", "RR", "RL"],
            "policy_backend": BpuInferLibPolicy.backend_name,
            "bpu_model": str(resolve_bpu_model(args)),
            "lowcmd_backend": LOWCMD_BACKEND,
        }
        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):
                if np.issubdtype(v.dtype, np.integer):
                    rec[k] = np.asarray(v, dtype=np.int32).reshape(-1).tolist()
                else:
                    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}")


def fmt_rc(state):
    r = state.remote
    btns = ",".join(r.pressed) if r.pressed else "none"
    return (
        f"lx={r.lx:+.2f} ly={r.ly:+.2f} rx={r.rx:+.2f} ry={r.ry:+.2f} "
        f"L2={r.L2:.2f} btns={btns}"
    )


def get_fast_lowcmd_builder():
    global _FAST_LOW_CMD_BUILDER
    if _FAST_LOW_CMD_BUILDER is None:
        _FAST_LOW_CMD_BUILDER = FastLowCmdBuilder()
    return _FAST_LOW_CMD_BUILDER


def send_damping(client):
    raw = get_fast_lowcmd_builder().build_encrypted_damping()
    client.send_raw(raw)


def build_servo12_raw(targets, kp, kd, state=None, position_protect_limit=None):
    actual_q = None if state is None else motor_pos(state)
    actual_tau = None if state is None else motor_tau(state)
    pp_limit = 0.0 if position_protect_limit is None else float(position_protect_limit)
    return get_fast_lowcmd_builder().build_encrypted_servo12_checked(
        np.asarray(targets, dtype=np.float32),
        float(kp),
        float(kd),
        actual_q=actual_q,
        actual_tau=actual_tau,
        position_protect_limit=pp_limit,
    )


def send_servo12_fast(client, targets, kp, kd, state=None, position_protect_limit=None):
    raw = build_servo12_raw(targets, kp, kd, state, position_protect_limit)
    client.send_raw(raw)


def send_hold_cmd(client, state, args):
    send_servo12_fast(client, DEFAULT_DOF_POS, args.kp, args.kd, state=state)


def send_position_cmd(client, state, targets, args):
    pp_limit = args.position_protect_limit if args.position_protect_limit > 0 else None
    send_servo12_fast(client, targets, args.kp, args.kd, state=state, position_protect_limit=pp_limit)


def ramp_to_default(client, args, state, logger=None, step_base=0):
    print("[INFO] Ramping to default pose...")
    current = motor_pos(state)
    error = current - DEFAULT_DOF_POS
    max_error = float(np.max(np.abs(error)))
    print(f"[INFO] Current max default-pose error: {max_error:.3f} rad")
    if max_error < 0.05:
        print("[INFO] Already near default pose.")
        return state

    ramp_steps = max(1, int(args.ramp_time * args.ramp_hz))
    dt = 1.0 / args.ramp_hz
    log_every = max(1, int(args.ramp_hz / 5.0))
    next_t = time.perf_counter()

    for i in range(ramp_steps):
        if EXIT:
            return state
        new_state = client.recv_latest()
        if new_state is not None:
            state = new_state

        ratio = float(i + 1) / float(ramp_steps)
        target = current + ratio * (DEFAULT_DOF_POS - current)
        send_servo12_fast(client, target, args.kp_cal, args.kd_cal, state=state)
        if logger is not None and (i % log_every == 0 or i == ramp_steps - 1):
            log_state(logger, step_base * 100000 + i, "RAMP", state, target=target)
        if i % max(1, ramp_steps // 4) == 0:
            actual = motor_pos(state)
            print(
                f"  ramp {i:4d}/{ramp_steps} "
                f"target_err={np.max(np.abs(target - DEFAULT_DOF_POS)):.3f} "
                f"actual_err={np.max(np.abs(actual - DEFAULT_DOF_POS)):.3f}"
            )
        next_t += dt
        sleep = next_t - time.perf_counter()
        if sleep > 0:
            time.sleep(sleep)
        else:
            next_t = time.perf_counter()

    hold_steps = max(1, int(0.5 * args.ramp_hz))
    for i in range(hold_steps):
        if EXIT:
            return state
        new_state = client.recv_latest()
        if new_state is not None:
            state = new_state
        send_hold_cmd(client, state, args)
        if logger is not None and (i % log_every == 0 or i == hold_steps - 1):
            log_state(logger, step_base * 100000 + ramp_steps + i, "RAMP_HOLD", state, target=DEFAULT_DOF_POS)
        next_t += dt
        sleep = next_t - time.perf_counter()
        if sleep > 0:
            time.sleep(sleep)
        else:
            next_t = time.perf_counter()

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


def state_ok(state, args):
    q = motor_pos(state)
    dq = motor_vel(state)
    gyro = as_np(state.imu.gyroscope)
    quat = as_np(state.imu.quaternion)
    grav = get_projected_gravity(quat)
    rpy_deg = np.degrees(as_np(state.imu.rpy))

    checks = [
        (np.all(np.isfinite(q)), "joint position is non-finite"),
        (np.all(np.isfinite(dq)), "joint velocity is non-finite"),
        (np.all(np.isfinite(gyro)), "gyro is non-finite"),
        (np.all(np.isfinite(quat)), "quaternion is non-finite"),
        (0.5 <= np.linalg.norm(grav) <= 1.5, f"gravity norm={np.linalg.norm(grav):.3f}"),
        (np.max(np.abs(dq)) <= args.max_dof_vel, f"max dof vel={np.max(np.abs(dq)):.2f}"),
        (np.max(np.abs(gyro)) <= args.max_gyro, f"max gyro={np.max(np.abs(gyro)):.2f}"),
        (abs(rpy_deg[0]) <= args.max_roll_deg, f"roll={rpy_deg[0]:.1f} deg"),
        (abs(rpy_deg[1]) <= args.max_pitch_deg, f"pitch={rpy_deg[1]:.1f} deg"),
    ]
    for ok, reason in checks:
        if not ok:
            return False, reason
    return True, "ok"


def action_ok(action_raw, args):
    if not np.all(np.isfinite(action_raw)):
        return False, "action is non-finite"
    max_abs = float(np.max(np.abs(action_raw)))
    hard_limit = float(args.action_hard_trip_limit)
    if hard_limit > 0 and max_abs > hard_limit:
        return False, f"action abs {max_abs:.2f} > hard trip {hard_limit:.2f}"
    if max_abs > args.action_trip_limit:
        return True, (
            f"action abs {max_abs:.2f} > soft trip {args.action_trip_limit:.2f}; "
            f"clipped to {args.action_clip:.2f}"
        )
    return True, "ok"


def clamp_action(action_raw, args):
    return np.clip(action_raw, -args.action_clip, args.action_clip).astype(np.float32)


def smooth_action(action, prev_action, args):
    alpha = float(args.action_ema_alpha)
    if 0.0 < alpha < 1.0:
        return (alpha * action + (1.0 - alpha) * prev_action).astype(np.float32)
    return action.astype(np.float32)


def limit_target_step(target, prev_target, args):
    limit = float(args.max_target_step)
    if limit <= 0:
        return target.astype(np.float32)
    delta = np.clip(target - prev_target, -limit, limit)
    return (prev_target + delta).astype(np.float32)


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}" if user else host
    print(f"[INFO] Equivalent manual command: {SPORT_KILL_CMD}")
    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:
            print("[INFO] Sport processes killed.")
            return True
        stderr = result.stderr.strip()
        if "no process" in stderr.lower() or not stderr:
            print("[INFO] No sport processes found.")
            return True
        print(f"[WARN] SSH returned {result.returncode}: {stderr}")
    except FileNotFoundError:
        print("[WARN] ssh command not found; kill sport processes manually.")
    except subprocess.TimeoutExpired:
        print("[WARN] SSH timed out; check Pi network.")
    except Exception as exc:
        print(f"[WARN] Failed to kill sport processes: {exc}")
    return False


def connect_client(args):
    validate_joint_order()

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

    print(f"[INFO] LowCmd backend: {LOWCMD_BACKEND}")
    print("[INFO] Connecting to MCU...")
    client = FastMCUClient()
    print("[INFO] Waking MCU...")
    client.wake_mcu(n_frames=50, dt=0.01)
    state = client.recv_state(timeout=2.0)
    if state is None:
        client.close()
        raise RuntimeError("No LowState received. Check robot network and sport processes.")

    print(f"[INFO] Connected. Battery={state.bms.SOC}%")
    print(f"[INFO] RPY deg: {np.round(np.degrees(as_np(state.imu.rpy)), 1)}")
    print("[INFO] Initial joint positions (rad):")
    q = motor_pos(state)
    for i, name in enumerate(JOINT_NAMES):
        print(f"  [{i:02d}] {name:4s}: q={q[i]:+7.3f}, default={DEFAULT_DOF_POS[i]:+7.3f}")
    return client, state


def log_state(logger, step, mode, state, cmd=None, cmd_raw=None, obs_single=None, action_raw=None,
              action_safe=None, target=None, state_reason="ok", timing=None):
    timing = {} if timing is None else timing
    logger.log(
        step,
        mode=mode,
        battery_soc=state.bms.SOC,
        rc_lx=state.remote.lx,
        rc_ly=state.remote.ly,
        rc_rx=state.remote.rx,
        rc_ry=state.remote.ry,
        rc_buttons=state.remote.pressed,
        imu_rpy_deg=np.degrees(as_np(state.imu.rpy)),
        imu_quat=as_np(state.imu.quaternion),
        base_ang_vel=as_np(state.imu.gyroscope),
        projected_gravity=get_projected_gravity(as_np(state.imu.quaternion)),
        dof_pos=motor_pos(state),
        dof_vel=motor_vel(state),
        tau_est=motor_tau(state),
        motor_mode=motor_mode(state),
        motor_temperature=motor_temperature(state),
        motor_reserve=motor_reserve(state),
        commands_raw=np.zeros(3, dtype=np.float32) if cmd_raw is None else cmd_raw,
        commands=np.zeros(3, dtype=np.float32) if cmd is None else cmd,
        obs_single=np.zeros(NUM_OBS, dtype=np.float32) if obs_single is None else obs_single,
        action_raw=np.zeros(NUM_ACTIONS, dtype=np.float32) if action_raw is None else action_raw,
        action_safe=np.zeros(NUM_ACTIONS, dtype=np.float32) if action_safe is None else action_safe,
        joint_targets=np.zeros(NUM_ACTIONS, dtype=np.float32) if target is None else target,
        state_reason=state_reason,
        loop_dt_ms=float(timing.get("loop_dt_ms", 0.0)),
        recv_ms=float(timing.get("recv_ms", 0.0)),
        policy_ms=float(timing.get("policy_ms", 0.0)),
        work_ms=float(timing.get("work_ms", 0.0)),
    )


def run_monitor(args):
    print(MONITOR_BANNER)
    input("Press Enter to start monitor...")
    logger = JsonlLogger(args.log_dir, args)
    client = None
    try:
        client, state = connect_client(args)
        edge = RCEdgeDetector()
        edge.update(state)
        step = 0
        dt = 1.0 / args.rate_hz
        next_t = time.perf_counter()
        while not EXIT:
            new_state = client.recv_latest()
            if new_state is not None:
                state = new_state
            rising, falling = edge.update(state)
            if step % args.print_every == 0:
                rpy = np.degrees(as_np(state.imu.rpy))
                print(f"\n[MONITOR {step}] bat={state.bms.SOC}% rpy={np.round(rpy, 1)}")
                print(f"  RC: {fmt_rc(state)}")
                if rising:
                    print(f"  rising: {sorted(rising)}")
                if falling:
                    print(f"  falling: {sorted(falling)}")
                print(f"  q:  {np.round(motor_pos(state), 3)}")
                print(f"  dq: {np.round(motor_vel(state), 3)}")
            log_state(logger, step, "MONITOR", state)
            step += 1
            if args.max_steps > 0 and step >= args.max_steps:
                break
            next_t += dt
            sleep = next_t - time.perf_counter()
            if sleep > 0:
                time.sleep(sleep)
            else:
                next_t = time.perf_counter()
    finally:
        logger.close()
        if client is not None:
            client.close()


def run_obs_check(args):
    print(OBS_BANNER)
    input("Press Enter to start obs-check...")
    logger = JsonlLogger(args.log_dir, args)
    client = None
    try:
        client, state = connect_client(args)
        obs_builder = ObsHistoryBuilder()
        cmd_filter = CommandFilter(args)
        last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
        step = 0
        dt = 1.0 / args.rate_hz
        next_t = time.perf_counter()
        while not EXIT:
            new_state = client.recv_latest()
            if new_state is not None:
                state = new_state
            cmd_raw = get_command(state, args)
            cmd = cmd_filter.update(cmd_raw)
            obs_single = obs_builder.build_single(state, cmd, last_action)
            policy_input = obs_builder.build_policy_input(obs_single)
            if step % args.print_every == 0:
                print(f"\n[OBS {step}] bat={state.bms.SOC}% {fmt_rc(state)}")
                print(f"  obs[0:3] gyro      {np.round(obs_single[0:3], 4)}")
                print(f"  obs[3:6] gravity   {np.round(obs_single[3:6], 4)}")
                print(f"  obs[6:9] cmd       {np.round(obs_single[6:9], 4)} raw={np.round(cmd, 3)}")
                print(f"  obs[9:21] q-qd     max={np.max(np.abs(obs_single[9:21])):.3f}")
                print(f"  obs[21:33] dq      max={np.max(np.abs(obs_single[21:33])):.3f}")
                print(f"  policy_input shape={policy_input.shape} min={policy_input.min():.3f} max={policy_input.max():.3f}")
            log_state(logger, step, "OBS_CHECK", state, cmd=cmd, cmd_raw=cmd_raw, obs_single=obs_single)
            step += 1
            if args.max_steps > 0 and step >= args.max_steps:
                break
            next_t += dt
            sleep = next_t - time.perf_counter()
            if sleep > 0:
                time.sleep(sleep)
            else:
                next_t = time.perf_counter()
    finally:
        logger.close()
        if client is not None:
            client.close()


def run_infer_check(args):
    print(INFER_BANNER)
    input("Press Enter to start infer-check...")
    logger = JsonlLogger(args.log_dir, args)
    client = None
    try:
        policy = make_policy(args)
        client, state = connect_client(args)
        obs_builder = ObsHistoryBuilder()
        cmd_filter = CommandFilter(args)
        last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
        prev_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
        step = 0
        dt = 1.0 / args.rate_hz
        next_t = time.perf_counter()
        while not EXIT:
            new_state = client.recv_latest()
            if new_state is not None:
                state = new_state
            cmd_raw = get_command(state, args)
            cmd = cmd_filter.update(cmd_raw)
            obs_single = obs_builder.build_single(state, cmd, last_action)
            policy_input = obs_builder.build_policy_input(obs_single)
            action_raw = policy(policy_input)
            ok, reason = action_ok(action_raw, args)
            action_safe = smooth_action(clamp_action(action_raw, args), prev_action, args)
            prev_action = action_safe.copy()
            last_action = action_safe.copy()
            if step % args.print_every == 0:
                print(f"\n[INFER {step}] ok={ok} reason={reason}")
                print(f"  cmd={np.round(cmd, 3)} action_raw={np.round(action_raw[:4], 3)} max={np.max(np.abs(action_raw)):.3f}")
                print(f"  action_safe max={np.max(np.abs(action_safe)):.3f}")
            log_state(
                logger, step, "INFER_CHECK", state, cmd=cmd, obs_single=obs_single,
                cmd_raw=cmd_raw,
                action_raw=action_raw, action_safe=action_safe, state_reason=reason,
            )
            if not ok and args.trip_on_infer_check:
                print(f"[FAULT] {reason}")
                break
            step += 1
            if args.max_steps > 0 and step >= args.max_steps:
                break
            next_t += dt
            sleep = next_t - time.perf_counter()
            if sleep > 0:
                time.sleep(sleep)
            else:
                next_t = time.perf_counter()
    finally:
        logger.close()
        if client is not None:
            client.close()


STARTUP_BANNER = """
============================================================
 RoboGauge Go1 RL-Gym BPU deployment

 State machine:
   IDLE --R2--> CALIBRATE --> HOLD --R2--> OBS_TEST
   OBS_TEST --R2--> INFER_TEST --R2--> RL
   Any active state --L2--> IDLE damping

 RL motor commands require --enable-rl. Without it, R2 at INFER_TEST
 will stay in INFER_TEST.

 Initial tests should be done with the robot suspended.
 Use --kill-sport to run the Pi sport-process kill step before MCU control.
 Manual equivalent:
   ssh pi@192.168.123.161 "sudo pkill -9 -f keep_sport_alive; sudo pkill -9 -f Legged_sport; sudo pkill -9 -f appTransit"
 After sport processes are killed, keep battery removal available as
 the final stop method; the original sport-mode remote combo is not active.
============================================================
"""

MONITOR_BANNER = """
============================================================
 MONITOR: read RC, IMU, and joint state only. No motor command.
============================================================
"""

OBS_BANNER = """
============================================================
 OBS-CHECK: build 45-dim obs and 225-dim history only.
 No motor command.
============================================================
"""

INFER_BANNER = """
============================================================
 INFER-CHECK: build obs and run BPU only.
 No motor command.
============================================================
"""


def run_deploy(args):
    print(STARTUP_BANNER)
    input("Press Enter when ready...")

    policy = make_policy(args)
    logger = JsonlLogger(args.log_dir, args)
    client = None
    state = None
    sm_state = State.IDLE
    step = 0
    cmd_raw = np.zeros(3, dtype=np.float32)
    cmd = np.zeros(3, dtype=np.float32)
    obs_single = None
    action_raw = np.zeros(NUM_ACTIONS, dtype=np.float32)
    action_safe = np.zeros(NUM_ACTIONS, dtype=np.float32)
    target = DEFAULT_DOF_POS.copy()
    reason = "ok"

    try:
        client, state = connect_client(args)
        edge = RCEdgeDetector()
        edge.update(state)
        obs_builder = ObsHistoryBuilder()
        cmd_filter = CommandFilter(args)
        servo_guard = MotorServoGuard(args.motor_mode_fault_frames)

        sm_state = State.IDLE
        last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
        prev_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
        prev_target = DEFAULT_DOF_POS.copy()
        step = 0
        rl_step = 0
        dt = 1.0 / args.rate_hz
        next_t = time.perf_counter()
        prev_loop_t = None

        print("[INFO] R2 advances layers. L2 emergency-stops to IDLE.")
        print("[INFO] Ctrl+C exits with safe_stop.")

        while not EXIT:
            loop_t0 = time.perf_counter()
            timing = None
            if args.log_timing:
                timing = {
                    "loop_dt_ms": 0.0 if prev_loop_t is None else (loop_t0 - prev_loop_t) * 1000.0,
                    "recv_ms": 0.0,
                    "policy_ms": 0.0,
                    "work_ms": 0.0,
                }
            prev_loop_t = loop_t0

            recv_t0 = time.perf_counter()
            new_state = client.recv_latest()
            if timing is not None:
                timing["recv_ms"] = (time.perf_counter() - recv_t0) * 1000.0
            if new_state is not None:
                state = new_state
            if state is None:
                time.sleep(0.001)
                continue

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

            if l2_rose and sm_state != State.IDLE:
                print(f"\n[L2] {sm_state.value} -> IDLE damping")
                sm_state = State.IDLE
                obs_builder.reset()
                cmd_filter.reset()
                last_action[:] = 0.0
                prev_action[:] = 0.0
                prev_target = DEFAULT_DOF_POS.copy()
                rl_step = 0
                servo_guard.reset()
                send_damping(client)

            if sm_state in (State.HOLD, State.OBS_TEST, State.INFER_TEST, State.RL):
                servo_fault = servo_guard.update(state)
            else:
                servo_guard.reset()
                servo_fault = None

            cmd_raw = get_command(state, args)
            if sm_state in (State.OBS_TEST, State.INFER_TEST, State.RL):
                cmd = cmd_filter.update(cmd_raw)
            else:
                cmd_filter.reset()
                cmd = cmd_raw
            obs_single = None
            action_raw = np.zeros(NUM_ACTIONS, dtype=np.float32)
            action_safe = np.zeros(NUM_ACTIONS, dtype=np.float32)
            target = DEFAULT_DOF_POS.copy()
            reason = "ok"

            if sm_state == State.IDLE:
                if step % 10 == 0:
                    send_damping(client)
                if r2_rose:
                    print("\n[R2] IDLE -> CALIBRATE")
                    sm_state = State.CALIBRATE
                    state = ramp_to_default(client, args, state, logger=logger, step_base=step)
                    obs_builder.reset()
                    cmd_filter.reset()
                    last_action[:] = 0.0
                    prev_action[:] = 0.0
                    prev_target = DEFAULT_DOF_POS.copy()
                    sm_state = State.HOLD
                    print("[STATE] HOLD")

            elif sm_state == State.HOLD:
                send_hold_cmd(client, state, args)
                if r2_rose:
                    if servo_fault:
                        reason = servo_fault
                        print(f"\n[FAULT] HOLD -> OBS_TEST blocked: {reason}")
                        sm_state = State.FAULT
                        send_damping(client)
                    else:
                        print("\n[R2] HOLD -> OBS_TEST")
                        obs_builder.reset()
                        cmd_filter.reset()
                        last_action[:] = 0.0
                        prev_action[:] = 0.0
                        sm_state = State.OBS_TEST

            elif sm_state == State.OBS_TEST:
                send_hold_cmd(client, state, args)
                obs_single = obs_builder.build_single(state, cmd, last_action)
                obs_builder.build_policy_input(obs_single)
                if servo_fault:
                    ok, reason = False, servo_fault
                else:
                    ok, reason = state_ok(state, args)
                if not ok:
                    print(f"\n[FAULT] OBS_TEST state check failed: {reason}")
                    sm_state = State.FAULT
                    send_damping(client)
                elif r2_rose:
                    print("\n[R2] OBS_TEST -> INFER_TEST")
                    obs_builder.reset()
                    cmd_filter.reset()
                    last_action[:] = 0.0
                    prev_action[:] = 0.0
                    sm_state = State.INFER_TEST

            elif sm_state == State.INFER_TEST:
                send_hold_cmd(client, state, args)
                obs_single = obs_builder.build_single(state, cmd, last_action)
                policy_input = obs_builder.build_policy_input(obs_single)
                if servo_fault:
                    ok, reason = False, servo_fault
                else:
                    ok, reason = state_ok(state, args)
                if ok:
                    policy_t0 = time.perf_counter()
                    action_raw = policy(policy_input)
                    if timing is not None:
                        timing["policy_ms"] += (time.perf_counter() - policy_t0) * 1000.0
                    ok, reason = action_ok(action_raw, args)
                if ok:
                    action_safe = smooth_action(clamp_action(action_raw, args), prev_action, args)
                    prev_action = action_safe.copy()
                    last_action = action_safe.copy()
                else:
                    print(f"\n[FAULT] INFER_TEST failed: {reason}")
                    sm_state = State.FAULT

                if r2_rose and sm_state == State.INFER_TEST:
                    if not args.enable_rl:
                        print("\n[GUARD] RL blocked. Re-run with --enable-rl after OBS/INFER logs look safe.")
                    else:
                        print("\n[R2] INFER_TEST -> RL")
                        obs_builder.reset()
                        cmd_filter.reset()
                        last_action[:] = 0.0
                        prev_action[:] = 0.0
                        prev_target = DEFAULT_DOF_POS.copy()
                        rl_step = 0
                        sm_state = State.RL

            elif sm_state == State.RL:
                if r2_rose:
                    print("\n[R2] RL -> HOLD")
                    sm_state = State.HOLD
                    obs_builder.reset()
                    cmd_filter.reset()
                    last_action[:] = 0.0
                    prev_action[:] = 0.0
                    state = ramp_to_default(client, args, state, logger=logger, step_base=step)
                    prev_target = DEFAULT_DOF_POS.copy()
                    rl_step = 0
                    continue

                obs_single = obs_builder.build_single(state, cmd, last_action)
                policy_input = obs_builder.build_policy_input(obs_single)
                if servo_fault:
                    ok, reason = False, servo_fault
                else:
                    ok, reason = state_ok(state, args)
                if ok:
                    policy_t0 = time.perf_counter()
                    action_raw = policy(policy_input)
                    if timing is not None:
                        timing["policy_ms"] += (time.perf_counter() - policy_t0) * 1000.0
                    ok, reason = action_ok(action_raw, args)
                if not ok:
                    print(f"\n[FAULT] RL failed: {reason}")
                    sm_state = State.FAULT
                    send_damping(client)
                else:
                    action_clipped = clamp_action(action_raw, args)
                    action_safe = smooth_action(action_clipped, prev_action, args)
                    target_raw = DEFAULT_DOF_POS + action_safe * ACTION_SCALE
                    target = limit_target_step(target_raw, prev_target, args)
                    prev_action = action_safe.copy()
                    last_action = action_safe.copy()
                    if rl_step >= args.warmup_steps:
                        prev_target = target.copy()
                        send_position_cmd(client, state, target, args)
                    else:
                        target = DEFAULT_DOF_POS.copy()
                        prev_target = DEFAULT_DOF_POS.copy()
                        send_hold_cmd(client, state, args)
                    rl_step += 1

            elif sm_state == State.FAULT:
                send_damping(client)
                obs_builder.reset()
                cmd_filter.reset()
                last_action[:] = 0.0
                prev_action[:] = 0.0
                if r2_rose:
                    print("\n[R2] FAULT -> CALIBRATE")
                    sm_state = State.CALIBRATE
                    state = ramp_to_default(client, args, state, logger=logger, step_base=step)
                    prev_target = DEFAULT_DOF_POS.copy()
                    sm_state = State.HOLD
                    print("[STATE] HOLD")

            if timing is not None:
                timing["work_ms"] = (time.perf_counter() - loop_t0) * 1000.0

            log_state(
                logger,
                step,
                sm_state.value,
                state,
                cmd=cmd,
                cmd_raw=cmd_raw,
                obs_single=obs_single,
                action_raw=action_raw,
                action_safe=action_safe,
                target=target,
                state_reason=reason,
                timing=timing,
            )

            if step % args.print_every == 0:
                rpy = np.degrees(as_np(state.imu.rpy))
                print(
                    f"\n[STEP {step}] state={sm_state.value} bat={state.bms.SOC}% "
                    f"rpy={np.round(rpy, 1)}"
                )
                print(f"  RC: {fmt_rc(state)}")
                print(f"  cmd={np.round(cmd, 3)} q={np.round(motor_pos(state), 2)}")
                print(
                    f"  motor_mode[FR]={motor_mode(state)[:3].tolist()} "
                    f"temp[FR]={motor_temperature(state)[:3].tolist()}"
                )
                if sm_state in (State.INFER_TEST, State.RL, State.FAULT):
                    print(
                        f"  action_raw_max={np.max(np.abs(action_raw)):.3f} "
                        f"action_safe_max={np.max(np.abs(action_safe)):.3f} reason={reason}"
                    )
                if sm_state == State.RL:
                    print(f"  target={np.round(target, 2)} rl_step={rl_step}")

            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()

    except PowerProtectViolation as exc:
        reason = f"power protect: {exc}"
        print(f"\n[FAULT] {reason}")
        if client is not None:
            send_damping(client)
        if state is not None:
            log_state(
                logger,
                step,
                State.FAULT.value,
                state,
                cmd=cmd,
                cmd_raw=cmd_raw,
                obs_single=obs_single,
                action_raw=action_raw,
                action_safe=action_safe,
                target=target,
                state_reason=reason,
            )

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


def build_arg_parser():
    parser = argparse.ArgumentParser(description="Deploy RoboGauge Go1 RL-Gym BPU on Go1 PRO")
    parser.add_argument("--bpu-model", default="",
                        help="Path to a compiled Gym 5-frame BPU .bin; overrides --bpu-round")
    parser.add_argument("--bpu-round", default=DEFAULT_BPU_ROUND,
                        choices=sorted(BPU_MODEL_REGISTRY),
                        help="Quick-select mapper_output_<round>_gemm/policy_<round>_int16_gemm.bin")
    parser.add_argument("--bpu-priority", type=int, default=0)
    parser.add_argument("--bpu-cores", type=int, nargs="+", default=[0],
                        help="Reserved BPU core ids for future runtime scheduling")

    parser.add_argument("--kill-sport", action="store_true", help="Kill Pi sport processes via SSH")
    parser.add_argument("--pi-host", default="192.168.123.161")
    parser.add_argument("--pi-user", default="pi")

    parser.add_argument("--monitor", action="store_true", help="Read RC/IMU/joints only; no motor command")
    parser.add_argument("--obs-check", action="store_true", help="Build obs/history only; no motor command")
    parser.add_argument("--infer-check", action="store_true", help="Run BPU inference only; no motor command")
    parser.add_argument("--enable-rl", action="store_true", help="Allow state machine to enter RL motor-control state")
    parser.add_argument("--no-rc", action="store_true", help="Use fixed --cmd-* instead of RC sticks")
    parser.add_argument("--swap-vy-yaw", action="store_true", help="Map left stick x to yaw and right stick x to vy")

    parser.add_argument("--kp", type=float, default=80.0)
    parser.add_argument("--kd", type=float, default=1.0)
    parser.add_argument("--kp-cal", type=float, default=20.0)
    parser.add_argument("--kd-cal", type=float, default=1.0)
    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=MAX_LIN_VEL_X)
    parser.add_argument("--rc-vy-scale", type=float, default=MAX_LIN_VEL_Y)
    parser.add_argument("--rc-wz-scale", type=float, default=MAX_ANG_VEL_YAW)
    parser.add_argument("--rc-deadzone", type=float, default=0.05)
    parser.add_argument("--cmd-ema-alpha", type=float, default=1.0)
    parser.add_argument("--max-cmd-step-x", type=float, default=0.0)
    parser.add_argument("--max-cmd-step-y", type=float, default=0.0)
    parser.add_argument("--max-cmd-step-yaw", type=float, default=0.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("--ramp-time", type=float, default=5.0)
    parser.add_argument("--ramp-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("--max-roll-deg", type=float, default=35.0)
    parser.add_argument("--max-pitch-deg", type=float, default=35.0)
    parser.add_argument("--max-dof-vel", type=float, default=30.0)
    parser.add_argument("--max-gyro", type=float, default=15.0)
    parser.add_argument("--action-trip-limit", type=float, default=12.0,
                        help="Soft raw-action warning threshold; output is still clipped to --action-clip")
    parser.add_argument("--action-hard-trip-limit", type=float, default=16.0,
                        help="Hard raw-action fault threshold; <=0 disables the hard trip")
    parser.add_argument("--action-clip", type=float, default=4.0)
    parser.add_argument("--action-ema-alpha", type=float, default=0.0)
    parser.add_argument("--max-target-step", type=float, default=0.08)
    parser.add_argument("--trip-on-infer-check", action="store_true")

    parser.add_argument("--log-dir", default=str(HERE / "logs"), help="JSONL log directory; empty disables logging")
    parser.add_argument("--log-flush-every", type=int, default=50)
    parser.add_argument("--motor-mode-fault-frames", type=int, default=3,
                        help="Consecutive non-servo feedback frames required before a motor-mode fault")
    parser.add_argument("--log-timing", action="store_true",
                        help="Log per-loop timing fields for deploy-mode latency diagnosis")
    return parser


def main():
    args = build_arg_parser().parse_args()
    if args.monitor:
        return run_monitor(args)
    if args.obs_check:
        return run_obs_check(args)
    if args.infer_check:
        return run_infer_check(args)
    return run_deploy(args)


if __name__ == "__main__":
    main()
