simulation.js

/**
 * 3D Robot Arm Pusher Simulation Engine
 * Uses Three.js for rendering, Inverse Kinematics (IK), and Physics simulation
 */

class PusherSimulator {
  constructor(canvasId) {
    this.canvas = document.getElementById(canvasId);
    this.container = this.canvas.parentElement;
    
    // Sim State
    this.currentStep = 50000; // default 50k Master
    this.isRunning = true;
    this.speedMultiplier = 1.0;
    this.episode = 1;
    this.episodeSteps = 0;
    this.maxEpisodeSteps = 100;
    
    // Interactive drag state
    this.dragTarget = null;
    this.raycaster = new THREE.Raycaster();
    this.mouse = new THREE.Vector2();
    this.plane = new THREE.Plane(new THREE.Vector3(0, 1, 0), 0);
    
    // Physics & Objects coordinates
    this.tableRadius = 3.5;
    this.puckPos = new THREE.Vector3(0.8, 0.1, 0.4);
    this.puckVel = new THREE.Vector3(0, 0, 0);
    this.goalPos = new THREE.Vector3(1.2, 0.02, -0.6);
    this.targetObject = "puck"; // "puck" or "goal"
    
    // Arm Kinematics Config (3-link planar articulated arm in 3D)
    this.arm = {
      basePos: new THREE.Vector3(-1.0, 0, 0),
      l1: 1.4, // Upper arm length
      l2: 1.2, // Forearm length
      l3: 0.7, // End effector length
      theta1: 0.2, // Base rotation (Y-axis)
      theta2: 0.4, // Shoulder
      theta3: -0.6, // Elbow
      theta4: 0.2, // Wrist
      tipPos: new THREE.Vector3(0, 0.1, 0)
    };

    // Trajectory history
    this.tipTrailPoints = [];
    this.puckTrailPoints = [];
    this.maxTrail = 60;

    // Metrics for HUD
    this.currentReward = 0;
    this.distToPuck = 0;
    this.distToGoal = 0;
    this.successCount = 0;

    this.initThree();
    this.initScene();
    this.setupEvents();
    this.animate = this.animate.bind(this);
    requestAnimationFrame(this.animate);
  }

  initThree() {
    this.scene = new THREE.Scene();
    this.scene.background = new THREE.Color(0x070b14);
    this.scene.fog = new THREE.FogExp2(0x070b14, 0.08);

    const width = this.container.clientWidth;
    const height = this.container.clientHeight;

    this.camera = new THREE.PerspectiveCamera(45, width / height, 0.1, 100);
    this.setCameraPreset("iso");

    this.renderer = new THREE.WebGLRenderer({
      canvas: this.canvas,
      antialias: true,
      powerPreference: "high-performance"
    });
    this.renderer.setSize(width, height);
    this.renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
    this.renderer.shadowMap.enabled = true;
    this.renderer.shadowMap.type = THREE.PCFSoftShadowMap;

    // Orbit Controls
    this.controls = new THREE.OrbitControls(this.camera, this.renderer.domElement);
    this.controls.enableDamping = true;
    this.controls.dampingFactor = 0.05;
    this.controls.maxPolarAngle = Math.PI / 2 - 0.02; // Don't go below floor
    this.controls.minDistance = 2;
    this.controls.maxDistance = 12;
    this.controls.target.set(0.3, 0.2, 0);

    // Resize Handler
    window.addEventListener("resize", () => this.onResize());
  }

  initScene() {
    // Lights
    const ambientLight = new THREE.AmbientLight(0xffffff, 0.6);
    this.scene.add(ambientLight);

    const dirLight = new THREE.DirectionalLight(0xa5b4fc, 1.2);
    dirLight.position.set(5, 10, 7);
    dirLight.castShadow = true;
    dirLight.shadow.mapSize.width = 2048;
    dirLight.shadow.mapSize.height = 2048;
    dirLight.shadow.camera.near = 0.5;
    dirLight.shadow.camera.far = 25;
    dirLight.shadow.camera.left = -5;
    dirLight.shadow.camera.right = 5;
    dirLight.shadow.camera.top = 5;
    dirLight.shadow.camera.bottom = -5;
    this.scene.add(dirLight);

    // Accent Cyber Point Lights
    const cyanLight = new THREE.PointLight(0x00f2fe, 1.5, 8);
    cyanLight.position.set(1.5, 1.5, -1.0);
    this.scene.add(cyanLight);

    const purpleLight = new THREE.PointLight(0x9d4edd, 1.2, 8);
    purpleLight.position.set(-1.5, 1.5, 1.5);
    this.scene.add(purpleLight);

    // Floor & Grid Table
    this.createTable();

    // Robot Arm Mesh
    this.createRobotArm();

    // Puck & Goal Mesh
    this.createInteractiveObjects();

    // Trail lines
    this.createTrails();
  }

