import argparse
import math
from pathlib import Path

import torch
import yaml

# AMP feature order must match `_build_amp_features` in `t1_env_cfg.py`:
# [joint_pos_rel(23), joint_vel(23), root_lin_vel(3), root_ang_vel(3), projected_gravity(3)].
T1_JOINT_NAMES = [
    "AAHead_yaw",
    "Head_pitch",
    "Left_Shoulder_Pitch",
    "Left_Shoulder_Roll",
    "Left_Elbow_Pitch",
    "Left_Elbow_Yaw",
    "Right_Shoulder_Pitch",
    "Right_Shoulder_Roll",
    "Right_Elbow_Pitch",
    "Right_Elbow_Yaw",
    "Waist",
    "Left_Hip_Pitch",
    "Left_Hip_Roll",
    "Left_Hip_Yaw",
    "Left_Knee_Pitch",
    "Left_Ankle_Pitch",
    "Left_Ankle_Roll",
    "Right_Hip_Pitch",
    "Right_Hip_Roll",
    "Right_Hip_Yaw",
    "Right_Knee_Pitch",
    "Right_Ankle_Pitch",
    "Right_Ankle_Roll",
]

JOINT_TO_IDX = {name: i for i, name in enumerate(T1_JOINT_NAMES)}

# Mirror rules aligned with `behaviors/custom/keyframe/keyframe.py`.
MOTOR_SYMMETRY = {
    "Head_yaw": (("Head_yaw",), False),
    "Head_pitch": (("Head_pitch",), False),
    "Shoulder_Pitch": (("Left_Shoulder_Pitch", "Right_Shoulder_Pitch"), False),
    "Shoulder_Roll": (("Left_Shoulder_Roll", "Right_Shoulder_Roll"), True),
    "Elbow_Pitch": (("Left_Elbow_Pitch", "Right_Elbow_Pitch"), False),
    "Elbow_Yaw": (("Left_Elbow_Yaw", "Right_Elbow_Yaw"), True),
    "Waist": (("Waist",), False),
    "Hip_Pitch": (("Left_Hip_Pitch", "Right_Hip_Pitch"), False),
    "Hip_Roll": (("Left_Hip_Roll", "Right_Hip_Roll"), True),
    "Hip_Yaw": (("Left_Hip_Yaw", "Right_Hip_Yaw"), True),
    "Knee_Pitch": (("Left_Knee_Pitch", "Right_Knee_Pitch"), False),
    "Ankle_Pitch": (("Left_Ankle_Pitch", "Right_Ankle_Pitch"), False),
    "Ankle_Roll": (("Left_Ankle_Roll", "Right_Ankle_Roll"), True),
}
READABLE_TO_POLICY = {"Head_yaw": "AAHead_yaw"}


def decode_keyframe_motor_positions(raw_motor_positions: dict[str, float]) -> dict[str, float]:
    """Decode one keyframe into per-joint radians."""
    out: dict[str, float] = {}
    deg_to_rad = math.pi / 180.0
    for readable_name, position_deg in raw_motor_positions.items():
        if readable_name in MOTOR_SYMMETRY:
            motor_names, is_inverse_direction = MOTOR_SYMMETRY[readable_name]
            invert_state = bool(is_inverse_direction)
            for motor_name in motor_names:
                signed_deg = position_deg if invert_state else -position_deg
                invert_state = False
                out_name = READABLE_TO_POLICY.get(motor_name, motor_name)
                out[out_name] = float(signed_deg) * deg_to_rad
        else:
            out_name = READABLE_TO_POLICY.get(readable_name, readable_name)
            out[out_name] = float(position_deg) * deg_to_rad
    return out


def load_sequence(yaml_path: Path) -> list[tuple[float, torch.Tensor]]:
    """Load yaml keyframes -> list[(delta_seconds, joint_pos_vec)]."""
    with yaml_path.open("r", encoding="utf-8") as f:
        desc = yaml.safe_load(f) or {}
    out: list[tuple[float, torch.Tensor]] = []
    for keyframe in desc.get("keyframes", []):
        delta_s = max(float(keyframe.get("delta", 0.1)), 1e-3)
        raw = keyframe.get("motor_positions", {}) or {}
        decoded = decode_keyframe_motor_positions(raw)
        joint_pos = torch.zeros(len(T1_JOINT_NAMES), dtype=torch.float32)
        for j_name, j_val in decoded.items():
            idx = JOINT_TO_IDX.get(j_name, None)
            if idx is not None:
                joint_pos[idx] = float(j_val)
        out.append((delta_s, joint_pos))
    return out


