#!/usr/bin/env python3
# /// script
# dependencies = [
#     "torch>=2.2.0",
#     "transformers>=4.48.0,<5.3.0",
#     "datasets>=3.0.0",
#     "peft>=0.14.0",
#     "trl>=0.15.0",
#     "bitsandbytes>=0.45.0",
#     "accelerate>=1.2.0",
#     "huggingface_hub",
#     "einops>=0.7.0",
#     "pillow>=10.0.0",
#     "numpy>=1.24.0",
#     "pyyaml>=6.0",
# ]
# ///
"""
🤖 Neon VLA Training Script for HuggingFace GPU Jobs.

Trains the G1 action heads on mixed data soup using a frozen video backbone.

Backbone Selection Guide:
  - Qwen2.5-Omni-7B  → Audio-native, best for spoken robot commands (default)
  - Cosmos-Reason2-8B → Physics reasoning, best for sim2real + manipulation
  - Qwen2.5-VL-7B    → Text-only instructions, most stable
  - Qwen2.5-VL-3B    → Fast iteration, fits on L4/Jetson

Flash Attention 2:
  - Auto-detected at runtime. Falls back to eager if unavailable.
  - Install: pip install flash-attn --no-build-isolation
  - 2-3× faster forward pass, ~30% less VRAM

Launch (A100 — best quality, Omni + FA2):
    hf jobs uv run --flavor a100-large --secrets HF_TOKEN --timeout 8h \\
        scripts/train_neon.py --backbone Qwen/Qwen2.5-Omni-7B --audio

Launch (A100 — physics reasoning):
    hf jobs uv run --flavor a100-large --secrets HF_TOKEN --timeout 8h \\
        scripts/train_neon.py --backbone nvidia/Cosmos-Reason2-8B

Launch (L4 — smaller, faster):
    hf jobs uv run --flavor l4x1 --secrets HF_TOKEN --timeout 8h \\
        scripts/train_neon.py --backbone Qwen/Qwen2.5-VL-3B-Instruct

Local:
    python scripts/train_neon.py --backbone Qwen/Qwen2.5-Omni-7B --audio --mode arms_only
"""

import os
import sys
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def _detect_flash_attention() -> str:
    """Detect if flash_attention_2 is available. Returns best attn implementation.

    FA2 gives 2-3× speedup and ~30% VRAM savings on Ampere+ GPUs (A100, L40S, RTX 4090).
    Falls back to "eager" (standard attention) if FA2 is not installed.
    SDPA is an alternative but less compatible with quantized models.
    """
    try:
        import flash_attn  # noqa: F401
        log.info(f"✅ flash_attention_2 available (v{flash_attn.__version__})")
        return "flash_attention_2"
    except ImportError:
        pass

    # Check if torch SDPA is available as middle ground
    try:
        import torch
        if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
            log.info("⚡ Using SDPA (scaled_dot_product_attention) — install flash-attn for 2× speedup")
            # Note: SDPA can have issues with 4-bit quantized models, stick with eager for safety
            # return "sdpa"
    except Exception:
        pass

    log.warning("⚠️ flash_attention_2 not available. Training will be slower.")
    log.warning("   Install: pip install flash-attn --no-build-isolation")
    return "eager"


def _detect_backbone_type(model_id: str) -> dict:
    """Detect backbone capabilities from model ID.

    Returns dict with:
        - is_omni: bool — native audio support
        - is_cosmos: bool — physics reasoning model
        - is_vl: bool — standard vision-language
        - recommended_audio: str — best audio encoder for this backbone
    """
    model_lower = model_id.lower()
    return {
        "is_omni": "omni" in model_lower,
        "is_cosmos": "cosmos" in model_lower,
        "is_vl": "vl" in model_lower and "omni" not in model_lower,
        "recommended_audio": (
            "omni" if "omni" in model_lower
            else "whisper"  # Cosmos and VL need external audio encoder
        ),
    }


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Train Neon VLA",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Backbone Comparison:
  ┌───────────────────────────┬───────┬─────────┬──────────┬────────────────┐
  │ Model                     │ Params│ VRAM(4b)│ Audio    │ Best For       │
  ├───────────────────────────┼───────┼─────────┼──────────┼────────────────┤
  │ Qwen/Qwen2.5-Omni-7B     │ 7B    │ ~8 GB   │ Native   │ Spoken commands│
  │ nvidia/Cosmos-Reason2-8B  │ 8B    │ ~10 GB  │ Whisper  │ Physics tasks  │
  │ Qwen/Qwen2.5-VL-7B       │ 7B    │ ~8 GB   │ Whisper  │ Text instruct  │
  │ Qwen/Qwen2.5-VL-3B       │ 3B    │ ~4 GB   │ Whisper  │ Fast iteration │
  │ Qwen/Qwen2.5-Omni-3B     │ 3B    │ ~4 GB   │ Native   │ Edge + audio   │
  │ nvidia/Cosmos-Reason2-2B  │ 2B    │ ~3 GB   │ Whisper  │ Edge + physics │
  └───────────────────────────┴───────┴─────────┴──────────┴────────────────┘