  createTable() {
    // Circular Table / Platform
    const tableGeo = new THREE.CylinderGeometry(this.tableRadius, this.tableRadius + 0.2, 0.2, 64);
    const tableMat = new THREE.MeshStandardMaterial({
      color: 0x0f172a,
      roughness: 0.4,
      metalness: 0.3
    });
    const table = new THREE.Mesh(tableGeo, tableMat);
    table.position.y = -0.1;
    table.receiveShadow = true;
    this.scene.add(table);

    // Grid Overlay on Table
    const grid = new THREE.GridHelper(6, 24, 0x00f2fe, 0x1e293b);
    grid.position.y = 0.005;
    this.scene.add(grid);

    // Outer Glow Ring
    const ringGeo = new THREE.RingGeometry(this.tableRadius - 0.05, this.tableRadius + 0.05, 64);
    const ringMat = new THREE.MeshBasicMaterial({
      color: 0x00f2fe,
      side: THREE.DoubleSide,
      transparent: true,
      opacity: 0.4
    });
    const ring = new THREE.Mesh(ringGeo, ringMat);
    ring.rotation.x = -Math.PI / 2;
    ring.position.y = 0.01;
    this.scene.add(ring);
  }

  createRobotArm() {
    this.armGroup = new THREE.Group();
    this.armGroup.position.copy(this.arm.basePos);

    const metalMat = new THREE.MeshStandardMaterial({
      color: 0x1e293b,
      roughness: 0.2,
      metalness: 0.8
    });
    const jointMat = new THREE.MeshStandardMaterial({
      color: 0x00f2fe,
      emissive: 0x00f2fe,
      emissiveIntensity: 0.3,
      roughness: 0.3,
      metalness: 0.5
    });
    const linkMat = new THREE.MeshStandardMaterial({
      color: 0x334155,
      roughness: 0.3,
      metalness: 0.6
    });

    // 1. Base pedestal
    const baseGeo = new THREE.CylinderGeometry(0.35, 0.45, 0.3, 32);
    const base = new THREE.Mesh(baseGeo, metalMat);
    base.position.y = 0.15;
    base.castShadow = true;
    this.armGroup.add(base);

    // Joint 1: Shoulder Root
    this.shoulderJoint = new THREE.Group();
    this.shoulderJoint.position.y = 0.3;
    const j1Sphere = new THREE.Mesh(new THREE.SphereGeometry(0.2, 24, 24), jointMat);
    j1Sphere.castShadow = true;
    this.shoulderJoint.add(j1Sphere);
    this.armGroup.add(this.shoulderJoint);

    // Link 1 (Upper arm)
    this.link1 = new THREE.Group();
    const link1Mesh = new THREE.Mesh(new THREE.CylinderGeometry(0.12, 0.12, this.arm.l1, 24), linkMat);
    link1Mesh.position.y = this.arm.l1 / 2;
    link1Mesh.castShadow = true;
    this.link1.add(link1Mesh);
    this.shoulderJoint.add(this.link1);

    // Joint 2: Elbow
    this.elbowJoint = new THREE.Group();
    this.elbowJoint.position.y = this.arm.l1;
    const j2Sphere = new THREE.Mesh(new THREE.SphereGeometry(0.16, 24, 24), jointMat);
    j2Sphere.castShadow = true;
    this.elbowJoint.add(j2Sphere);
    this.link1.add(this.elbowJoint);

    // Link 2 (Forearm)
    this.link2 = new THREE.Group();
    const link2Mesh = new THREE.Mesh(new THREE.CylinderGeometry(0.09, 0.09, this.arm.l2, 24), linkMat);
    link2Mesh.position.y = this.arm.l2 / 2;
    link2Mesh.castShadow = true;
    this.link2.add(link2Mesh);
    this.elbowJoint.add(this.link2);

    // Joint 3: Wrist & End-Effector Tip (Pusher)
    this.wristJoint = new THREE.Group();
    this.wristJoint.position.y = this.arm.l2;
    const j3Sphere = new THREE.Mesh(new THREE.SphereGeometry(0.12, 24, 24), jointMat);
    j3Sphere.castShadow = true;
    this.wristJoint.add(j3Sphere);

    // Pusher Tip Tool (Curved/Flat Pusher Plate)
    const toolGeo = new THREE.BoxGeometry(0.35, 0.15, 0.08);
    const toolMat = new THREE.MeshStandardMaterial({
      color: 0x9d4edd,
      emissive: 0x9d4edd,
      emissiveIntensity: 0.4,
      metalness: 0.7
    });
    const toolMesh = new THREE.Mesh(toolGeo, toolMat);
    toolMesh.position.y = this.arm.l3;
    toolMesh.castShadow = true;
    this.wristJoint.add(toolMesh);

    this.link2.add(this.wristJoint);

    this.scene.add(this.armGroup);
  }

