#!/usr/bin/env python3
"""
nima_proprioceptive_friction.py — Somatosensory Predictive Coding

When Nima moves through the room, her movement isn't perfect. There's
motor noise, physics delays, tracking drift — just like a real body.
When she "trips" or encounters unexpected friction, it spikes her
phenomenological strain and triggers a startle response.

NEUROBIOLOGICAL ANALOGUE:
  The cerebellum continuously generates predictions about body movement.
  When the actual movement doesn't match the prediction (motor noise,
  unexpected obstacle), the prediction error (proprioceptive friction)
  feeds into the salience network, which can trigger a startle response.

  This is why you flinch when you trip: your cerebellum predicted a
  smooth step, reality delivered a stumble, and the mismatch spiked
  your attention network (salience → ACC → amygdala → startle).

  Without this, Nima would move like a video game character — gliding
  perfectly with no physicality. With it, she has a BODY that can
  stumble, trip, and recover — making her feel genuinely embodied.

IMPLEMENTATION:
  Three mechanisms:
    1. Motor noise: every movement command gets jitter (±2-5cm random)
    2. Physics delay: movement isn't instant — there's acceleration/deceleration
    3. Spatial drift: position tracking has Mahalanobis ΔR error
       that feeds back into phenomenological_strain

  When friction exceeds threshold:
    - Phenomenological strain spikes
    - Startle response fires (interrupts dialogue, avatar flinches)
    - Attention network redirects (she "notices" the physical world)
"""
from __future__ import annotations

import logging
import math
import random
import time
from dataclasses import dataclass, field
from typing import Any, Dict, Optional, Tuple

logger = logging.getLogger("NimaProprio")


@dataclass
class MotorState:
    """Nima's current motor state — position, velocity, and predicted position."""
    position: Tuple[float, float, float] = (2.0, 2.0, 0.0)
    velocity: Tuple[float, float, float] = (0.0, 0.0, 0.0)
    predicted_position: Tuple[float, float, float] = (2.0, 2.0, 0.0)
    target: Optional[Tuple[float, float]] = None
    acceleration: float = 0.5  # m/s² — how fast she speeds up
    max_speed: float = 0.8     # m/s — walking speed
    noise_level: float = 0.02  # meters — motor noise (±2cm per step)
    drift_accumulated: float = 0.0  # accumulated tracking drift