def sequence_to_amp_features(
    sequence: list[tuple[float, torch.Tensor]],
    sample_fps: float,
    projected_gravity: tuple[float, float, float],
) -> torch.Tensor:
    """Convert decoded sequence into AMP features tensor (N, 55)."""
    if len(sequence) == 0:
        raise ValueError("Empty keyframe sequence.")
    dt = 1.0 / max(sample_fps, 1e-6)
    grav = torch.tensor(projected_gravity, dtype=torch.float32)

    frames_joint_pos: list[torch.Tensor] = []
    for delta_s, joint_pos in sequence:
        repeat = max(int(round(delta_s / dt)), 1)
        for _ in range(repeat):
            frames_joint_pos.append(joint_pos.clone())
    if len(frames_joint_pos) < 2:
        frames_joint_pos.append(frames_joint_pos[0].clone())

    pos = torch.stack(frames_joint_pos, dim=0)
    vel = torch.zeros_like(pos)
    vel[1:] = (pos[1:] - pos[:-1]) / dt
    vel[0] = vel[1]

    root_lin = torch.zeros((pos.shape[0], 3), dtype=torch.float32)
    root_ang = torch.zeros((pos.shape[0], 3), dtype=torch.float32)
    grav_batch = grav.unsqueeze(0).repeat(pos.shape[0], 1)
    return torch.cat([pos, vel, root_lin, root_ang, grav_batch], dim=-1)


def main():
    parser = argparse.ArgumentParser(description="Build AMP expert features from get_up keyframe YAML files.")
    parser.add_argument(
        "--front_yaml",
        type=str,
        default="behaviors/custom/keyframe/get_up/get_up_front.yaml",
        help="Path to front get-up YAML.",
    )
    parser.add_argument(
        "--back_yaml",
        type=str,
        default="behaviors/custom/keyframe/get_up/get_up_back.yaml",
        help="Path to back get-up YAML.",
    )
    parser.add_argument(
        "--sample_fps",
        type=float,
        default=50.0,
        help="Sampling fps when expanding keyframe durations.",
    )
    parser.add_argument(
        "--repeat_cycles",
        type=int,
        default=200,
        help="How many times to repeat front+back sequences to enlarge dataset.",
    )
    parser.add_argument(
        "--projected_gravity",
        type=float,
        nargs=3,
        default=(0.0, 0.0, -1.0),
        help="Projected gravity feature used for synthesized expert data.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="rl_game/get_up/amp/expert_features.pt",
        help="Output expert feature file path.",
    )
    args = parser.parse_args()

    front_path = Path(args.front_yaml).expanduser().resolve()
    back_path = Path(args.back_yaml).expanduser().resolve()
    if not front_path.is_file():
        raise FileNotFoundError(f"Front YAML not found: {front_path}")
    if not back_path.is_file():
        raise FileNotFoundError(f"Back YAML not found: {back_path}")

    front_seq = load_sequence(front_path)
    back_seq = load_sequence(back_path)
    front_feat = sequence_to_amp_features(front_seq, args.sample_fps, tuple(args.projected_gravity))
    back_feat = sequence_to_amp_features(back_seq, args.sample_fps, tuple(args.projected_gravity))
    base_feat = torch.cat([front_feat, back_feat], dim=0)

    repeat_cycles = max(int(args.repeat_cycles), 1)
    expert_features = base_feat.repeat(repeat_cycles, 1).contiguous()

    out_path = Path(args.output).expanduser()
    if not out_path.is_absolute():
        out_path = Path.cwd() / out_path
    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "expert_features": expert_features,
            "feature_dim": int(expert_features.shape[1]),
            "num_samples": int(expert_features.shape[0]),
            "source": "get_up_keyframe_yaml",
            "front_yaml": str(front_path),
            "back_yaml": str(back_path),
            "sample_fps": float(args.sample_fps),
            "repeat_cycles": repeat_cycles,
            "projected_gravity": [float(v) for v in args.projected_gravity],
        },
        str(out_path),
    )
    print(f"[INFO]: saved expert features -> {out_path}")
    print(f"[INFO]: shape={tuple(expert_features.shape)}")


if __name__ == "__main__":
    main()
