"""
FORGE-3B: Master Configuration
All hyperparameters, paths, and GPU optimization settings in one place.
"""

from __future__ import annotations
import os
import json
import math
from dataclasses import dataclass, field, asdict
from typing import Optional, List, Literal
from pathlib import Path


# ─────────────────────────────────────────────────────────────────────────────
# GPU GLOBAL OPTIMIZATION FLAGS
# These are applied at import time across the entire process.
# ─────────────────────────────────────────────────────────────────────────────

import torch

# Allow TF32 for matmuls and convolutions — massive speedup on Ampere/Hopper
# with negligible precision loss for language model training.
if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.backends.cudnn.benchmark = True          # auto-tune cuDNN kernels
    try:
        torch.backends.cuda.enable_flash_sdp(True)     # Flash SDP in PyTorch 2.x
        torch.backends.cuda.enable_mem_efficient_sdp(True)
        torch.backends.cuda.enable_math_sdp(False)     # disable slow math SDP
    except Exception:
        pass

# NCCL tuning for 64x H100 GPU cluster (8x 8-GPU nodes via NVLink + 3.2 Tbps InfiniBand)
os.environ.setdefault("NCCL_MIN_NCHANNELS", "32")
os.environ.setdefault("NCCL_MAX_NCHANNELS", "64")
os.environ.setdefault("NCCL_NSOCKS_PERTHREAD", "8")
os.environ.setdefault("NCCL_SOCKET_NTHREADS", "8")
os.environ.setdefault("NCCL_BUFFSIZE", "16777216")      # 16MB ring buffer for high-bandwidth AllReduce
os.environ.setdefault("NCCL_NET_GDR_LEVEL", "5")        # GPUDirect RDMA Level 5 for NVLink + InfiniBand
os.environ.setdefault("NCCL_CROSS_NIC", "1")            # Multi-NIC rail-optimized routing
os.environ.setdefault("NCCL_IB_DISABLE", "0")           # InfiniBand enabled
os.environ.setdefault("NCCL_P2P_DISABLE", "0")          # NVLink P2P enabled
os.environ.setdefault("NCCL_NVLS_ENABLE", "1")          # NVLink SHARP collective offload
os.environ.setdefault("CUDA_DEVICE_MAX_CONNECTIONS", "1") # Precise CUDA stream scheduling

# CUDA memory allocator — reduce fragmentation for large model + optimizer states
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF",
    "expandable_segments:True,max_split_size_mb:128,garbage_collection_threshold:0.8")

# Tokenizer parallelism — disable HF tokenizer parallelism (we use CRAYON)
os.environ["TOKENIZERS_PARALLELISM"] = "false"