""",
    )
    parser.add_argument("--backbone", type=str, default="Qwen/Qwen2.5-Omni-7B",
                        help="Video backbone model ID (default: Qwen2.5-Omni-7B for audio-native)")
    parser.add_argument("--mode", type=str, default="arms_only",
                        choices=["arms_only", "upper_body", "whole_body"])
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--output", type=str, default="cagataydev/neon-g1-v1")
    parser.add_argument("--dataset", type=str, default="lerobot/xvla-agibot-world",
                        help="Primary training dataset")
    parser.add_argument("--chunk-size", type=int, default=16, help="Action chunk size")
    parser.add_argument("--audio", action="store_true",
                        help="Enable audio encoder (auto-enabled for Omni backbones)")
    parser.add_argument("--no-flash-attn", action="store_true",
                        help="Force eager attention (skip flash_attention_2 detection)")
    parser.add_argument("--attn", type=str, default=None,
                        choices=["flash_attention_2", "sdpa", "eager"],
                        help="Override attention implementation")
    args = parser.parse_args()

    import torch

    # GPU info
    if torch.cuda.is_available():
        gpu = torch.cuda.get_device_name(0)
        vram = torch.cuda.get_device_properties(0).total_memory / 1024**3
        log.info(f"GPU: {gpu}, VRAM: {vram:.1f} GB")

        # Check compute capability for FA2 (needs sm_80+ = Ampere+)
        cc = torch.cuda.get_device_capability(0)
        log.info(f"Compute capability: {cc[0]}.{cc[1]}")
        if cc[0] < 8:
            log.warning(f"GPU compute capability {cc[0]}.{cc[1]} < 8.0 — flash_attention_2 requires Ampere+")
    else:
        log.warning("No GPU! Training will be slow.")

    # Import neon (install if needed)
    try:
        from neon.model.neon_vla import NeonConfig, NeonVLA
        from neon.model.video_backbone import BackboneConfig
        from neon.data.data_soup import DataSourceConfig, DataSoupConfig
        from neon.training.config import TrainConfig
        from neon.training.train import NeonTrainer
    except ImportError:
        log.info("Installing neon-vla...")
        import subprocess
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "."])
        from neon.model.neon_vla import NeonConfig, NeonVLA
        from neon.model.video_backbone import BackboneConfig
        from neon.data.data_soup import DataSourceConfig, DataSoupConfig
        from neon.training.config import TrainConfig
        from neon.training.train import NeonTrainer

    # Detect attention implementation
    if args.attn:
        attn_impl = args.attn
        log.info(f"Attention: {attn_impl} (user override)")
    elif args.no_flash_attn:
        attn_impl = "eager"
        log.info("Attention: eager (--no-flash-attn flag)")
    else:
        attn_impl = _detect_flash_attention()

    # Detect backbone capabilities
    backbone_info = _detect_backbone_type(args.backbone)
    log.info(f"Backbone: {args.backbone}")
    log.info(f"  Omni (native audio): {backbone_info['is_omni']}")
    log.info(f"  Cosmos (physics):    {backbone_info['is_cosmos']}")
    log.info(f"  Recommended audio:   {backbone_info['recommended_audio']}")

    # Enable audio if Omni backbone or --audio flag
    enable_audio = args.audio or backbone_info["is_omni"]
    audio_config = None
    if enable_audio:
        audio_config = {
            "enabled": True,
            "encoder_type": backbone_info["recommended_audio"],
            "whisper_model": "openai/whisper-base",
            "freeze_encoder": True,
            "enable_speech_output": False,  # No TTS during training
        }
        log.info(f"Audio: ENABLED (encoder={backbone_info['recommended_audio']})")
    else:
        log.info("Audio: DISABLED (use --audio or Omni backbone to enable)")

    # Build config
    config = TrainConfig(
        model=NeonConfig(
            backbone=BackboneConfig(
                model_id=args.backbone,
                load_in_4bit=True,
                freeze_backbone=True,
                attn_implementation=attn_impl,
            ),
            control_mode=args.mode,
            include_locomotion=(args.mode != "arms_only"),
            num_action_steps=args.chunk_size,
            use_action_chunking=True,
            use_separate_heads=True,
            audio=audio_config,
        ),
        data=DataSoupConfig(
            action_chunk_size=args.chunk_size,
            sources=[
                DataSourceConfig(
                    name="primary",
                    type="lerobot",
                    path=args.dataset,
                    weight=1.0,
                ),
            ],
        ),
        epochs=args.epochs,
        batch_size=args.batch_size,
        learning_rate=args.lr,
        hub_model_id=args.output,
        push_to_hub=True,
    )

    log.info(f"Training config:")
    log.info(f"  Backbone:   {args.backbone}")
    log.info(f"  Attention:  {attn_impl}")
    log.info(f"  Mode:       {args.mode}")
    log.info(f"  Audio:      {enable_audio}")
    log.info(f"  Dataset:    {args.dataset}")
    log.info(f"  Chunk size: {args.chunk_size}")
    log.info(f"  Epochs:     {args.epochs}")
    log.info(f"  Batch:      {args.batch_size}")
    log.info(f"  Output:     {args.output}")

    # Train
    trainer = NeonTrainer(config)
    stats = trainer.train()

    log.info(f"✅ Training complete!")
    log.info(f"   Best loss: {stats.get('best_loss', 'N/A')}")
    log.info(f"   Steps: {stats.get('total_steps', 'N/A')}")
    log.info(f"   Output: https://huggingface.co/{args.output}")


if __name__ == "__main__":
    main()
