// homing.ts: the source of a port of Hello Robot's stretch_body, offered under the GNU Lesser General Public License v3.0.
//
// This is the file as the page at joosep.parts/stretch is built from it. It is a work based on stretch_body 0.7.31
// (https://github.com/hello-robot/stretch_body, robot.py, lift.py, arm.py, prismatic_joint.py, head.py, end_of_arm.py,
// wrist_yaw.py, wrist_pitch.py, wrist_roll.py, stretch_gripper.py, dynamixel_hello_XL430.py and the homing parameters of
// robot_params_SE3.py), Copyright 2020-2024 Hello Robot Inc., which is free software under the GNU Lesser General Public
// License v3.0. The changes are the translation into TypeScript and the differences listed in the header of the code
// below. As a work based on it, this file is offered under the same licence, and without any warranty. The page's script
// holds this code compiled into it.
//
// Licence texts, served next to this file: stretch_body-LICENSE.txt (Hello Robot's notice), LGPL-3.0.txt and GPL-3.0.txt
// (the GNU General Public License v3.0, which the LGPL adds to).
//
// Hello Robot's homing in the twin: Robot.home() of stretch_body 0.7.31 (master @92c9db0) for a Stretch 3 with the Dex Wrist 3
// and the SG3 gripper (eoa_wrist_dw3_tool_sg3), in its order, directions, speeds and end positions, in simulated time,
// through the firmware's own calls (firmware.ts). It is a port of robot.py, lift.py, arm.py, prismatic_joint.py, head.py,
// end_of_arm.py, wrist_yaw.py, wrist_pitch.py, wrist_roll.py, stretch_gripper.py, dynamixel_hello_XL430.py and the homing
// parameters of robot_params_SE3.py.
//
// The teleop calls robot.home() on Start while the robot is not homed (teleop.ts); the original blocks in that call, so here
// PadPath starts this sequence and runs `tick` once per teleop cycle, and `homed` is what robot.is_homed() reads back (the
// lift, the arm, the wrist yaw and the gripper all found their stop).
//
// Sources, in stretch_body 0.7.31, body/stretch_body/:
//   robot.py 609-650            home(): the head, a thread that pitches the wrist 1 s later, the lift, the arm, the end of
//                               arm, then a beep
//   head.py 67-77               the head joints need no homing (req_calibration 0, RP head_pan, head_tilt): move_to 0 each,
//                               not waited for
//   lift.py 29-30, arm.py 26-27, prismatic_joint.py 691-790
//                               a stepper's homing: a guarded move_by "well past the stop" at the joint's default profile
//                               (prismatic_joint.py 32-33, 392-397), the contact (15 s), the stop marked, 1 s, move_to the end
//                               position, 1 s, until near the setpoint (stepper.py 501-512, 15 s)
//   end_of_arm.py 20, 91-102    the tool's joints in the order of its devices (RP SE3_eoa_wrist_dw3_tool_sg3): wrist_pitch,
//                               wrist_roll, wrist_yaw, stretch_gripper
//   wrist_pitch.py 19-23, wrist_roll.py 36-40   pass
//   wrist_yaw.py 20-25, stretch_gripper.py 33-34 (StretchGripper3 for the SG3), dynamixel_hello_XL430.py 1101-1229
//                               a Dynamixel's homing: PWM towards the first stop, 1 s, until it stops moving (15 s), the
//                               gripper 2.25 s more at the stop, the stop marked, move_to 0 until its profile ends (387-401,
//                               6 s); the wrist yaw then waits 2 s more
//   robot_params_SE3.py (RP)    the speeds, limits and PWM, quoted at each step
//
// Where it differs from the original, and why:
//   * A stepper's move "well past the stop" (5 m) is clamped by the firmware to the end of the joint's range and lands there
//     on its profile; the real stepper runs into the stop at the profile's speed and stops on its motor effort
//     (contact_thresh_homing). Contact here: the setpoint has reached that end and the joint is within its distance_tol.
//   * The sim has no PWM: a Dynamixel goes to the end of its range at the PWM's share of its no-load speed (pwm_homing over
//     pwm_limit 885) and the joint's 'max' acceleration, and it stops moving below the servo's moving threshold (10 x 0.229
//     rpm at the servo, dynamixel_XL430.py 975, over the joint's gear: RP gr 2.4 for the wrist yaw, 1.0 for the gripper).
//   * Marking a stop changes nothing: the sim's joint positions are right from the start.
//   * Collision management, off while the original homes (robot.py 614, 645), is not modelled in the sim.
//
// The robot is switched on (a page load) in POWER_ON_POSE, not homed (`park`).
import type { MainModule, MjData, MjModel } from '@mujoco/mujoco'
import { GRIPPER_SLIDE_PER_RAD, LIMITS, type AxisName, type Firmware } from './firmware'
import type { Command } from './teleop'
import type { HomeStep } from './protocol'