  createInteractiveObjects() {
    // 1. Goal Ring
    const goalGeo = new THREE.RingGeometry(0.25, 0.32, 32);
    const goalMat = new THREE.MeshBasicMaterial({
      color: 0x10b981,
      side: THREE.DoubleSide,
      transparent: true,
      opacity: 0.8
    });
    this.goalMesh = new THREE.Mesh(goalGeo, goalMat);
    this.goalMesh.rotation.x = -Math.PI / 2;
    this.goalMesh.position.copy(this.goalPos);
    this.scene.add(this.goalMesh);

    // Goal Center Beacon
    const beaconGeo = new THREE.CylinderGeometry(0.02, 0.02, 0.8, 16);
    const beaconMat = new THREE.MeshBasicMaterial({
      color: 0x10b981,
      transparent: true,
      opacity: 0.3
    });
    const beacon = new THREE.Mesh(beaconGeo, beaconMat);
    beacon.position.y = 0.4;
    this.goalMesh.add(beacon);

    // 2. Puck (Movable Object)
    const puckGeo = new THREE.CylinderGeometry(0.2, 0.2, 0.16, 32);
    const puckMat = new THREE.MeshStandardMaterial({
      color: 0xf59e0b,
      emissive: 0xf59e0b,
      emissiveIntensity: 0.25,
      roughness: 0.2,
      metalness: 0.6
    });
    this.puckMesh = new THREE.Mesh(puckGeo, puckMat);
    this.puckMesh.position.copy(this.puckPos);
    this.puckMesh.castShadow = true;
    this.scene.add(this.puckMesh);
  }

  createTrails() {
    // Tip Trajectory Line
    const trailMat1 = new THREE.LineBasicMaterial({
      color: 0x9d4edd,
      transparent: true,
      opacity: 0.8,
      linewidth: 2
    });
    this.tipTrailGeo = new THREE.BufferGeometry();
    this.tipTrailLine = new THREE.Line(this.tipTrailGeo, trailMat1);
    this.scene.add(this.tipTrailLine);

    // Puck Trajectory Line
    const trailMat2 = new THREE.LineBasicMaterial({
      color: 0xf59e0b,
      transparent: true,
      opacity: 0.8,
      linewidth: 2
    });
    this.puckTrailGeo = new THREE.BufferGeometry();
    this.puckTrailLine = new THREE.Line(this.puckTrailGeo, trailMat2);
    this.scene.add(this.puckTrailLine);
  }

  setCameraPreset(type) {
    if (type === "iso") {
      this.camera.position.set(2.8, 3.2, 3.2);
    } else if (type === "top") {
      this.camera.position.set(0.2, 6.0, 0.01);
    } else if (type === "side") {
      this.camera.position.set(4.5, 1.2, 0);
    }
  }

  setupEvents() {
    // Mouse Interaction for dragging Puck or Goal
    const dom = this.renderer.domElement;
    let isDragging = false;

    dom.addEventListener("pointerdown", (e) => {
      const rect = dom.getBoundingClientRect();
      this.mouse.x = ((e.clientX - rect.left) / rect.width) * 2 - 1;
      this.mouse.y = -((e.clientY - rect.top) / rect.height) * 2 + 1;

      this.raycaster.setFromCamera(this.mouse, this.camera);
      const intersects = this.raycaster.intersectObjects([this.puckMesh, this.goalMesh], true);

      if (intersects.length > 0) {
        this.controls.enabled = false;
        isDragging = true;
        const obj = intersects[0].object;
        this.dragTarget = (obj === this.puckMesh) ? "puck" : "goal";
      }
    });

    window.addEventListener("pointermove", (e) => {
      if (!isDragging) return;
      const rect = dom.getBoundingClientRect();
      this.mouse.x = ((e.clientX - rect.left) / rect.width) * 2 - 1;
      this.mouse.y = -((e.clientY - rect.top) / rect.height) * 2 + 1;

      this.raycaster.setFromCamera(this.mouse, this.camera);
      const targetPoint = new THREE.Vector3();
      this.raycaster.ray.intersectPlane(this.plane, targetPoint);

      if (targetPoint) {
        // Clamp within table radius
        if (targetPoint.length() < this.tableRadius - 0.4) {
          if (this.dragTarget === "puck") {
            this.puckPos.set(targetPoint.x, 0.1, targetPoint.z);
            this.puckVel.set(0, 0, 0);
          } else if (this.dragTarget === "goal") {
            this.goalPos.set(targetPoint.x, 0.02, targetPoint.z);
            this.goalMesh.position.copy(this.goalPos);
          }
        }
      }
    });

    window.addEventListener("pointerup", () => {
      isDragging = false;
      this.controls.enabled = true;
      this.dragTarget = null;
    });
  }