# ─────────────────────────────────────────────────────────────────────────────
# MODEL CONFIGURATION
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class ForgeModelConfig:
    """Complete FORGE-500M architecture configuration. 500.02M total params."""
    
    # ── Core Dimensions ──────────────────────────────────────────────────────
    vocab_size: int = 206_464          # CRAYON standard profile (206,373) padded to ×128
    d_model: int = 1024               # 500M: 1024 (down from 1280)
    n_layers: int = 18                # 500M: 18 layers (15 ARG + 3 MHA)
    max_seq_len: int = 4096
    
    # ── Layer Distribution ────────────────────────────────────────────────────
    # Layer 0-4 = ARG; Layer 5 = MHA; repeat ×3 across 18 layers
    mha_layer_indices: List[int] = field(default_factory=lambda: [5, 11, 17])
    # Dense FFN on odd layers, HSE on even layers
    dense_ffn_layer_indices: List[int] = field(default_factory=lambda: list(range(1, 18, 2)))
    hse_ffn_layer_indices: List[int] = field(default_factory=lambda: list(range(0, 18, 2)))
    
    # ── ARG (Adaptive Recurrent Gating) ──────────────────────────────────────
    arg_d_inner: int = 1024           # matches d_model for 500M
    arg_d_state: int = 48
    arg_d_rank: int = 48               # dt projection rank
    arg_conv_kernel: int = 4
    arg_local_window: int = 64         # local attention window (phase 1+2)
    arg_local_n_heads: int = 8
    arg_local_n_kv_heads: int = 2
    arg_head_dim: int = 64            # 1024 / 16 heads
    arg_scalar_gate: bool = True       # scalar gate per token (vs d_model-dim)
    
    # ── MHA (Multi-Head Attention) ────────────────────────────────────────────
    mha_n_heads: int = 16
    mha_n_kv_heads: int = 4            # GQA: 4:1 ratio
    mha_head_dim: int = 64            # 1024 / 16 heads
    rope_theta: float = 500_000.0      # long-context RoPE base
    rope_scaling_type: Optional[str] = None  # None | "yarn"
    rope_scaling_factor: float = 1.0
    
    # ── Dense FFN (SwiGLU) ────────────────────────────────────────────────────
    dense_d_ff: int = 2560             # 2.5 × d_model, multiple of 128
    
    # ── HSE FFN (Hierarchical Sparse Expert) ──────────────────────────────────
    hse_n_domains: int = 4
    hse_n_experts_per_domain: int = 6  # total 24 experts (down from 32)
    hse_top_k: int = 2                 # active per token
    hse_d_ff_expert: int = 210        # scaled for 500M total params
    hse_capacity_factor: float = 1.25
    hse_expert_dropout: float = 0.1
    hse_aux_loss_alpha: float = 0.01
    hse_gumbel_tau_init: float = 1.0
    hse_gumbel_tau_final: float = 0.1
    hse_gumbel_tau_warmup_tokens: int = 10_000_000_000  # anneal over 10B tokens
    
    # ── Normalization ─────────────────────────────────────────────────────────
    norm_type: str = "dgn"             # "dgn" or "rmsnorm"
    dgn_n_groups: int = 16            # 1024 / 16 = 64 per group
    norm_eps: float = 1e-6
    
    # ── Initialization ────────────────────────────────────────────────────────
    init_std: float = 0.02
    init_std_residual_scale: bool = True  # scale by 1/sqrt(2n_layers)
    
    # ── Misc ──────────────────────────────────────────────────────────────────
    dropout: float = 0.0               # 0.0 during pretraining
    tie_word_embeddings: bool = True
    use_cache: bool = True             # KV cache for inference
    pad_token_id: int = 0
    bos_token_id: int = 1
    eos_token_id: int = 2
    
    # ── Compile/Optimization Flags ────────────────────────────────────
    use_flash_attention: bool = True
    use_triton_kernels: bool = True    # Triton fused ops
    use_torch_compile: bool = True
    compile_mode: str = "max-autotune"  # max throughput on H100
    use_gradient_checkpointing: bool = False  # 500M fits easily in 80GB without checkpointing
    gradient_checkpointing_ratio: float = 0.0
    
    def n_params_total(self) -> int:
        """Approximate total parameter count."""
        embed = self.vocab_size * self.d_model
        n_arg = len([i for i in range(self.n_layers) if i not in self.mha_layer_indices])
        n_mha = len(self.mha_layer_indices)
        n_dense = len(self.dense_ffn_layer_indices)
        n_hse = len(self.hse_ffn_layer_indices)
        
        arg_per = (
            self.d_model * 2 * self.arg_d_inner +      # in_proj
            self.arg_d_inner * self.arg_conv_kernel +   # conv1d
            self.arg_d_inner * (self.arg_d_rank + 2 * self.arg_d_state) +  # x_proj
            self.arg_d_rank * self.arg_d_inner +        # dt_proj
            self.arg_d_state * 2 +                      # nu, theta
            self.arg_d_inner +                          # D
            self.arg_d_inner * self.d_model +           # out_proj
            self.arg_local_n_heads * self.arg_head_dim * self.d_model +  # Q
            self.arg_local_n_kv_heads * self.arg_head_dim * self.d_model * 2 +  # KV
            self.arg_local_n_heads * self.arg_head_dim * self.d_model +  # O
            self.d_model * self.arg_d_state             # CPB
        )
        
        mha_per = (
            self.d_model * self.mha_n_heads * self.mha_head_dim +      # Q
            self.d_model * self.mha_n_kv_heads * self.mha_head_dim * 2 +  # KV
            self.d_model * self.mha_n_heads * self.mha_head_dim         # O
        )
        
        dense_per = 3 * self.d_model * self.dense_d_ff
        
        n_experts_total = self.hse_n_domains * self.hse_n_experts_per_domain
        hse_per = n_experts_total * 3 * self.d_model * self.hse_d_ff_expert
        
        total = embed + n_arg * arg_per + n_mha * mha_per + n_dense * dense_per + n_hse * hse_per
        return total
    
    def to_json(self, path: str):
        with open(path, "w") as f:
            json.dump(asdict(self), f, indent=2)
    
    @classmethod
    def from_json(cls, path: str) -> "ForgeModelConfig":
        import inspect
        with open(path) as f:
            data = json.load(f)
        sig = inspect.signature(cls)
        filtered_data = {k: v for k, v in data.items() if k in sig.parameters}
        return cls(**filtered_data)


