from pathlib import Path

BASE_DIR = Path(__file__).parent

MODELS = {
    "mm-10m-1": {
        "label": "MM-10M",
        "accent": "#9a6237",
        "description": "10M-parameter policy.",
    },
    "mm-base-e": {
        "label": "MM-Base-E",
        "accent": "#2f6a73",
        "description": "Trained with a strong energy regularization.",
    },
    "mm-100M-s1": {
        "label": "MM-100M",
        "accent": "#7a4c8f",
        "description": "100M-parameter policy.",
    },
    "mm-fullbody-base": {
        "label": "MM-FullBody-Base",
        "accent": "#c13a5b",
        "description": "Released baseline checkpoint (50M parameters).",
    },
}

MODEL_KEYS = list(MODELS.keys())

MOTIONS = {
    "KIT_314_walking_slow03_poses": {
        "label": "Walking · Slow",
        "category": "Walking",
        "dataset": "KIT / KINESIS",
    },
    "KIT_167_walking_medium03_poses": {
        "label": "Walking · Medium",
        "category": "Walking",
        "dataset": "KIT / KINESIS",
    },
    "KIT_424_walking_fast10_poses": {
        "label": "Walking · Fast",
        "category": "Walking",
        "dataset": "KIT / KINESIS",
    },
    "KIT_314_walking_run09_poses": {
        "label": "Running",
        "category": "Walking",
        "dataset": "KIT / KINESIS",
    },
    "KIT_11_WalkingStraightForwards05_poses": {
        "label": "Walk Straight Forward",
        "category": "Walking",
        "dataset": "KIT / KINESIS",
    },
    "KIT_4_WalkingStraightBackwards01_poses": {
        "label": "Walk Backwards · Subject 4",
        "category": "Backwards",
        "dataset": "KIT / KINESIS",
    },
    "KIT_6_WalkingStraightBackwards04_poses": {
        "label": "Walk Backwards · Subject 6",
        "category": "Backwards",
        "dataset": "KIT / KINESIS",
    },
    "KIT_7_WalkingStraightBackwards04_poses": {
        "label": "Walk Backwards · Subject 7",
        "category": "Backwards",
        "dataset": "KIT / KINESIS",
    },
    "KIT_10_WalkingStraightBackwards07_poses": {
        "label": "Walk Backwards · Subject 10",
        "category": "Backwards",
        "dataset": "KIT / KINESIS",
    },
    "KIT_11_WalkingStraightBackwards09_1_poses": {
        "label": "Walk Backwards · Subject 11",
        "category": "Backwards",
        "dataset": "KIT / KINESIS",
    },
    "KIT_12_WalkingStraightBackwards05_poses": {
        "label": "Walk Backwards · Subject 12",
        "category": "Backwards",
        "dataset": "KIT / KINESIS",
    },
    "KIT_3_turn_right04_poses": {
        "label": "Turn Right · Subject 3",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_8_RightTurn10_poses": {
        "label": "Turn Right · Subject 8",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_10_RightTurn02_poses": {
        "label": "Turn Right · Subject 10",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_348_turn_right06_poses": {
        "label": "Turn Right · Subject 348",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_424_turn_right02_poses": {
        "label": "Turn Right · Subject 424",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_7_LeftTurn06_poses": {
        "label": "Turn Left · Subject 7",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_11_LeftTurn05_poses": {
        "label": "Turn Left · Subject 11",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_12_LeftTurn03_poses": {
        "label": "Turn Left · Subject 12",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_359_turn_left05_poses": {
        "label": "Turn Left · Subject 359",
        "category": "Turning",
        "dataset": "KIT / KINESIS",
    },
    "KIT_4_WalkInClockwiseCircle04_poses": {
        "label": "Clockwise Circle · Subject 4",
        "category": "Circular",
        "dataset": "KIT / KINESIS",
    },
    "KIT_9_WalkInClockwiseCircle04_poses": {
        "label": "Clockwise Circle · Subject 9a",
        "category": "Circular",
        "dataset": "KIT / KINESIS",
    },
    "KIT_9_WalkInClockwiseCircle06_poses": {
        "label": "Clockwise Circle · Subject 9b",
        "category": "Circular",
        "dataset": "KIT / KINESIS",
    },
    "KIT_12_WalkInClockwiseCircle11_poses": {
        "label": "Clockwise Circle · Subject 12",
        "category": "Circular",
        "dataset": "KIT / KINESIS",
    },
    "KIT_6_WalkInCounterClockwiseCircle04_1_poses": {
        "label": "Counter-Clockwise Circle · Subject 6a",
        "category": "Circular",
        "dataset": "KIT / KINESIS",
    },
    "KIT_6_WalkInCounterClockwiseCircle06_1_poses": {
        "label": "Counter-Clockwise Circle · Subject 6b",
        "category": "Circular",
        "dataset": "KIT / KINESIS",
    },
    "KIT_11_WalkInCounterClockwiseCircle02_poses": {
        "label": "Counter-Clockwise Circle · Subject 11",
        "category": "Circular",
        "dataset": "KIT / KINESIS",
    },
}