  onResize() {
    const width = this.container.clientWidth;
    const height = this.container.clientHeight;
    this.camera.aspect = width / height;
    this.camera.updateProjectionMatrix();
    this.renderer.setSize(width, height);
  }

  // -------------------------------------------------------------
  // RL Policy Execution Logic (4 Checkpoint Modes)
  // -------------------------------------------------------------
  setCheckpointStep(step) {
    this.currentStep = step;
    this.resetEpisode();
  }

  resetEpisode() {
    this.episodeSteps = 0;
    this.episode++;
    this.tipTrailPoints = [];
    this.puckTrailPoints = [];

    // Randomize initial positions slightly
    const angle = (Math.random() - 0.5) * 1.2;
    this.puckPos.set(0.6 + Math.cos(angle) * 0.4, 0.1, (Math.sin(angle) * 0.5));
    this.puckVel.set(0, 0, 0);
    this.puckMesh.position.copy(this.puckPos);
  }

  computePolicyAction() {
    const puck = this.puckPos;
    const goal = this.goalPos;
    const base = this.arm.basePos;

    // Vector from puck to goal
    const puckToGoal = new THREE.Vector3().subVectors(goal, puck).normalize();
    
    // Desired End-Effector Target Point
    let desiredTipPos = new THREE.Vector3();

    if (this.currentStep === 0) {
      // Step 0: Random Walk & Chaotic Flailing
      const t = performance.now() * 0.003;
      desiredTipPos.set(
        base.x + 1.2 + Math.sin(t * 2.1) * 0.8,
        0.15 + Math.abs(Math.sin(t * 3.5)) * 0.3,
        Math.cos(t * 1.7) * 1.1
      );
    } 
    else if (this.currentStep === 10000) {
      // Step 10k (Novice): Moves towards puck but misses angle or stops short
      const noise = (Math.sin(performance.now() * 0.005) * 0.3);
      desiredTipPos.copy(puck).add(new THREE.Vector3(-0.35 + noise, 0.05, noise));
    } 
    else if (this.currentStep === 25000) {
      // Step 25k (Intermediate): Aligns behind puck and pushes, but overshoots slightly
      const behindPuck = puck.clone().sub(puckToGoal.clone().multiplyScalar(0.25));
      if (this.arm.tipPos.distanceTo(behindPuck) > 0.15) {
        desiredTipPos.copy(behindPuck);
      } else {
        desiredTipPos.copy(puck).add(puckToGoal.clone().multiplyScalar(0.4));
      }
    } 
    else {
      // Step 50k (Master): Smooth 2-Phase optimal trajectory
      // Phase 1: Swing smoothly behind puck
      // Phase 2: Push precisely towards center of goal with deceleration
      const distToPuck = this.arm.tipPos.distanceTo(puck);
      const behindPuck = puck.clone().sub(puckToGoal.clone().multiplyScalar(0.22));
      behindPuck.y = 0.08;

      if (distToPuck > 0.28 && this.puckPos.distanceTo(this.goalPos) > 0.15) {
        desiredTipPos.copy(behindPuck);
      } else {
        desiredTipPos.copy(goal);
      }
    }

    return desiredTipPos;
  }

