saltylab-firmware/ui/social-bot/src/components/ImuPanel.jsx

/**
 * ImuPanel.jsx — IMU attitude visualization.
 *
 * Topics:
 *   /saltybot/imu           (sensor_msgs/Imu)         — quaternion orientation
 *   /saltybot/balance_state (std_msgs/String JSON)    — pitch/roll/yaw deg,
 *                                                        motor_cmd, state, mode
 *
 * Displays:
 *   - Artificial horizon canvas (pitch / roll)
 *   - Compass tape (yaw)
 *   - Three.js 3D robot orientation cube
 *   - Numeric readouts + angular velocity
 */

import { useEffect, useRef, useState, useCallback } from 'react';
import * as THREE from 'three';

// ── 2D canvas helpers ──────────────────────────────────────────────────────────

function drawHorizon(ctx, W, H, pitch, roll) {
  const cx = W / 2, cy = H / 2;
  const rollRad = roll * Math.PI / 180;
  const pitchPx = pitch * (H / 60);

  ctx.clearRect(0, 0, W, H);
  ctx.fillStyle = '#051a30';
  ctx.fillRect(0, 0, W, H);

  ctx.save();
  ctx.translate(cx, cy);
  ctx.rotate(-rollRad);

  // Ground
  ctx.fillStyle = '#1a0f00';
  ctx.fillRect(-W, pitchPx, W * 2, H * 2);

  // Horizon
  ctx.strokeStyle = '#00ffff';
  ctx.lineWidth = 1.5;
  ctx.beginPath();
  ctx.moveTo(-W, pitchPx);
  ctx.lineTo(W, pitchPx);
  ctx.stroke();

  // Pitch ladder
  for (let d = -30; d <= 30; d += 10) {
    if (d === 0) continue;
    const y = pitchPx + d * (H / 60);
    const lw = Math.abs(d) % 20 === 0 ? 22 : 14;
    ctx.strokeStyle = 'rgba(0,210,210,0.4)';
    ctx.lineWidth = 0.7;
    ctx.beginPath();
    ctx.moveTo(-lw, y); ctx.lineTo(lw, y);
    ctx.stroke();
    ctx.fillStyle = 'rgba(0,210,210,0.5)';
    ctx.font = '7px monospace';
    ctx.textAlign = 'left';
    ctx.fillText((-d).toString(), lw + 2, y + 3);
  }
  ctx.restore();

  // Reticle
  ctx.strokeStyle = '#f97316';
  ctx.lineWidth = 1.5;
  ctx.beginPath();
  ctx.moveTo(cx - 28, cy); ctx.lineTo(cx - 8, cy);
  ctx.moveTo(cx + 8, cy);  ctx.lineTo(cx + 28, cy);
  ctx.moveTo(cx, cy - 4);  ctx.lineTo(cx, cy + 4);
  ctx.stroke();
}

function drawCompass(ctx, W, H, yaw) {
  const cx = W / 2;
  ctx.clearRect(0, 0, W, H);
  ctx.fillStyle = '#050510';
  ctx.fillRect(0, 0, W, H);

  const degPerPx = W / 70;
  const cardinals = { 0:'N', 45:'NE', 90:'E', 135:'SE', 180:'S', 225:'SW', 270:'W', 315:'NW' };

  for (let i = -35; i <= 35; i++) {
    const deg = ((Math.round(yaw) + i) % 360 + 360) % 360;
    const x = cx + i * degPerPx;
    const isMaj = deg % 45 === 0;
    const isMed = deg % 15 === 0;
    if (!isMed && !isMaj) continue;

    ctx.strokeStyle = isMaj ? '#00cccc' : 'rgba(0,200,200,0.3)';
    ctx.lineWidth = isMaj ? 1.5 : 0.5;
    const tH = isMaj ? 16 : 7;
    ctx.beginPath();
    ctx.moveTo(x, 0); ctx.lineTo(x, tH);
    ctx.stroke();

    if (isMaj && cardinals[deg] !== undefined) {
      ctx.fillStyle = deg === 0 ? '#ff4444' : '#00cccc';
      ctx.font = 'bold 9px monospace';
      ctx.textAlign = 'center';
      ctx.fillText(cardinals[deg], x, 28);
    }
  }

  const hdg = ((Math.round(yaw) % 360) + 360) % 360;
  ctx.fillStyle = '#00ffff';
  ctx.font = 'bold 11px monospace';
  ctx.textAlign = 'center';
  ctx.fillText(hdg + '°', cx, H - 4);

  // Pointer
  ctx.strokeStyle = '#f97316';
  ctx.lineWidth = 2;
  ctx.beginPath();
  ctx.moveTo(cx, 0); ctx.lineTo(cx, 10);
  ctx.stroke();
}

// ── Three.js hook ──────────────────────────────────────────────────────────────

