"""
Training configuration for Vortex models.
Covers both 7B and 13B variants with hardware-specific optimizations.
"""

import torch

TRAINING_CONFIG = {
    # Training hyperparameters
    "learning_rate": 3e-4,
    "weight_decay": 0.1,
    "beta1": 0.9,
    "beta2": 0.95,
    "clip_grad_norm": 1.0,

    # Batch sizing
    "global_batch_size": 512,  # tokens per batch
    "micro_batch_size": 8,     # per GPU
    "gradient_accumulation_steps": 4,

    # Training schedule
    "max_steps": 100000,
    "warmup_steps": 2000,
    "save_interval": 5000,
    "eval_interval": 1000,
    "log_interval": 100,

    # Mixed precision
    "use_amp": True,
    "amp_dtype": torch.bfloat16,

    # Optimizer
    "optimizer": "AdamW",
    "use_fused": True,  # fused AdamW if available

    # Curriculum learning stages (as fractions of max_steps)
    "curriculum_stages": [
        {"name": "foundation", "start": 0.0, "end": 0.2},    # 0-20%
        {"name": "domain", "start": 0.2, "end": 0.5},       # 20-50%
        {"name": "reasoning", "start": 0.5, "end": 0.8},    # 50-80%
        {"name": "integration", "start": 0.8, "end": 1.0},  # 80-100%
    ],

    # Loss weights (science-aware loss)
    "loss_weights": {
        "lm_loss": 1.0,
        "equation_loss": 0.3,
        "domain_loss": 0.1,
        "citation_loss": 0.1,
        "numerical_loss": 0.2,
    },

    # Checkpointing
    "checkpoint_dir": "checkpoints",
    "save_optimizer_state": True,
    "save_scheduler_state": True,

    # Logging
    "log_dir": "logs",
    "use_wandb": False,
    "wandb_project": "vortex-scientific",

    # Data loading
    "num_workers": 8,
    "prefetch_factor": 2,
    "pin_memory": True,

    # Device configuration
    "device": "cuda",  # or "mps" for Apple Silicon
    "use_mps": False,

    # Quantization (for 13B on 8GB VRAM)
    "quantization": None,  # None, "int8", "int4"
}

# Hardware-specific overrides
TRAINING_CONFIG_7B_CUDA = TRAINING_CONFIG.copy()
TRAINING_CONFIG_7B_CUDA.update({
    "device": "cuda",
    "quantization": None,
    "micro_batch_size": 8,
})

TRAINING_CONFIG_13B_CUDA = TRAINING_CONFIG.copy()
TRAINING_CONFIG_13B_CUDA.update({
    "device": "cuda",
    "quantization": "int8",  # 13B needs INT8 on 8GB
    "micro_batch_size": 4,
})

TRAINING_CONFIG_MPS = TRAINING_CONFIG.copy()
TRAINING_CONFIG_MPS.update({
    "device": "mps",
    "use_mps": True,
    "use_amp": False,  # MPS doesn't support bfloat16 AMP well
    "micro_batch_size": 4,
})