#!/usr/bin/env python3
"""
Test MrBalance directly from Hugging Face Hub.

Requirements:
    pip install -U torch transformers safetensors mujoco

Files:
    balance_plate_rl.py
    test_mrbalance_hf.py

Usage:
    python test_mrbalance_hf.py
    python test_mrbalance_hf.py --object sphere
    python test_mrbalance_hf.py --object egg
    python test_mrbalance_hf.py --object heavy_ball

The script:
    1. Downloads MrBalance from Hugging Face.
    2. Loads it through AutoModel with trust_remote_code=True.
    3. Creates the original MuJoCo environment.
    4. Uses the Hugging Face policy to control the plate.
    5. Prints episode statistics.
    6. Optionally launches the MuJoCo viewer.
"""

from __future__ import annotations

import argparse
import time

import numpy as np
import torch
from transformers import AutoModel

from balance_plate_rl import (
    DEVICE,
    MAX_EPISODE_STEPS,
    OBJECT_TYPES,
    PlateBalanceEnv,
)


# ============================================================================
# CONFIG
# ============================================================================

MODEL_ID = "fromziro/MrBalance"

DEFAULT_EPISODES = 5
DEFAULT_OBJECT = "sphere"

DEVICE_TO_USE = torch.device(
    DEVICE
)


# ============================================================================
# MODEL LOADING
# ============================================================================

def load_model():
    print("=" * 72)
    print("Loading MrBalance from Hugging Face")
    print("=" * 72)

    print(f"Model:  {MODEL_ID}")
    print(f"Device: {DEVICE_TO_USE}")
    print()

    model = AutoModel.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
    )

    model = model.to(DEVICE_TO_USE)
    model.eval()

    print("[load] Model loaded successfully.")
    print()

    # Print basic architecture information.
    print(
        f"[model] Observation size: "
        f"{model.config.observation_size}"
    )

    print(
        f"[model] Hidden size:       "
        f"{model.config.hidden_size}"
    )

    print(
        f"[model] Intermediate size: "
        f"{model.config.intermediate_size}"
    )

    print(
        f"[model] Bottleneck size:   "
        f"{model.config.bottleneck_size}"
    )

    print(
        f"[model] Action size:       "
        f"{model.config.action_size}"
    )

    print()

    return model


# ============================================================================
# SINGLE EPISODE
# ============================================================================

def run_episode(
    model,
    object_name: str,
    seed: int,
    render: bool = False,
):
    env = PlateBalanceEnv(
        seed=seed,
        render=render,
        active_objects=[object_name],
    )

    obs = env.reset(
        specific_object=object_name
    )

    total_reward = 0.0
    episode_length = 0

    if render:
        try:
            import mujoco.viewer

            viewer_context = (
                mujoco.viewer.launch_passive(
                    env.model,
                    env.data,
                )
            )
        except Exception as exc:
            env.close()
            raise RuntimeError(
                f"Could not launch MuJoCo viewer: {exc}"
            ) from exc
    else:
        viewer_context = None

    try:
        if viewer_context is not None:
            with viewer_context as viewer:

                while viewer.is_running():

                    step_start = time.time()

                    obs_tensor = torch.as_tensor(
                        obs,
                        dtype=torch.float32,
                        device=DEVICE_TO_USE,
                    ).unsqueeze(0)

                    with torch.no_grad():
                        output = model(
                            obs_tensor,
                            deterministic=True,
                        )

                    action = (
                        output.action[0]
                        .detach()
                        .cpu()
                        .numpy()
                        .astype(np.float64)
                    )

                    obs, reward, terminated, truncated, info = (
                        env.step(action)
                    )

                    total_reward += reward
                    episode_length += 1

                    viewer.sync()

                    if terminated or truncated:
                        break

                    # Match roughly the environment's 100 Hz control rate.
                    target_dt = 0.01
                    elapsed = time.time() - step_start

                    if elapsed < target_dt:
                        time.sleep(
                            target_dt - elapsed
                        )

        else:
            while True:

                obs_tensor = torch.as_tensor(
                    obs,
                    dtype=torch.float32,
                    device=DEVICE_TO_USE,
                ).unsqueeze(0)

                with torch.no_grad():
                    output = model(
                        obs_tensor,
                        deterministic=True,
                    )

                action = (
                    output.action[0]
                    .detach()
                    .cpu()
                    .numpy()
                    .astype(np.float64)
                )

                obs, reward, terminated, truncated, info = (
                    env.step(action)
                )

                total_reward += reward
                episode_length += 1

                if terminated or truncated:
                    break

    finally:
        env.close()

    success = (
        not info["fallen"]
        and episode_length >= MAX_EPISODE_STEPS
    )

    return {
        "object": object_name,
        "reward": float(total_reward),
        "length": int(episode_length),
        "success": bool(success),
        "fallen": bool(info["fallen"]),
        "distance": float(info["distance"]),
        "velocity": float(info["object_velocity"]),
    }