function useThreeOrientation(containerRef) {
  const sceneRef = useRef(null);

  useEffect(() => {
    const el = containerRef.current;
    if (!el) return;

    const scene    = new THREE.Scene();
    const camera   = new THREE.PerspectiveCamera(45, el.clientWidth / el.clientHeight, 0.1, 100);
    camera.position.set(2.5, 2, 3);
    camera.lookAt(0, 0, 0);

    const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: true });
    renderer.setClearColor(0x070712, 1);
    renderer.setPixelRatio(window.devicePixelRatio);
    renderer.setSize(el.clientWidth, el.clientHeight);
    el.appendChild(renderer.domElement);

    // Lighting
    scene.add(new THREE.AmbientLight(0x404060, 3));
    const dir = new THREE.DirectionalLight(0xffffff, 4);
    dir.position.set(5, 8, 6);
    scene.add(dir);

    // Robot body group
    const group = new THREE.Group();

    // Body
    const body = new THREE.Mesh(
      new THREE.BoxGeometry(0.8, 1.4, 0.4),
      new THREE.MeshPhongMaterial({ color: 0x1a2a4a, specular: 0x334466 })
    );
    body.position.y = 0.5;
    group.add(body);

    // Wheels
    const wheelGeo = new THREE.CylinderGeometry(0.35, 0.35, 0.12, 18);
    const wheelMat = new THREE.MeshPhongMaterial({ color: 0x1a1a1a, specular: 0x333333 });
    [-0.5, 0.5].forEach(x => {
      const pivot = new THREE.Group();
      pivot.position.set(x, -0.2, 0);
      pivot.rotation.z = Math.PI / 2;
      pivot.add(new THREE.Mesh(wheelGeo, wheelMat));
      group.add(pivot);

      const rim = new THREE.Mesh(
        new THREE.TorusGeometry(0.3, 0.02, 8, 16),
        new THREE.MeshPhongMaterial({ color: 0x0055cc })
      );
      rim.rotation.z = Math.PI / 2;
      rim.position.set(x * 1.07, -0.2, 0);
      group.add(rim);
    });

    // Forward indicator arrow (red tip at +Z front)
    const arrow = new THREE.Mesh(
      new THREE.ConeGeometry(0.08, 0.25, 8),
      new THREE.MeshBasicMaterial({ color: 0xff3030 })
    );
    arrow.rotation.x = -Math.PI / 2;
    arrow.position.set(0, 0.5, 0.35);
    group.add(arrow);

    // Sensor head
    const head = new THREE.Mesh(
      new THREE.BoxGeometry(0.32, 0.18, 0.32),
      new THREE.MeshPhongMaterial({ color: 0x111122 })
    );
    head.position.set(0, 1.35, 0);
    group.add(head);

    // Axis helper
    scene.add(new THREE.AxesHelper(1.4));
    scene.add(group);

    const q = new THREE.Quaternion();
    let curQ = new THREE.Quaternion();

    sceneRef.current = { group, q, curQ };

    let animId;
    const animate = () => {
      animId = requestAnimationFrame(animate);
      curQ.slerp(q, 0.12);
      group.quaternion.copy(curQ);
      renderer.render(scene, camera);
    };
    animate();

    const ro = new ResizeObserver(() => {
      camera.aspect = el.clientWidth / el.clientHeight;
      camera.updateProjectionMatrix();
      renderer.setSize(el.clientWidth, el.clientHeight);
    });
    ro.observe(el);

    return () => {
      cancelAnimationFrame(animId);
      ro.disconnect();
      renderer.dispose();
      el.removeChild(renderer.domElement);
      sceneRef.current = null;
    };
  }, [containerRef]);

  const updateOrientation = useCallback((qx, qy, qz, qw) => {
    if (!sceneRef.current) return;
    sceneRef.current.q.set(qx, qy, qz, qw);
  }, []);

  return updateOrientation;
}

// ── Component ──────────────────────────────────────────────────────────────────

function Readout({ label, value, unit, warn }) {
  return (
    <div className="bg-gray-950 rounded p-2 text-center">
      <div className="text-gray-600 text-xs">{label}</div>
      <div className={`text-lg font-bold ${warn ? 'text-amber-400' : 'text-cyan-400'}`}>{value}</div>
      <div className="text-gray-700 text-xs">{unit}</div>
    </div>
  );
}