# ─────────────────────────────────────────────────────────────────────────────
# TRAINING CONFIGURATION
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class PretrainConfig:
    """Phase-aware pretraining configuration for FORGE-500M.
    
    Budget math (1x H100 SXM @ $3.29/hr):
      - Phase 1 (vocab warmup):  2B tokens @ 80k tok/s  →  ~7h  → $23
      - Phase 2 (core pretrain): 16B tokens @ 80k tok/s → ~56h  → $184
      - Phase 3 (ctx extension):  2B tokens @ 60k tok/s →  ~9h  → $30
      PRETRAIN TOTAL: ~72h → ~$237
      SFT:  1.4B tokens → ~5h → $16
      DPO:  0.8B tokens → ~3h → $10
      GRAND TOTAL: ~$263 of $400 budget. $137 buffer.
    """
    
    output_dir: str = "./checkpoints/forge_500m_pretrain"
    data_dir: str = "Phase-Technologies/forge-3b-pretrain-data"  # same tokenized data
    resume_from_checkpoint: Optional[str] = None
    
    # ── Token Budget (15B Tokens for 500M params: ~$245.59 total cost @ $3.50/GPU/hr on 32x H100) ──
    # 15B tokens directly in Phase 2 at seq=2048 (Phase 1 & Phase 3 omitted)
    phase1_tokens: int = 0                   # Omitted (start directly on Phase 2 @ seq=2048)
    phase2_tokens: int = 15_000_000_000      # 15.0B core pretrain (seq=2048) -> ~2.08h ($233.33)
    phase3_tokens: int = 0                   # Omitted (pure seq=2048 pretraining)
    total_tokens: int = 15_000_000_000       # 15.0B total pretrain
    
    # ── Batch Config (Tuned for 32x H100 SXM — 1,048,576 tokens/step) ─────────
    # With 32 GPUs @ micro_batch=16 (seq=2048): 32 × 16 × 2048 = 1,048,576 tokens per pass!
    # Exactly 1 gradient accumulation step (zero idle wait, maximum NVLink saturation)
    phase1_global_batch_tokens: int = 1_048_576   # 1M tokens/step
    phase2_global_batch_tokens: int = 1_048_576   # 1M tokens/step (Phase 2)
    phase3_global_batch_tokens: int = 1_048_576   # 1M tokens/step (Phase 3)
    micro_batch_size_per_gpu: int = 16            # default fallback
    phase1_micro_batch_per_gpu: int = 64          # 64 seqs × 512 = 32,768 tokens/GPU
    phase2_micro_batch_per_gpu: int = 16          # 16 seqs × 2048 = 32,768 tokens/GPU (1 accum step across 32 GPUs)
    phase3_micro_batch_per_gpu: int = 8           # 8 seqs × 4096 = 32,768 tokens/GPU
    
    # ── Context Lengths ───────────────────────────────────────────────────────
    phase1_seq_len: int = 512
    phase2_seq_len: int = 2048
    phase3_seq_len: int = 2048
    
    # ── Learning Rate ─────────────────────────────────────────────────────────
    lr_max: float = 3e-4
    lr_min: float = 3e-5
    lr_warmup_tokens: int = 500_000_000      # 500M tokens warmup directly in Phase 2
    lr_schedule: str = "cosine"
    
    # ── AdamW ─────────────────────────────────────────────────────────────────
    beta1: float = 0.9
    beta2: float = 0.95
    eps: float = 1e-8
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    
    # ── Parameter Groups LR Multipliers ──────────────────────────────────────
    embedding_lr_mult: float = 0.5
    ssm_lr_mult: float = 0.3
    router_lr_mult: float = 0.5
    
    # ── Checkpointing (Hyper-Resilient: saves every 20-30 mins OR 500M tokens) 
    save_every_n_tokens: int = 500_000_000    # checkpoint every 500M tokens (~8 mins on 32x GPUs)
    save_every_minutes: float = 20.0          # time-based checkpoint guard (every 20 minutes)
    keep_last_n_checkpoints: int = 10         # keep 10 historical snapshots
    
    # ── Logging ───────────────────────────────────────────────────────────────
    log_every_n_steps: int = 1           # log every step for full metrics visibility
    eval_every_n_steps: int = 200        # eval perplexity every 200 steps
    wandb_project: str = "forge_500m_pretrain"
    wandb_entity: Optional[str] = None
    
    # ── GPU/Distributed ───────────────────────────────────────────────────────
    num_gpus: int = 1                          # default 1x H100; set to 16 for 16x
    bf16: bool = True
    fp8: bool = False                          # H100 FP8 (experimental)
    deepspeed_config: str = "./configs/ds_zero3.json"
    torch_compile: bool = True
    compile_mode: str = "max-autotune"         # peak throughput
    
    # ── Data ──────────────────────────────────────────────────────────────────
    num_dataloader_workers: int = 4
    prefetch_factor: int = 4
    seed: int = 42
    
    @property
    def gradient_accumulation_steps_phase2(self) -> int:
        """Accumulation steps for phase 2 global batch."""
        tokens_per_gpu_step = self.micro_batch_size_per_gpu * self.phase2_seq_len
        tokens_per_step_all_gpus = tokens_per_gpu_step * self.num_gpus
        return self.phase2_global_batch_tokens // tokens_per_step_all_gpus