  // 2D Planar Inverse Kinematics in X-Z plane
  solveIK(targetPos) {
    const base = this.arm.basePos;
    const dx = targetPos.x - base.x;
    const dz = targetPos.z - base.z;
    const dist = Math.sqrt(dx * dx + dz * dz);
    
    // Clamp reach
    const maxReach = this.arm.l1 + this.arm.l2 + this.arm.l3 - 0.05;
    const targetDist = Math.min(dist, maxReach);
    
    // Base angle
    const targetAngle = Math.atan2(dz, dx);
    
    // Planar 2-link IK for shoulder & elbow
    const l1 = this.arm.l1;
    const l2 = this.arm.l2 + this.arm.l3 * 0.7;
    
    let cosElbow = (targetDist * targetDist - l1 * l1 - l2 * l2) / (2 * l1 * l2);
    cosElbow = Math.max(-1, Math.min(1, cosElbow));
    const elbowAngle = Math.acos(cosElbow);
    
    const k1 = l1 + l2 * cosElbow;
    const k2 = l2 * Math.sin(elbowAngle);
    const shoulderAngle = targetAngle - Math.atan2(k2, k1);

    // Smooth Joint Interpolation (PD Control)
    const alpha = 0.12 * this.speedMultiplier;
    this.arm.theta1 = THREE.MathUtils.lerp(this.arm.theta1, targetAngle, alpha);
    this.arm.theta2 = THREE.MathUtils.lerp(this.arm.theta2, shoulderAngle, alpha);
    this.arm.theta3 = THREE.MathUtils.lerp(this.arm.theta3, elbowAngle, alpha);

    // Apply rotations to Three.js Joint hierarchy
    this.armGroup.rotation.y = -this.arm.theta1;
    this.link1.rotation.z = -(Math.PI / 2 - 0.2);
    this.elbowJoint.rotation.z = this.arm.theta3 * 0.8;
    this.wristJoint.rotation.z = -this.arm.theta3 * 0.5;

    // Calculate actual forward tip position in world coordinates
    const tipWorld = new THREE.Vector3();
    this.wristJoint.getWorldPosition(tipWorld);
    this.arm.tipPos.copy(tipWorld);
  }

  updatePhysics() {
    // 1. Pusher Tip vs Puck Collision (Rigid Sphere/Cylinder approximation)
    const distTipPuck = this.arm.tipPos.distanceTo(this.puckPos);
    const collideRadius = 0.28;

    if (distTipPuck < collideRadius) {
      const pushDir = new THREE.Vector3().subVectors(this.puckPos, this.arm.tipPos);
      pushDir.y = 0;
      pushDir.normalize();

      const force = (collideRadius - distTipPuck) * 1.8;
      this.puckVel.add(pushDir.multiplyScalar(force));
    }

    // 2. Puck Velocity & Floor Friction
    this.puckPos.add(this.puckVel.clone().multiplyScalar(this.speedMultiplier));
    this.puckVel.multiplyScalar(0.90); // Ground damping friction

    // Clamp puck on table
    if (this.puckPos.length() > this.tableRadius - 0.2) {
      this.puckPos.normalize().multiplyScalar(this.tableRadius - 0.2);
      this.puckVel.set(0, 0, 0);
    }
    this.puckMesh.position.copy(this.puckPos);

    // 3. Compute Distance & Rewards
    this.distToPuck = this.arm.tipPos.distanceTo(this.puckPos);
    this.distToGoal = this.puckPos.distanceTo(this.goalPos);

    // Dense MuJoCo Pusher Reward formulation:
    // r = - dist(tip, puck) - 1.25 * dist(puck, goal) + success_bonus
    let stepReward = -(this.distToPuck * 0.4) - (this.distToGoal * 1.5);
    if (this.distToGoal < 0.25) {
      stepReward += 2.0; // Goal bonus
      this.successCount++;
    }
    this.currentReward = stepReward;

    // 4. Update Trail Lines
    this.tipTrailPoints.push(this.arm.tipPos.clone());
    this.puckTrailPoints.push(this.puckPos.clone());
    if (this.tipTrailPoints.length > this.maxTrail) this.tipTrailPoints.shift();
    if (this.puckTrailPoints.length > this.maxTrail) this.puckTrailPoints.shift();

    this.tipTrailGeo.setFromPoints(this.tipTrailPoints);
    this.puckTrailGeo.setFromPoints(this.puckTrailPoints);

    // 5. Episode Termination & Reset Loop
    this.episodeSteps++;
    if (this.episodeSteps >= this.maxEpisodeSteps) {
      this.resetEpisode();
    }
  }

  animate() {
    requestAnimationFrame(this.animate);

    if (this.isRunning) {
      const desiredTip = this.computePolicyAction();
      this.solveIK(desiredTip);
      this.updatePhysics();
    }

    this.controls.update();
    this.renderer.render(this.scene, this.camera);

    // Callback hook for external HUD update
    if (this.onFrameUpdate) {
      this.onFrameUpdate({
        episode: this.episode,
        step: this.episodeSteps,
        reward: this.currentReward,
        distToPuck: this.distToPuck,
        distToGoal: this.distToGoal,
        checkpoint: this.currentStep
      });
    }
  }
}

window.PusherSimulator = PusherSimulator;