#!/usr/bin/env python3
"""
Launch the humanoid brain simulation.

Examples
--------
# Headless stub run (no GPU, no weights)
python scripts/run_sim.py --duration 10

# With interactive MuJoCo viewer
python scripts/run_sim.py --render --task "raise both arms" --duration 30

# With real model weights (CUDA required)
python scripts/run_sim.py --real --render --task "stand upright"
"""
from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

# EGL is the reliable headless backend on Linux; set before importing mujoco.
os.environ.setdefault("MUJOCO_GL", "egl")

import yaml

# Allow imports from project root
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from brain.brain import Brain
from env.humanoid_env import HumanoidTouchEnv
from hal.hal import HAL
from models.diffusion_policy import DiffusionPathPlanner
from models.language import OpenVLALanguageModel
from models.policy import CosmosPolicy
from models.world_model import CosmosWorldModel


def _load_yaml(path: str) -> dict:
    with open(path, encoding="utf-8") as f:
        return yaml.safe_load(f)


def main():
    parser = argparse.ArgumentParser(description="Run the humanoid brain in MuJoCo.")
    parser.add_argument(
        "--robot-config", default="configs/robot.yaml",
        help="Path to robot YAML config",
    )
    parser.add_argument(
        "--model-config", default="configs/models.yaml",
        help="Path to model YAML config",
    )
    parser.add_argument(
        "--task", default=None,
        help="Natural-language task instruction (overrides config default)",
    )
    parser.add_argument(
        "--render", action="store_true",
        help="Open interactive MuJoCo viewer (requires display)",
    )
    parser.add_argument(
        "--duration", type=float, default=30.0,
        help="Simulation wall-clock duration in seconds",
    )
    parser.add_argument(
        "--real", action="store_true",
        help="Disable stub mode and load actual model weights (needs GPU + HuggingFace)",
    )
    parser.add_argument(
        "--stub", action="store_true",
        help="Force stub mode (fast, CPU-only, no weights needed)",
    )
    parser.add_argument(
        "--log-level", default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s  %(name)-20s  %(levelname)s  %(message)s",
        datefmt="%H:%M:%S",
    )

    robot_cfg = _load_yaml(args.robot_config)
    model_cfg = _load_yaml(args.model_config)

    if args.stub:
        model_cfg["world_model"]["stub"] = True
        model_cfg["policy"]["stub"] = True
        model_cfg["language"]["stub"] = True
        model_cfg["diffusion_planner"]["stub"] = True
        logging.getLogger().info("Stub mode forced.")
    elif args.real:
        model_cfg["world_model"]["stub"] = False
        model_cfg["policy"]["stub"] = False
        model_cfg["language"]["stub"] = False
        model_cfg["diffusion_planner"]["stub"] = False
        logging.getLogger().info("Real model weights requested.")

    # ── build components
    env = HumanoidTouchEnv(robot_cfg)
    world_model = CosmosWorldModel(model_cfg["world_model"])
    policy = CosmosPolicy(model_cfg["policy"])
    language_model = OpenVLALanguageModel(model_cfg["language"])
    path_planner = DiffusionPathPlanner(model_cfg["diffusion_planner"])
    hal = HAL(env, robot_cfg["hal"])
    brain = Brain(
        hal, world_model, policy, language_model, model_cfg["brain"],
        path_planner=path_planner,
    )

    task = args.task or model_cfg["brain"]["default_task"]

    # ── run
    if args.render:
        brain.run_with_viewer(task=task, duration=args.duration)
    else:
        brain.run(task=task, duration=args.duration)

    print("\nFinal stats:", brain.stats)


if __name__ == "__main__":
    main()