"""
Multimodal sensory fusion.

Combines vision, touch, IMU, and proprioception into a unified
PerceptionOutput that the brain feeds to the world model and policy.

All processing here is purely NumPy so it runs on the brain thread
without GPU contention.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from env.humanoid_env import RobotObs


@dataclass
class PerceptionOutput:
    # Normalised inputs for world model ingestion
    frames: np.ndarray          # (T, H, W, 3) float32 [0,1]
    touch_seq: np.ndarray       # (T, N_touch) float32 normalised
    imu_seq: np.ndarray         # (T, 6)       float32 normalised
    proprio_seq: np.ndarray     # (T, proprio_dim) float32 normalised

    # Compact summary vectors (for lightweight policy/language conditioning)
    touch_now: np.ndarray       # (N_touch,) latest touch frame
    imu_now: np.ndarray         # (6,)       latest IMU frame
    proprio_now: np.ndarray     # (proprio_dim,)


# ── normalisation statistics (approximate, will drift — use online stats in prod)
_TOUCH_MAX   = 500.0    # typical max contact force [N] in sim
_ACCEL_MAX   = 30.0     # g-ish
_GYRO_MAX    = 10.0     # rad/s
_QPOS_CLIP   = 5.0      # joint position clip
_QVEL_CLIP   = 20.0     # joint velocity clip


class PerceptionModule:
    """
    Stateless fusion of a list of RobotObs into a PerceptionOutput.

    The most recent frame with a non-None vision field is propagated
    forward to fill in frames where the camera was skipped (vision runs
    at a fraction of the control rate).
    """

    def process(self, obs_history: list[RobotObs]) -> PerceptionOutput:
        if not obs_history:
            raise ValueError("obs_history must be non-empty")

        # ── vision: fill forward from last rendered frame
        frames = self._build_frame_sequence(obs_history)

        # ── touch
        touch_seq = np.stack(
            [o.touch / _TOUCH_MAX for o in obs_history], axis=0
        ).clip(-1, 1).astype(np.float32)

        # ── IMU
        imu_seq = np.stack([
            np.concatenate([
                o.imu[:3] / _ACCEL_MAX,
                o.imu[3:] / _GYRO_MAX,
            ]) for o in obs_history
        ], axis=0).clip(-1, 1).astype(np.float32)

        # ── proprioception (qpos | qvel)
        proprio_seq = np.stack([
            self._normalise_proprio(o.proprio) for o in obs_history
        ], axis=0).astype(np.float32)

        latest = obs_history[-1]
        return PerceptionOutput(
            frames=frames,
            touch_seq=touch_seq,
            imu_seq=imu_seq,
            proprio_seq=proprio_seq,
            touch_now=touch_seq[-1],
            imu_now=imu_seq[-1],
            proprio_now=proprio_seq[-1],
        )

    # ── helpers

    def _build_frame_sequence(self, obs_history: list[RobotObs]) -> np.ndarray:
        T = len(obs_history)
        last_frame = None
        for o in reversed(obs_history):
            if o.vision is not None:
                last_frame = o.vision
                break

        if last_frame is None:
            H, W = 224, 224
            last_frame = np.zeros((H, W, 3), dtype=np.uint8)

        H, W, C = last_frame.shape
        frames = np.zeros((T, H, W, C), dtype=np.float32)
        carry = last_frame.astype(np.float32) / 255.0

        # Walk backwards, filling None slots with the most-recent known frame
        frame_list = []
        fill = carry
        for o in reversed(obs_history):
            if o.vision is not None:
                fill = o.vision.astype(np.float32) / 255.0
            frame_list.append(fill)
        frame_list.reverse()
        return np.stack(frame_list, axis=0)   # (T, H, W, 3)

    @staticmethod
    def _normalise_proprio(proprio: np.ndarray) -> np.ndarray:
        # Split into qpos (first nq) and qvel (last nv).
        # Root joint (free): qpos[0:7] (xyz + quaternion), qvel[0:6] (vel + angvel)
        # Rest: joint angles + velocities → clip and normalise.
        p = proprio.copy()
        p[:3] = np.clip(p[:3], -3.0, 3.0) / 3.0            # root xyz position
        p[3:7] = p[3:7]                                       # quaternion already unit
        p[7:] = np.clip(p[7:], -_QPOS_CLIP, _QPOS_CLIP) / _QPOS_CLIP
        return p