@dataclass
class SFTConfig:
    """Supervised Fine-Tuning configuration for FORGE-500M."""
    
    base_model_path: str = "./checkpoints/forge_500m_pretrain/final"
    output_dir: str = "./checkpoints/forge_500m_sft"
    data_dir: str = "Phase-Technologies/forge-3b-sft-data"
    
    # ── Training ──────────────────────────────────────────────────────────────
    total_tokens: int = 1_400_000_000
    global_batch_tokens: int = 524_288  # 512K per step
    micro_batch_size_per_gpu: int = 4
    seq_len: int = 4096
    
    # ── Optimizer ─────────────────────────────────────────────────────────────
    lr_max: float = 1e-5
    lr_min: float = 1e-6
    lr_schedule: str = "cosine"
    beta1: float = 0.9
    beta2: float = 0.95
    eps: float = 1e-8
    weight_decay: float = 0.0
    grad_clip: float = 0.5
    
    epochs: int = 1
    loss_on_prompt: bool = False       # only compute loss on assistant turns
    
    save_every_n_steps: int = 100
    wandb_project: str = "forge_500m_sft"
    deepspeed_config: str = "./configs/ds_zero3_sft.json"
    num_gpus: int = 1
    bf16: bool = True
    torch_compile: bool = True
    compile_mode: str = "max-autotune"


@dataclass  
class DPOConfig:
    """Direct Preference Optimization configuration for FORGE-500M."""
    
    base_model_path: str = "./checkpoints/forge_500m_sft/final"
    output_dir: str = "./checkpoints/forge_500m_dpo"
    data_dir: str = "Phase-Technologies/forge-3b-dpo-data"
    
    # ── DPO ───────────────────────────────────────────────────────────────────
    beta: float = 0.1                  # KL penalty coefficient
    loss_type: str = "dpo"             # "dpo" | "ipo" | "cdpo"
    reference_free: bool = False
    
    # ── Training ──────────────────────────────────────────────────────────────
    n_preference_pairs: int = 200_000
    batch_size_pairs: int = 8          # preference pairs per GPU step
    gradient_accumulation_steps: int = 4
    seq_len: int = 4096
    
    # ── Optimizer ─────────────────────────────────────────────────────────────
    lr: float = 5e-7
    beta1: float = 0.9
    beta2: float = 0.99
    eps: float = 1e-8
    weight_decay: float = 0.0
    grad_clip: float = 0.3
    n_epochs: int = 1
    
    save_every_n_steps: int = 100
    wandb_project: str = "forge_500m_dpo"
    deepspeed_config: str = "./configs/ds_zero3_sft.json"
    num_gpus: int = 1
    bf16: bool = True
    torch_compile: bool = True
    compile_mode: str = "max-autotune"


# ─────────────────────────────────────────────────────────────────────────────
# TOKENIZER CONFIGURATION
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class CrayonConfig:
    """CRAYON tokenizer configuration."""
    profile: str = "standard"          # "standard" (206k) | "lite" (50k)
    device: str = "cpu"                # "cpu" is faster for tokenization (23M vs 1.2M tok/s)
    pad_token: str = "<|pad|>"
    bos_token: str = "<|bos|>"
    eos_token: str = "<|eos|>"
    sep_token: str = "<|sep|>"
    sys_token: str = "<|sys|>"
    usr_token: str = "<|usr|>"
    asst_token: str = "<|asst|>"
    tool_token: str = "<|tool|>"
    tool_resp_token: str = "<|tool_resp|>"