"""Electric Slide commands for the official physics-based walking policy."""

from __future__ import annotations

from dataclasses import dataclass
import math
from pathlib import Path

import numpy as np

from choreography import BEAT_SECONDS, COUNTS_PER_WALL, SHOW_BEATS, SHOW_SECONDS


@dataclass(frozen=True)
class PolicyCommand:
    twist: tuple[float, float, float]
    head_pose: tuple[float, float, float, float]
    body_pose: tuple[float, float, float, float, float, float]


# Body-frame velocity commands for each count: forward, lateral, yaw rate.
# Pollen's current simulator deliberately leaves the lateral slot at zero,
# and this published policy does not produce a backward gait.  The dance
# therefore uses supported forward footfalls plus mirrored arcs for its
# right/left phrases.  The second half bends through one physical wall turn.
COUNT_TWISTS = (
    (0.22, 0.00, 0.65),
    (0.22, 0.00, 0.65),
    (0.22, 0.00, 0.65),
    (0.00, 0.00, 0.00),
    (0.22, 0.00, -0.65),
    (0.22, 0.00, -0.65),
    (0.22, 0.00, -0.65),
    (0.00, 0.00, 0.00),
    (0.18, 0.00, 1.00),
    (0.18, 0.00, 1.00),
    (0.18, 0.00, 1.00),
    (0.00, 0.00, 1.00),
    (0.00, 0.00, 0.00),
    (0.14, 0.00, 1.00),
    (0.00, 0.00, 0.00),
    (0.10, 0.00, 1.00),
    (0.10, 0.00, 1.00),
    (0.10, 0.00, 1.00),
)

MOTION_TIMELINE = (
    Path(__file__).resolve().parent
    / "generated/learned_motion_retarget_50hz.npz"
)
_TIMELINE_CACHE: dict[str, np.ndarray] | None = None


def _timeline() -> dict[str, np.ndarray] | None:
    global _TIMELINE_CACHE
    if _TIMELINE_CACHE is None and MOTION_TIMELINE.exists():
        data = np.load(MOTION_TIMELINE, allow_pickle=True)
        _TIMELINE_CACHE = {name: data[name] for name in data.files}
    return _TIMELINE_CACHE


def _motion_sample(time_seconds: float) -> tuple[np.ndarray, int] | None:
    data = _timeline()
    if data is None:
        return None
    commands = data["command_13d"]
    control_hz = float(data["control_hz"])
    phase = (time_seconds % SHOW_SECONDS) * control_hz
    lower = int(math.floor(phase)) % len(commands)
    upper = (lower + 1) % len(commands)
    blend = phase - math.floor(phase)
    command = commands[lower] * (1.0 - blend) + commands[upper] * blend
    policy_code = int(data["policy_code"][lower])
    return command, policy_code


def _smoothstep(value: float) -> float:
    value = max(0.0, min(1.0, value))
    return value * value * (3.0 - 2.0 * value)


def _blended_twist(beat: float) -> tuple[float, float, float]:
    count = int(math.floor(beat)) % COUNTS_PER_WALL
    phase = beat - math.floor(beat)
    current = COUNT_TWISTS[count]
    previous = COUNT_TWISTS[(count - 1) % COUNTS_PER_WALL]
    # Do not carry the prior wall's turn into the next grapevine.
    if count == 0:
        previous = current
    blend = _smoothstep(phase / 0.10)
    return tuple(previous[i] + (current[i] - previous[i]) * blend for i in range(3))


def command_for_time(time_seconds: float, duck_index: int = 0) -> PolicyCommand:
    """Return the 61D-policy command blocks for one show instant."""
    motion = _motion_sample(time_seconds)
    if motion is not None:
        command, _ = motion
        return PolicyCommand(
            twist=tuple(float(value) for value in command[0:3]),
            head_pose=tuple(float(value) for value in command[3:7]),
            body_pose=tuple(float(value) for value in command[7:13]),
        )
    beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
    count = int(math.floor(beat)) % COUNTS_PER_WALL
    phase = beat - math.floor(beat)
    pulse = math.sin(math.pi * phase)
    # Keep the physics commands exactly synchronized.  Even small pose-command
    # phase offsets can amplify into formation drift through the learned gait.
    flock_phase = 0.0
    groove = math.sin(2.0 * math.pi * (beat / 2.0 + flock_phase))

    twist = _blended_twist(beat)

    # Head attitude stays within the command ranges used in training.
    neck_pitch = 0.05 * groove
    head_pitch = -0.05 * groove
    head_yaw = 0.10 * math.sin(2.0 * math.pi * (beat / 4.0 + flock_phase))
    head_roll = 0.09 * groove

    body_z = -0.004 * abs(groove)
    body_roll = 0.035 * groove
    body_pitch = 0.018 * math.sin(2.0 * math.pi * beat / 4.0)

    # Character section: grounded scoop, relaxed lean, then head leads turn.
    if count in (12, 13):
        amount = _smoothstep(phase) if count == 12 else pulse
        body_z -= 0.010 * amount
        body_pitch += 0.035 * amount
        neck_pitch += 0.08 * amount
        head_pitch -= 0.06 * amount
        head_yaw += 0.08 * amount
    elif count in (14, 15):
        amount = _smoothstep(phase) if count == 14 else pulse
        body_pitch -= 0.035 * amount
        neck_pitch -= 0.06 * amount
        head_pitch += 0.05 * amount
    elif count in (16, 17):
        head_yaw += 0.20 * pulse
        head_roll += 0.04 * groove

    return PolicyCommand(
        twist=twist,
        head_pose=(neck_pitch, head_pitch, head_yaw, head_roll),
        body_pose=(0.0, 0.0, body_z, body_roll, body_pitch, 0.0),
    )


def policy_for_time(time_seconds: float) -> str:
    """Choose official one-shot kick policies for the two accent counts."""
    motion = _motion_sample(time_seconds)
    if motion is not None:
        _, policy_code = motion
        return ("walk", "kick_left", "kick_right")[policy_code]
    beat = (time_seconds / BEAT_SECONDS) % SHOW_BEATS
    count = int(math.floor(beat)) % COUNTS_PER_WALL
    if count == 12:
        return "kick_left"
    if count == 14:
        return "kick_right"
    return "walk"


def integrated_turn_radians(samples: int = 10000) -> float:
    """Numerically integrate one phrase's yaw command for a regression test."""
    duration = COUNTS_PER_WALL * BEAT_SECONDS
    dt = duration / samples
    return sum(command_for_time((index + 0.5) * dt).twist[2] * dt for index in range(samples))


__all__ = [
    "COUNT_TWISTS",
    "PolicyCommand",
    "SHOW_SECONDS",
    "command_for_time",
    "integrated_turn_radians",
    "policy_for_time",
]