/** the pose the robot is switched on in (Proposal): stowed, as Hello Robot's shutdown leaves it (Back held 2 s stows the
 *  robot, then shuts its computer down: gamepad_teleop.py 435-454; Robot.stow, robot.py 559-607, with the tool's stow pose,
 *  RP SE3_eoa_wrist_dw3_tool_sg3: lift 0.3 m, arm in, wrist yaw 3.0 rad, pitch -0.62 rad, roll 0, the head at 0). The
 *  gripper keeps its opening of the scene's start: its fingers are tied to its slide by equality constraints that a pose set
 *  by hand would break. */
export const POWER_ON_POSE: Readonly<Record<string, number>> = {
  joint_lift: 0.3, joint_arm_l0: 0, joint_arm_l1: 0, joint_arm_l2: 0, joint_arm_l3: 0,
  joint_wrist_yaw: 3.0, joint_wrist_pitch: -0.62, joint_wrist_roll: 0, joint_head_pan: 0, joint_head_tilt: 0,
}

/** put the robot's joints in POWER_ON_POSE at rest (`joints`: the model's joint names, in order) */
export function park(mj: MainModule, m: MjModel, d: MjData, joints: readonly string[]): void {
  const qadr = m.jnt_qposadr as Int32Array, dadr = m.jnt_dofadr as Int32Array
  const q = d.qpos as Float64Array, v = d.qvel as Float64Array
  for (const [name, x] of Object.entries(POWER_ON_POSE)) {
    const j = joints.indexOf(name)
    if (j < 0) continue
    q[qadr[j]] = x
    v[dadr[j]] = 0
  }
  mj.mj_forward(m, d)
}

/** a Dynamixel's moving threshold at the joint (rad/s; the gripper's in m/s of slide) */
const MOVING = { wrist_yaw: 10 * 0.229 * 2 * Math.PI / 60 / 2.4, gripper: 10 * 0.229 * 2 * Math.PI / 60 * GRIPPER_SLIDE_PER_RAD }

/** the original's waits (s) */
export const HOME_WAIT = {
  /** the pitch thread's sleep before it moves the wrist pitch (robot.py 623) */
  pitch: 1.0,
  /** wait_for_contact (prismatic_joint.py 735) and wait_until_at_setpoint (stepper.py 501) */
  stepper: 15.0,
  /** after the contact, and after the move to the end position (prismatic_joint.py 748, 769) */
  settle: 1.0,
  /** a Dynamixel's PWM before it is asked whether it moves, and the most it moves for (dynamixel_hello_XL430.py 1138-1153) */
  dxlFirst: 1.0,
  dxlStop: 15.0,
  /** the gripper's delay_at_stop (stretch_gripper.py 34) */
  gripperHold: 2.25,
  /** wait_until_at_setpoint after move_to(0) (dynamixel_hello_XL430.py 1221) */
  dxlZero: 6.0,
  /** WristYaw.home's extra time to get back (wrist_yaw.py 25) */
  yaw: 2.0,
}

/** the end positions: lift.py 29 (0.6 m), arm.py 26 (0.1 m), the tool's stow pitch (RP SE3_eoa_wrist_dw3_tool_sg3 stow,
 *  -0.62 rad) */
export const HOME_END = { lift: 0.6, arm: 0.1, wrist_pitch: -0.62 }

/** a Dynamixel's homing PWM as a share of its full PWM (RP pwm_homing[0] over pwm_limit 885): wrist yaw -300, SG3 -400 */
const PWM_SHARE = { wrist_yaw: 300 / 885, gripper: 400 / 885 }
/** their speeds there (rad/s at the joint; the gripper's servo): that share of the no-load speed (firmware.ts LIMITS.CLAMP:
 *  the yaw 3.05 rad/s, the gripper's servo 6.0 rad/s), and the joint's 'max' acceleration */
export const HOME_PWM = {
  wrist_yaw: { v: PWM_SHARE.wrist_yaw * LIMITS.CLAMP.wrist_yaw[0], a: LIMITS.CLAMP.wrist_yaw[1] },
  gripper: { v: PWM_SHARE.gripper * LIMITS.CLAMP.gripper[0] / GRIPPER_SLIDE_PER_RAD, a: LIMITS.CLAMP.gripper[1] / GRIPPER_SLIDE_PER_RAD },
}

type Joint = 'lift' | 'arm' | 'wrist_yaw' | 'gripper'

const PAST_STOP = 5.0              // m, "well past the stop" (prismatic_joint.py 727-731)

/** the sequence, one step after another; a joint whose stop is not found skips the rest of its own steps, as the original */
const STEPS: readonly HomeStep[] = ['lift', 'lift_mark', 'lift_end', 'arm', 'arm_mark', 'arm_end', 'yaw', 'yaw_end', 'yaw_wait',
  'gripper', 'gripper_hold', 'gripper_end']

