"""Policy observation assembly (``PolicyCfg`` of the Online-CLoSD ball task).

Order and semantics (N = number of actuated joints):

    task_command 9 | ref joint_pos N | ref joint_vel N | root xy 2 | ref pelvis xy 2 |
    ref pelvis 6-D 6 | projected gravity 3 | heading xy 2 | base ang vel (body) 3 |
    joint_pos - default N | joint_vel N | previous training-clipped action N -> 27 + 5N

Noise is not added (deployment = clean sensors); enable ``noise=True`` for a
training-like corruption.
"""

from __future__ import annotations

import numpy as np

from .mathutil import quat_apply, quat_apply_inverse, rot6d_interleaved

_TASK_NOISE = np.array([0.005, 0.005, 0.008, 0.008, 0.05, 0.05, 0.05, 0.03, 0.03])


def build_observation(
    *,
    task_command: np.ndarray,
    ref_joint_pos: np.ndarray,
    ref_joint_vel: np.ndarray,
    root_xy: np.ndarray,
    ref_pelvis_xy: np.ndarray,
    ref_pelvis_quat: np.ndarray,
    pelvis_quat: np.ndarray,
    heading_quat: np.ndarray | None,
    base_ang_vel_body: np.ndarray,
    joint_pos: np.ndarray,
    joint_vel: np.ndarray,
    default_joint_pos: np.ndarray,
    prev_action: np.ndarray,
    noise: bool = False,
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    gravity = quat_apply_inverse(pelvis_quat, np.array([0.0, 0.0, -1.0]))
    # Gravity and angular velocity come from the robot IMU.  Heading can use
    # the sign-normalised mocap quaternion to avoid unobservable IMU yaw drift.
    heading_source = pelvis_quat if heading_quat is None else heading_quat
    heading = quat_apply(heading_source, np.array([1.0, 0.0, 0.0]))[:2]
    heading = heading / max(float(np.linalg.norm(heading)), 1e-6)
    parts = [
        np.asarray(task_command, dtype=np.float64),
        np.asarray(ref_joint_pos, dtype=np.float64),
        np.asarray(ref_joint_vel, dtype=np.float64),
        np.asarray(root_xy, dtype=np.float64),
        np.asarray(ref_pelvis_xy, dtype=np.float64),
        rot6d_interleaved(np.asarray(ref_pelvis_quat, dtype=np.float64)),
        gravity,
        heading,
        np.asarray(base_ang_vel_body, dtype=np.float64),
        np.asarray(joint_pos, dtype=np.float64) - np.asarray(default_joint_pos, dtype=np.float64),
        np.asarray(joint_vel, dtype=np.float64),
        np.asarray(prev_action, dtype=np.float64),
    ]
    if noise:
        rng = rng or np.random.default_rng()
        n = len(joint_pos)
        parts[0] = parts[0] + rng.uniform(-_TASK_NOISE, _TASK_NOISE)
        parts[3] = parts[3] + rng.uniform(-0.05, 0.05, 2)
        parts[6] = parts[6] + rng.uniform(-0.03, 0.03, 3)
        parts[7] = parts[7] + rng.uniform(-0.04, 0.04, 2)
        parts[8] = parts[8] + rng.uniform(-0.2, 0.2, 3)
        parts[9] = parts[9] + rng.uniform(-0.01, 0.01, n)
        parts[10] = parts[10] + rng.uniform(-1.0, 1.0, n)
    return np.concatenate(parts).astype(np.float32)


def expected_obs_dim(num_actions: int) -> int:
    return 27 + 5 * num_actions