"""
Training configuration for Neon VLA.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional

from neon.data.data_soup import DataSourceConfig, DataSoupConfig
from neon.model.neon_vla import NeonConfig


@dataclass
class TrainConfig:
    """Full training configuration."""

    # Model
    model: NeonConfig = field(default_factory=NeonConfig)

    # Data
    data: DataSoupConfig = field(default_factory=DataSoupConfig)

    # Training hyperparameters
    epochs: int = 3
    batch_size: int = 4
    gradient_accumulation_steps: int = 8
    learning_rate: float = 2e-4
    weight_decay: float = 0.01
    warmup_ratio: float = 0.1
    lr_scheduler: str = "cosine"
    max_grad_norm: float = 0.3  # Parameter Golf: tighter clipping for small heads

    # Optimization — Parameter Golf Adam betas for action prediction
    optim: str = "paged_adamw_8bit"
    adam_beta1: float = 0.85  # Lower β₁ — faster adaptation to shifting action distributions
    adam_beta2: float = 0.99  # Standard β₂ from MicroGPT engine
    bf16: bool = True
    gradient_checkpointing: bool = True

    # LoRA (for backbone fine-tuning, if not frozen)
    use_lora: bool = True
    lora_r: int = 32
    lora_alpha: int = 64
    lora_dropout: float = 0.05
    lora_target_modules: List[str] = field(
        default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
    )

    # Saving
    output_dir: str = "/tmp/neon_output"
    hub_model_id: str = "cagataydev/neon-g1-v1"
    push_to_hub: bool = True
    save_steps: int = 500
    save_total_limit: int = 3

    # Logging
    logging_steps: int = 10
    report_to: str = "none"  # "wandb" for W&B

    # Hardware
    num_workers: int = 4
    seed: int = 42


def default_arms_only_config() -> TrainConfig:
    """Default config for arms-only training (tabletop manipulation).

    Uses Qwen2.5-Omni-7B with flash_attention_2 auto-detection.
    Audio is enabled via Omni's native encoder.
    """
    return TrainConfig(
        model=NeonConfig(
            backbone=__import__("neon.model.video_backbone", fromlist=["BackboneConfig"]).BackboneConfig(
                model_id="Qwen/Qwen2.5-Omni-7B",
                load_in_4bit=True,
                freeze_backbone=True,
                attn_implementation="flash_attention_2",  # Auto-falls back to eager
            ),
            control_mode="arms_only",
            include_locomotion=False,
            num_action_steps=16,
            use_action_chunking=True,
            audio={
                "enabled": True,
                "encoder_type": "omni",
                "enable_speech_output": False,
            },
        ),
        epochs=3,
        batch_size=4,
        learning_rate=2e-4,
    )


def default_wholebody_config() -> TrainConfig:
    """Default config for whole-body training (locomotion + manipulation).

    Uses Qwen2.5-Omni-7B with flash_attention_2 auto-detection.
    """
    return TrainConfig(
        model=NeonConfig(
            backbone=__import__("neon.model.video_backbone", fromlist=["BackboneConfig"]).BackboneConfig(
                model_id="Qwen/Qwen2.5-Omni-7B",
                load_in_4bit=True,
                freeze_backbone=True,
                attn_implementation="flash_attention_2",
            ),
            control_mode="whole_body",
            include_locomotion=True,
            num_action_steps=16,
            use_action_chunking=True,
            audio={
                "enabled": True,
                "encoder_type": "omni",
                "enable_speech_output": False,
            },
        ),
        epochs=5,
        batch_size=2,
        gradient_accumulation_steps=16,
        learning_rate=1e-4,
    )


def cosmos_physics_config() -> TrainConfig:
    """Config for physics-focused training with Cosmos-Reason2-8B.

    Best for sim2real transfer, physics-heavy manipulation, world model tasks.
    Uses Whisper encoder for audio (Cosmos doesn't have native audio).
    Requires A100 (40GB+).
    """
    return TrainConfig(
        model=NeonConfig(
            backbone=__import__("neon.model.video_backbone", fromlist=["BackboneConfig"]).BackboneConfig(
                model_id="nvidia/Cosmos-Reason2-8B",
                load_in_4bit=True,
                freeze_backbone=True,
                attn_implementation="flash_attention_2",
            ),
            control_mode="arms_only",
            include_locomotion=True,
            num_action_steps=16,
            use_action_chunking=True,
            mlp_hidden=768,
            action_head_layers=4,
            audio={
                "enabled": True,
                "encoder_type": "whisper",
                "whisper_model": "openai/whisper-base",
                "freeze_encoder": True,
                "enable_speech_output": False,
            },
        ),
        epochs=5,
        batch_size=2,
        gradient_accumulation_steps=16,
        learning_rate=1e-4,
        hub_model_id="cagataydev/neon-g1-cosmos-v1",
    )


def edge_3b_config() -> TrainConfig:
    """Lightweight config for Jetson Orin / L4 deployment.

    Uses Qwen2.5-Omni-3B — smallest audio-native backbone.
    Fits in ~4GB VRAM (4-bit).
    """
    return TrainConfig(
        model=NeonConfig(
            backbone=__import__("neon.model.video_backbone", fromlist=["BackboneConfig"]).BackboneConfig(
                model_id="Qwen/Qwen2.5-Omni-3B",
                load_in_4bit=True,
                freeze_backbone=True,
                attn_implementation="flash_attention_2",
            ),
            control_mode="arms_only",
            include_locomotion=False,
            num_action_steps=8,
            use_action_chunking=True,
            mlp_hidden=256,
            action_head_layers=2,
            audio={
                "enabled": True,
                "encoder_type": "omni",
                "enable_speech_output": False,
            },
        ),
        epochs=3,
        batch_size=8,
        learning_rate=3e-4,
        hub_model_id="cagataydev/neon-g1-edge-v1",
    )