export class Homing {
  /** robot.is_homed(): every joint that needs homing found its stop since the robot was powered on (a page load) */
  homed = true
  /** robot.home() runs */
  running = false
  /** the step that runs now (null when not running) */
  step: HomeStep | null = null
  /** simulated time robot.home() began, and the step began */
  startedAt = 0
  private stepAt = 0
  private i = 0
  private found: Record<Joint, boolean> = { lift: false, arm: false, wrist_yaw: false, gripper: false }
  /** the pitch thread's call, at its time (robot.py 620-628) */
  private pitchAt: number | null = null
  /** the latest call per target: made again after the firmware is started anew (a hand-over during the homing) */
  private readonly issued = new Map<string, Command>()
  /** how the last homing ended: true every stop found, false a stop not found, null none ended yet */
  lastResult: boolean | null = null

  /** the robot is switched on: not homed, nothing running */
  powerOn(): void {
    this.homed = false
    this.abort()
  }

  /** homed without the sequence (the gate's suites start homed, as the Python environment's episodes do) */
  assumeHomed(): void {
    this.homed = true
    this.abort()
  }

  /** stop a homing that runs (a fresh scene: the robot is put back in its start pose and is still not homed) */
  abort(): void {
    this.running = false
    this.step = null
    this.pitchAt = null
    this.issued.clear()
  }

  /** robot.home() at simulated time t: the head to 0, the pitch thread, then the lift's homing (robot.py 609-631) */
  start(fw: Firmware, t: number): void {
    if (this.running) return
    this.running = true
    this.startedAt = t
    this.found = { lift: false, arm: false, wrist_yaw: false, gripper: false }
    this.issued.clear()
    // head.py 75-76: move_to 0 at the default profile (3.0 rad/s, 8.0 rad/s^2, RP head_pan and head_tilt motion default)
    this.call(fw, { target: 'head.head_pan', method: 'move_to', kwargs: { x_r: 0.0 } })
    this.call(fw, { target: 'head.head_tilt', method: 'move_to', kwargs: { x_r: 0.0 } })
    this.pitchAt = t + HOME_WAIT.pitch
    this.i = -1
    this.next(fw, t)
  }

  /** one teleop cycle at simulated time t (after the cycle's own calls): the pitch thread and the step that runs */
  tick(fw: Firmware, t: number): void {
    if (!this.running) return
    if (this.pitchAt !== null && t >= this.pitchAt) {
      // end_of_arm.move_to('wrist_pitch', stow pitch) at its default profile (2.0 rad/s, 6.0 rad/s^2, RP SE3_wrist_pitch_DW3)
      this.call(fw, { target: 'end_of_arm.wrist_pitch', method: 'move_to', kwargs: { x_r: HOME_END.wrist_pitch } })
      this.pitchAt = null
    }
    for (let k = 0; k < STEPS.length && this.running; k++) {
      const r = this.check(fw, t - this.stepAt)
      if (r === null) return
      this.next(fw, t, r)
    }
  }

  /** the firmware was started anew (fw.reset, a hand-over while homing): the homing's moves go on from where the joints are */
  reapply(fw: Firmware): void {
    if (!this.running || !this.issued.size) return
    fw.apply([...this.issued.values()])
  }

  /** the step's progress `dt` seconds into it: true done, false failed (its joint's stop not found), null still running */
  private check(fw: Firmware, dt: number): boolean | null {
    const s = this.step as HomeStep
    switch (s) {
      case 'lift': case 'arm':
        return arrived(fw, s, true) ? true : dt >= HOME_WAIT.stepper ? false : null
      case 'lift_mark': case 'arm_mark':
        return dt >= HOME_WAIT.settle ? true : null
      case 'lift_end': case 'arm_end': {
        // sleep 1.0, then wait_until_at_setpoint: a timeout there is a warning only, the joint stays homed
        if (dt < HOME_WAIT.settle) return null
        return arrived(fw, s === 'lift_end' ? 'lift' : 'arm', true) || dt >= HOME_WAIT.settle + HOME_WAIT.stepper ? true : null
      }
      case 'yaw': case 'gripper': {
        const a = s === 'yaw' ? 'wrist_yaw' : 'gripper'
        if (dt < HOME_WAIT.dxlFirst) return null
        return !moving(fw, a) ? true : dt >= HOME_WAIT.dxlStop ? false : null
      }
      case 'gripper_hold':
        return dt >= HOME_WAIT.gripperHold ? true : null
      case 'yaw_end': case 'gripper_end':
        return arrived(fw, s === 'yaw_end' ? 'wrist_yaw' : 'gripper', false) || dt >= HOME_WAIT.dxlZero ? true : null
      case 'yaw_wait':
        return dt >= HOME_WAIT.yaw ? true : null
    }
  }