export function ImuPanel({ subscribe }) {
  const horizonRef = useRef(null);
  const compassRef = useRef(null);
  const threeRef   = useRef(null);

  const [attitude, setAttitude] = useState({ pitch: 0, roll: 0, yaw: 0 });
  const [angVel, setAngVel]     = useState({ x: 0, y: 0, z: 0 });
  const [balState, setBalState] = useState(null);
  const [pktHz, setPktHz]       = useState(0);

  const threeUpdate = useThreeOrientation(threeRef);
  const pktRef = useRef({ count: 0, last: Date.now() });

  // Draw canvases on attitude change
  useEffect(() => {
    const hc = horizonRef.current;
    const cc = compassRef.current;
    if (hc) {
      const ctx = hc.getContext('2d');
      drawHorizon(ctx, hc.width, hc.height, attitude.pitch, attitude.roll);
    }
    if (cc) {
      const ctx = cc.getContext('2d');
      drawCompass(ctx, cc.width, cc.height, attitude.yaw);
    }
  }, [attitude]);

  // Subscribe to Imu
  useEffect(() => {
    const unsub = subscribe('/saltybot/imu', 'sensor_msgs/Imu', (msg) => {
      const o = msg.orientation;
      if (o) {
        // Convert quaternion to Euler (simple ZYX)
        const { x, y, z, w } = o;
        const pitchRad = Math.asin(2 * (w * y - z * x));
        const rollRad  = Math.atan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y));
        const yawRad   = Math.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z));
        setAttitude({
          pitch: pitchRad * 180 / Math.PI,
          roll:  rollRad  * 180 / Math.PI,
          yaw:   yawRad   * 180 / Math.PI,
        });
        threeUpdate(x, y, z, w);
      }
      const av = msg.angular_velocity;
      if (av) {
        setAngVel({
          x: av.x * 180 / Math.PI,
          y: av.y * 180 / Math.PI,
          z: av.z * 180 / Math.PI,
        });
      }
      // Hz counter
      pktRef.current.count++;
      const now = Date.now();
      if (now - pktRef.current.last >= 1000) {
        setPktHz(pktRef.current.count);
        pktRef.current.count = 0;
        pktRef.current.last = now;
      }
    });
    return unsub;
  }, [subscribe, threeUpdate]);

  // Subscribe to balance_state for motor/state info
  useEffect(() => {
    const unsub = subscribe('/saltybot/balance_state', 'std_msgs/String', (msg) => {
      try {
        setBalState(JSON.parse(msg.data));
      } catch { /* ignore */ }
    });
    return unsub;
  }, [subscribe]);

  const pitchWarn = Math.abs(attitude.pitch) > 20;
  const rollWarn  = Math.abs(attitude.roll) > 20;

  return (
    <div className="space-y-4">
      {/* Numeric readouts */}
      <div className="grid grid-cols-3 gap-2">
        <Readout label="PITCH" value={attitude.pitch.toFixed(1)} unit="°" warn={pitchWarn} />
        <Readout label="ROLL"  value={attitude.roll.toFixed(1)}  unit="°" warn={rollWarn} />
        <Readout label="YAW"   value={((attitude.yaw + 360) % 360).toFixed(1)} unit="°" />
      </div>

      {/* Angular velocity */}
      <div className="grid grid-cols-3 gap-2">
        <Readout label="ω PITCH" value={angVel.x.toFixed(1)} unit="°/s" />
        <Readout label="ω ROLL"  value={angVel.y.toFixed(1)} unit="°/s" />
        <Readout label="ω YAW"   value={angVel.z.toFixed(1)} unit="°/s" />
      </div>

      {/* 2D gauges */}
      <div className="grid grid-cols-1 sm:grid-cols-2 gap-4">
        <div className="bg-gray-950 rounded-lg border border-cyan-950 p-3">
          <div className="text-cyan-700 text-xs font-bold tracking-widest mb-2">ARTIFICIAL HORIZON</div>
          <canvas
            ref={horizonRef}
            width={280}
            height={120}
            className="w-full rounded border border-gray-900 block"
          />
        </div>
        <div className="bg-gray-950 rounded-lg border border-cyan-950 p-3">
          <div className="text-cyan-700 text-xs font-bold tracking-widest mb-2">COMPASS</div>
          <canvas
            ref={compassRef}
            width={280}
            height={56}
            className="w-full rounded border border-gray-900 block"
          />
          {/* Balance state info */}
          {balState && (
            <div className="mt-2 grid grid-cols-2 gap-1 text-xs">
              <div><span className="text-gray-600">State: </span><span className="text-cyan-400">{balState.state}</span></div>
              <div><span className="text-gray-600">Mode: </span><span className="text-amber-400">{balState.mode}</span></div>
              <div><span className="text-gray-600">Motor: </span><span className="text-orange-400">{balState.motor_cmd}</span></div>
              <div><span className="text-gray-600">Error: </span><span className="text-red-400">{balState.pid_error_deg}°</span></div>
            </div>
          )}
        </div>
      </div>

      {/* Three.js 3D orientation */}
      <div className="bg-gray-950 rounded-lg border border-cyan-950 p-3">
        <div className="flex items-center justify-between mb-2">
          <div className="text-cyan-700 text-xs font-bold tracking-widest">3D ORIENTATION</div>
          <div className="text-gray-600 text-xs">{pktHz} Hz</div>
        </div>
        <div ref={threeRef} className="w-full rounded border border-gray-900" style={{ height: '200px' }} />
      </div>
    </div>
  );
}