from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np


CONTROL_HZ = 50
OBSERVATION_HZ = 10
CONTROL_DECIMATION = CONTROL_HZ // OBSERVATION_HZ
# Two seconds of 50 Hz control predicted from every 10 Hz observation.
CHUNK_SIZE = 100
EPISODE_DURATION_S = 5.0
TERMINAL_STABILIZATION_S = 1.5
RECOVERY_STABLE_S = 0.5
STATE_DIM = 61
ACTION_DIM = 14
EMBODIMENT_ID = 24

POLICY_REVISION = "916c1e1ee11f8460506532e037a4cf1aa110aca6"
MICRODUCK_RL_REVISION = "d424a0c899f6b33cbd3daeb279913134349c0b63"
TINYVLA_REVISION = "565d139f4960174fdd3006ec57c25ac62e129245"

DEFAULT_POSE = np.array(
    [
        0.0,
        -0.0873,
        -0.4579,
        -0.0049,
        0.4530,
        0.3491,
        0.3491,
        0.0,
        0.0,
        0.0,
        0.0873,
        0.4579,
        0.0049,
        -0.4530,
    ],
    dtype=np.float32,
)


@dataclass(frozen=True)
class BehaviorSpec:
    name: str
    policy_file: str
    scene_file: str
    duration_s: float
    start_keyframe: str
    handoff_policy_file: str | None = None
    handoff_after_s: float | None = None
    phase_period_s: float | None = None
    kind: Literal[
        "stand",
        "walking",
        "sitstand",
        "phase",
        "kick",
        "roulade",
        "roller",
        "recovery",
    ] = "stand"


BEHAVIORS: tuple[BehaviorSpec, ...] = (
    BehaviorSpec(
        "stand",
        "alpha_stand.onnx",
        "scene.xml",
        EPISODE_DURATION_S,
        "STAND",
        kind="stand",
    ),
    BehaviorSpec(
        "walking",
        "alpha_walking.onnx",
        "scene.xml",
        EPISODE_DURATION_S,
        "STAND",
        handoff_policy_file="alpha_stand.onnx",
        handoff_after_s=4.0,
        kind="walking",
    ),
    BehaviorSpec(
        "sitstand",
        "alpha_sitstand.onnx",
        "scene.xml",
        EPISODE_DURATION_S,
        "STAND",
        kind="sitstand",
    ),
    BehaviorSpec(
        "ground_pick",
        "alpha_ground_pick.onnx",
        "scene.xml",
        EPISODE_DURATION_S,
        "STAND",
        handoff_policy_file="alpha_stand.onnx",
        handoff_after_s=2.8,
        phase_period_s=4.0,
        kind="phase",
    ),
    BehaviorSpec(
        "kick_left",
        "ball_kick_left.onnx",
        "scene_ball.xml",
        5.0,
        "STAND",
        handoff_policy_file="alpha_stand.onnx",
        handoff_after_s=3.0,
        kind="kick",
    ),
    BehaviorSpec(
        "kick_right",
        "ball_kick_right.onnx",
        "scene_ball.xml",
        5.0,
        "STAND",
        handoff_policy_file="alpha_stand.onnx",
        handoff_after_s=3.0,
        kind="kick",
    ),
    BehaviorSpec(
        "roller",
        "roller.onnx",
        "scene_rollers.xml",
        EPISODE_DURATION_S,
        "STAND",
        handoff_policy_file="alpha_stand.onnx",
        handoff_after_s=4.0,
        kind="roller",
    ),
    BehaviorSpec(
        "roller_crouch",
        "roller_crouch.onnx",
        "scene_rollers.xml",
        EPISODE_DURATION_S,
        "STAND",
        handoff_policy_file="roller.onnx",
        handoff_after_s=3.5,
        phase_period_s=5.0,
        kind="phase",
    ),
    BehaviorSpec(
        "roulade",
        "roulade.onnx",
        "scene.xml",
        5.0,
        "STAND",
        handoff_policy_file="alpha_stand.onnx",
        handoff_after_s=2.0,
        kind="roulade",
    ),
)


RECOVERY_BEHAVIORS: tuple[BehaviorSpec, ...] = (
    BehaviorSpec(
        "recovery",
        "alpha_stand.onnx",
        "scene.xml",
        EPISODE_DURATION_S,
        "STAND",
        kind="recovery",
    ),
    BehaviorSpec(
        "recovery_roller",
        "alpha_stand.onnx",
        "scene_rollers.xml",
        EPISODE_DURATION_S,
        "STAND",
        kind="recovery",
    ),
)

ALL_BEHAVIORS: tuple[BehaviorSpec, ...] = BEHAVIORS + RECOVERY_BEHAVIORS


def split_for_ordinal(ordinal: int, episodes_per_policy: int) -> str:
    """Put the final 10% of each teacher into validation."""
    n_validation = max(1, round(episodes_per_policy * 0.1))
    return "validation" if ordinal >= episodes_per_policy - n_validation else "train"


def validate_contract() -> None:
    if CONTROL_HZ % OBSERVATION_HZ:
        raise ValueError("observation rate must divide the control rate")
    if len(DEFAULT_POSE) != ACTION_DIM:
        raise ValueError("default pose and action dimensions disagree")
    if any(b.duration_s != EPISODE_DURATION_S for b in ALL_BEHAVIORS):
        raise ValueError("every episode must be exactly five seconds")
    names = [b.name for b in ALL_BEHAVIORS]
    if len(names) != len(set(names)):
        raise ValueError("behavior names must be unique")