  /** the step after this one (`ok`: how this one ended), and its calls */
  private next(fw: Firmware, t: number, ok = true): void {
    const was = this.i >= 0 ? STEPS[this.i] : null
    if (was === 'lift' || was === 'arm') this.found[was] = ok
    if (was === 'yaw') this.found.wrist_yaw = ok
    let i = this.i + 1
    // a joint whose stop was not found: the rest of its own steps are skipped (prismatic_joint.py 764-772; a Dynamixel
    // returns at 1162-1164, after the gripper's delay at the stop, and WristYaw.home still waits its 2 s)
    if (was === 'lift' && !ok) i = STEPS.indexOf('arm')
    if (was === 'arm' && !ok) i = STEPS.indexOf('yaw')
    if (was === 'yaw' && !ok) i = STEPS.indexOf('yaw_wait')
    if (was === 'gripper_hold' && !this.found.gripper) i = STEPS.length
    if (was === 'gripper') this.found.gripper = ok
    this.i = i
    this.stepAt = t
    if (i >= STEPS.length) return this.finish()
    const s = STEPS[i]
    this.step = s
    switch (s) {
      // lift.py 29: to the positive stop, at the default profile (0.11 m/s, 0.2 m/s^2, RP lift motion default)
      case 'lift': return this.call(fw, { target: 'lift', method: 'move_by', kwargs: { x_m: PAST_STOP } })
      case 'lift_end': return this.call(fw, { target: 'lift', method: 'move_to', kwargs: { x_m: HOME_END.lift } })
      // arm.py 26: to the negative stop (0.14 m/s, 0.14 m/s^2, RP arm motion default)
      case 'arm': return this.call(fw, { target: 'arm', method: 'move_by', kwargs: { x_m: -PAST_STOP } })
      case 'arm_end': return this.call(fw, { target: 'arm', method: 'move_to', kwargs: { x_m: HOME_END.arm } })
      // PWM -300: towards fewer ticks, the first stop, 0 ticks, which with zero_t 7136, gear 2.4 and the flipped encoder is
      // +4.56 rad of joint angle (calc), the top of the range (the sim's ends at 4.42 rad)
      case 'yaw': return this.call(fw, { target: 'end_of_arm.wrist_yaw', method: 'move_to',
                                        kwargs: { x_r: 10.0, v_r: HOME_PWM.wrist_yaw.v, a_r: HOME_PWM.wrist_yaw.a } })
      // move_to(0) at the default profile (2.0 rad/s, 3.0 rad/s^2, RP SE3_wrist_yaw_DW3)
      case 'yaw_end': return this.call(fw, { target: 'end_of_arm.wrist_yaw', method: 'move_to', kwargs: { x_r: 0.0 } })
      // PWM -400 closes the fingers onto their stop: squeezed shut, the -100 % end of range_t (firmware.ts GRIPPER_PCT)
      case 'gripper': return this.call(fw, { target: 'end_of_arm.stretch_gripper', method: 'move_to',
                                            kwargs: { pct: -100.0, v_r: HOME_PWM.gripper.v, a_r: HOME_PWM.gripper.a } })
      // StretchGripper.move_to(pct=0) at the default profile (6 rad/s, 10 rad/s^2, RP SE3_stretch_gripper_SG3)
      case 'gripper_end': return this.call(fw, { target: 'end_of_arm.stretch_gripper', method: 'move_to', kwargs: { pct: 0.0 } })
      default: return
    }
  }

  private finish(): void {
    const ok = this.found.lift && this.found.arm && this.found.wrist_yaw && this.found.gripper
    this.homed = this.homed || ok
    this.lastResult = ok
    this.abort()
  }

  private call(fw: Firmware, c: Command): void {
    this.issued.set(c.target, c)
    fw.apply([c])
  }
}

/** the axis's setpoint has reached its goal and stopped there (a Dynamixel's profile is over, the moving status the original
 *  polls); `near` also asks the joint to be within its distance_tol of it (a stepper's near_pos_setpoint, LIMITS.NEAR) */
function arrived(fw: Firmware, axis: AxisName, near: boolean): boolean {
  const s = fw.state().axes[axis]
  if (s.mode !== 'pos' || s.u !== 0.0 || s.s !== s.goal) return false
  if (!near) return true
  const q = fw.readout()[axis][0]
  return Math.abs(q - s.goal) < LIMITS.NEAR[axis]
}

/** a Dynamixel's Moving flag: the joint turns faster than the servo's moving threshold */
function moving(fw: Firmware, axis: 'wrist_yaw' | 'gripper'): boolean {
  return Math.abs(fw.readout()[axis][1]) > MOVING[axis]
}