# ============================================================================
# RANDOM OBSERVATION SANITY CHECK
# ============================================================================

def sanity_check(model):
    """
    Verify that the Hugging Face model can actually execute inference
    independently of MuJoCo.
    """

    print("=" * 72)
    print("Running model sanity check")
    print("=" * 72)

    x = torch.randn(
        4,
        64,
        dtype=torch.float32,
        device=DEVICE_TO_USE,
    )

    with torch.no_grad():
        output = model(
            x,
            deterministic=True,
        )

    print(
        "[sanity] action shape:",
        tuple(output.action.shape),
    )

    print(
        "[sanity] value shape:",
        tuple(output.value.shape),
    )

    print(
        "[sanity] action range:",
        float(output.action.min()),
        "to",
        float(output.action.max()),
    )

    print(
        "[sanity] example action:",
        output.action[0].detach().cpu().numpy(),
    )

    print("[sanity] PASS")
    print()


# ============================================================================
# MAIN
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Test MrBalance from Hugging Face."
    )

    parser.add_argument(
        "--object",
        type=str,
        default=DEFAULT_OBJECT,
        choices=OBJECT_TYPES,
        help="Object to balance.",
    )

    parser.add_argument(
        "--episodes",
        type=int,
        default=DEFAULT_EPISODES,
        help="Number of evaluation episodes.",
    )

    parser.add_argument(
        "--render",
        action="store_true",
        help="Launch interactive MuJoCo viewer.",
    )

    parser.add_argument(
        "--seed",
        type=int,
        default=12345,
        help="Evaluation seed.",
    )

    args = parser.parse_args()

    model = load_model()

    sanity_check(model)

    print("=" * 72)
    print(
        f"Testing object: {args.object}"
    )
    print(
        f"Episodes:       {args.episodes}"
    )
    print("=" * 72)
    print()

    results = []

    for episode in range(args.episodes):

        print(
            f"[episode {episode + 1}/{args.episodes}] "
            f"Running..."
        )

        result = run_episode(
            model=model,
            object_name=args.object,
            seed=args.seed + episode,
            render=args.render,
        )

        results.append(result)

        status = (
            "SUCCESS"
            if result["success"]
            else (
                "FALL"
                if result["fallen"]
                else "TIMEOUT"
            )
        )

        print(
            f"  status:   {status}"
        )

        print(
            f"  reward:   {result['reward']:.2f}"
        )

        print(
            f"  length:   {result['length']}"
        )

        print(
            f"  distance: {result['distance']:.4f} m"
        )

        print(
            f"  velocity: {result['velocity']:.4f} m/s"
        )

        print()

    # ------------------------------------------------------------------------
    # Summary
    # ------------------------------------------------------------------------

    rewards = [
        r["reward"]
        for r in results
    ]

    lengths = [
        r["length"]
        for r in results
    ]

    distances = [
        r["distance"]
        for r in results
    ]

    successes = sum(
        r["success"]
        for r in results
    )

    print("=" * 72)
    print("RESULTS")
    print("=" * 72)

    print(
        f"Object:            {args.object}"
    )

    print(
        f"Reward mean:       {np.mean(rewards):.2f}"
    )

    print(
        f"Reward std:        {np.std(rewards):.2f}"
    )

    print(
        f"Length mean:       {np.mean(lengths):.1f}"
    )

    print(
        f"Survival rate:     "
        f"{100.0 * sum(not r['fallen'] for r in results) / len(results):.1f}%"
    )

    print(
        f"Success rate:      "
        f"{100.0 * successes / len(results):.1f}%"
    )

    print(
        f"Tracking error:    "
        f"{np.mean(distances):.4f} m"
    )

    print("=" * 72)


if __name__ == "__main__":
    main()