CATEGORIES = ["All", "Walking", "Backwards", "Turning", "Circular"]
DEFAULT_MOTION = "Walking · Medium"

_LABEL_TO_KEY = {metadata["label"]: motion_key for motion_key, metadata in MOTIONS.items()}


def label_to_key(label: str) -> str:
    return _LABEL_TO_KEY[label]


def get_motion_metadata(label: str) -> dict[str, str]:
    return MOTIONS[label_to_key(label)]


def get_video_path(model_key: str, motion_key: str) -> str:
    path = BASE_DIR / "videos" / "demo" / model_key / motion_key / "eval_MyoFullBody.mp4"
    return str(path)


def get_motions_by_category(category: str) -> list[str]:
    if category == "All":
        return [metadata["label"] for metadata in MOTIONS.values()]
    return [
        metadata["label"]
        for metadata in MOTIONS.values()
        if metadata["category"] == category
    ]


CITATION_BIBTEX = """@article{li2026towards,
  title={Towards Embodied AI with MuscleMimic: Unlocking full-body musculoskeletal motor learning at scale},
  author={Li, Chengkun and Wang, Cheryl and Ziliotto, Bianca and Simos, Merkourios and Kovecses, Jozsef and Durandau, Guillaume and Mathis, Alexander},
  journal={arXiv preprint arXiv:2603.25544},
  year={2026}
}"""


RESOURCE_LINKS = [
    {
        "label": "Paper",
        "url": "https://arxiv.org/abs/2603.25544",
        "note": "Manuscript and supplementary figures",
    },
    {
        "label": "Code",
        "url": "https://github.com/amathislab/musclemimic",
        "note": "Training and evaluation codebase",
    },
    {
        "label": "Checkpoint",
        "url": "https://huggingface.co/amathislab/mm-fullbody-base",
        "note": "Released pretrained full-body model",
    },
    {
        "label": "Dataset",
        "url": "https://huggingface.co/datasets/amathislab/musclemimic-retargeted",
        "note": "Retargeted motion dataset",
    },
]

HIGHLIGHT_VIDEOS = [
    {"file": "walking_running.mp4", "label": "Walking And Running", "tag": "Locomotion"},
    {"file": "walking_turning.mp4", "label": "Walking And Turning", "tag": "Locomotion"},
    {"file": "lifting_box.mp4", "label": "Lifting Box", "tag": "Manipulation"},
    {"file": "drinking_water.mp4", "label": "Drinking Water", "tag": "Manipulation"},
    {"file": "dancing.mp4", "label": "Dancing", "tag": "Dynamic Motion"},
    {"file": "jumpingjack.mp4", "label": "Jumping Jack", "tag": "Dynamic Motion"},
]