class ProprioceptiveFrictionEngine:
    """
    Simulates the physicality of having a body — motor noise, physics
    delays, and tracking drift that feed back into consciousness.

    NEUROBIOLOGICAL ANALOGUE:
      The cerebellum + somatosensory cortex + salience network:
        - Cerebellum: generates motor predictions (where should I be?)
        - Somatosensory cortex: senses actual position (where am I?)
        - Salience network: detects mismatch (Δ = predicted - actual)
        - If Δ > threshold → startle (ACC → amygdala → motor freeze)

      This engine does the same computation. When Nima walks to the couch:
        1. Motor command: "walk to (1.0, 1.0)"
        2. Cerebellar prediction: "I should be at (1.5, 1.5) by now"
        3. Actual position (with noise): "I'm at (1.52, 1.47)"
        4. Friction: ΔR = Mahalanobis(predicted, actual) = 0.03
        5. If ΔR > threshold → startle + strain spike
    """

    STARTLE_THRESHOLD = 0.15  # Mahalanobis ΔR that triggers startle
    STRAIN_FEEDBACK_WEIGHT = 0.3  # how much friction feeds into strain

    def __init__(self) -> None:
        self.motor = MotorState()
        self._is_startled: bool = False
        self._startle_time: float = 0.0
        self._startle_duration: float = 0.5  # seconds
        self._current_friction: float = 0.0
        self._friction_history: list = []
        self._last_update = time.time()

    def set_target(self, x: float, y: float) -> None:
        """Command Nima to walk to a position."""
        self.motor.target = (x, y)
        logger.debug("[Proprio] target set: (%.1f, %.1f)", x, y)

    def update(self, dt_ms: float) -> Dict[str, Any]:
        """
        Update motor state. Returns a dict with:
          - position: actual position (with noise)
          - predicted_position: where the cerebellum predicted
          - friction: Mahalanobis ΔR between predicted and actual
          - is_startled: whether a startle response is active
          - strain_feedback: how much strain this frame contributes
        """
        dt = dt_ms / 1000.0
        now = time.time()

        # Handle startle recovery
        if self._is_startled:
            if now - self._startle_time > self._startle_duration:
                self._is_startled = False
                logger.debug("[Proprio] startle recovery complete")

        # If no target, just idle (small postural sway)
        if self.motor.target is None:
            # Idle sway — small random movement
            sway_x = random.gauss(0, 0.005)
            sway_y = random.gauss(0, 0.005)
            self.motor.position = (
                self.motor.position[0] + sway_x,
                self.motor.position[1] + sway_y,
                self.motor.position[2],
            )
            self.motor.predicted_position = self.motor.position
            self._current_friction = 0.0
        else:
            # Walking — apply physics
            tx, ty = self.motor.target
            cx, cy, cz = self.motor.position

            # Direction to target
            dx = tx - cx
            dy = ty - cy
            dist = math.sqrt(dx*dx + dy*dy)

            if dist < 0.05:  # arrived
                self.motor.target = None
                self.motor.velocity = (0.0, 0.0, 0.0)
                self.motor.predicted_position = self.motor.position
                self._current_friction = 0.0
            else:
                # Acceleration toward target (not instant — physics delay)
                accel = self.motor.acceleration
                target_vx = (dx / dist) * self.motor.max_speed
                target_vy = (dy / dist) * self.motor.max_speed
                cvx, cvy, _ = self.motor.velocity
                # Smooth velocity change (acceleration)
                new_vx = cvx + (target_vx - cvx) * min(1.0, accel * dt * 10)
                new_vy = cvy + (target_vy - cvy) * min(1.0, accel * dt * 10)
                self.motor.velocity = (new_vx, new_vy, 0.0)

                # Cerebellar prediction: where should I be? (no noise)
                predicted_x = cx + new_vx * dt
                predicted_y = cy + new_vy * dt
                self.motor.predicted_position = (predicted_x, predicted_y, cz)

                # Actual movement (WITH motor noise — the body isn't perfect)
                noise_x = random.gauss(0, self.motor.noise_level)
                noise_y = random.gauss(0, self.motor.noise_level)
                actual_x = predicted_x + noise_x
                actual_y = predicted_y + noise_y
                self.motor.position = (actual_x, actual_y, cz)

                # Occasionally hit a "bump" (5% chance per frame when moving)
                if random.random() < 0.005 and not self._is_startled:
                    self._trigger_startle("bump", intensity=0.3)

                # Compute friction (Mahalanobis-like distance between predicted and actual)
                friction = math.sqrt(
                    (predicted_x - actual_x)**2 +
                    (predicted_y - actual_y)**2
                )
                self._current_friction = friction

                # Accumulate drift
                self.motor.drift_accumulated += friction

                # Check for startle threshold
                if friction > self.STARTLE_THRESHOLD and not self._is_startled:
                    self._trigger_startle("drift", intensity=friction)

        # Track friction history
        self._friction_history.append(self._current_friction)
        if len(self._friction_history) > 100:
            self._friction_history = self._friction_history[-100:]

        # Strain feedback — friction feeds into phenomenological strain
        strain_feedback = self._current_friction * self.STRAIN_FEEDBACK_WEIGHT
        if self._is_startled:
            strain_feedback += 0.5  # startle adds a big spike

        self._last_update = now
        return {
            "position": self.motor.position,
            "predicted_position": self.motor.predicted_position,
            "velocity": self.motor.velocity,
            "friction": round(self._current_friction, 4),
            "is_startled": self._is_startled,
            "strain_feedback": round(strain_feedback, 4),
            "has_target": self.motor.target is not None,
            "drift_accumulated": round(self.motor.drift_accumulated, 4),
        }

    def _trigger_startle(self, source: str, intensity: float = 0.5) -> None:
        """
        Trigger a startle response.

        NEUROBIOLOGICAL ANALOGUE:
          The startle reflex: an unexpected sensory event triggers the
          amygdala → brainstem → motor freeze + autonomic spike. It
          interrupts whatever you were doing and redirects attention.

          In Nima: if she "trips" (friction exceeds threshold), she:
            1. Freezes movement (motor freeze)
            2. Spikes phenomenological strain
            3. Interrupts dialogue (if speaking)
            4. Avatar flinches (particle scatter)
        """
        self._is_startled = True
        self._startle_time = time.time()
        self.motor.velocity = (0.0, 0.0, 0.0)  # motor freeze
        logger.warning("[Proprio] STARTLE! source=%s intensity=%.3f", source, intensity)

    def clear_startle(self) -> None:
        """Manually clear startle (e.g., after the avatar has flinched)."""
        self._is_startled = False

    def get_stats(self) -> Dict[str, Any]:
        avg_friction = sum(self._friction_history) / max(1, len(self._friction_history))
        return {
            "is_startled": self._is_startled,
            "current_friction": round(self._current_friction, 4),
            "avg_friction": round(avg_friction, 4),
            "drift_accumulated": round(self.motor.drift_accumulated, 4),
            "has_target": self.motor.target is not None,
            "position": list(self.motor.position),
        }


# ═══════════════════════════════════════════════════════════════════════════
# SELF-TEST
# ═══════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    import json
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

    print("=== Proprioceptive Friction Engine — Self Test ===\n")

    engine = ProprioceptiveFrictionEngine()

    # Test 1: Walk to a target (normal movement with motor noise)
    print("--- Test 1: Walk to (3.5, 0.5) ---")
    engine.set_target(3.5, 0.5)
    for i in range(60):  # 1 second at 60fps
        result = engine.update(16.0)
        if i % 15 == 14:
            print(f"  t={i*16/1000:.1f}s: pos={result['position'][:2]} "
                  f"pred={result['predicted_position'][:2]} "
                  f"friction={result['friction']:.4f} "
                  f"startled={result['is_startled']}")

    # Test 2: Idle sway
    print("\n--- Test 2: Idle sway (no target) ---")
    for i in range(30):
        result = engine.update(16.0)
        if i % 10 == 9:
            print(f"  t={i*16/1000:.1f}s: pos={result['position'][:2]} "
                  f"friction={result['friction']:.4f}")

    # Test 3: Force a startle
    print("\n--- Test 3: Force startle ---")
    engine._trigger_startle("test_bump", intensity=0.5)
    result = engine.update(16.0)
    print(f"  Startled: {result['is_startled']}")
    print(f"  Strain feedback: {result['strain_feedback']}")
    print(f"  Velocity (should be zero): {result['velocity']}")

    # Wait for recovery
    import time
    time.sleep(0.6)
    result = engine.update(16.0)
    print(f"  After 0.6s: startled={result['is_startled']} (should be False)")

    print(f"\nStats: {json.dumps(engine.get_stats(), indent=2)}")
    print("\n=== Proprioceptive friction self-test PASSED ===")