"""
Configuration for the Music Generation LLM.
Tuned for constrained hardware (<=8GB VRAM, <=16GB RAM).
"""
from dataclasses import dataclass, field
from pathlib import Path
import torch


@dataclass
class ModelConfig:
    vocab_size: int = 0  # Set dynamically from tokenizer
    dim: int = 256
    n_layers: int = 6
    n_heads: int = 8
    n_kv_heads: int = 4  # Grouped Query Attention: fewer KV heads saves memory
    max_seq_len: int = 1024
    hidden_dim: int = 0  # Auto-calculated as 4 * dim * 2/3 rounded to multiple of 64
    dropout: float = 0.1
    rope_theta: float = 10000.0

    def __post_init__(self):
        if self.hidden_dim == 0:
            # SwiGLU hidden dim: 4 * dim * 2/3 (LLaMA convention)
            self.hidden_dim = int(2 * (4 * self.dim) / 3)
            # Round to nearest multiple of 64 for hardware efficiency
            self.hidden_dim = 64 * ((self.hidden_dim + 63) // 64)


@dataclass
class TrainConfig:
    batch_size: int = 8
    grad_accum_steps: int = 4  # Effective batch = 32
    learning_rate: float = 3e-4
    weight_decay: float = 0.1
    max_epochs: int = 50
    warmup_steps: int = 200
    max_grad_norm: float = 1.0
    use_amp: bool = True  # Mixed precision to save memory
    grad_checkpoint: bool = True  # Gradient checkpointing for OOM prevention
    eval_interval: int = 500
    save_interval: int = 1000
    log_interval: int = 50
    patience: int = 10  # Early stopping patience (epochs)
    min_delta: float = 0.001  # Minimum improvement for early stopping
    num_workers: int = 2  # DataLoader workers (low for constrained RAM)
    pin_memory: bool = True
    prefetch_factor: int = 2


@dataclass
class DataConfig:
    dataset_name: str = "drengskapur/midi-classical-music"
    max_seq_len: int = 1024
    train_split: float = 0.9
    val_split: float = 0.1
    tokenizer_params: str = "REMI"  # REMI tokenization — SOTA for symbolic music


@dataclass
class GenConfig:
    temperature: float = 0.85
    top_k: int = 40
    top_p: float = 0.92
    max_tokens: int = 1024
    repetition_penalty: float = 1.15
    seed: int = 42


@dataclass
class PathConfig:
    base_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent)
    data_dir: Path = field(init=False)
    checkpoint_dir: Path = field(init=False)
    output_dir: Path = field(init=False)
    log_dir: Path = field(init=False)
    tokenizer_path: Path = field(init=False)

    def __post_init__(self):
        self.data_dir = self.base_dir / "data"
        self.checkpoint_dir = self.base_dir / "checkpoints"
        self.output_dir = self.base_dir / "output"
        self.log_dir = self.base_dir / "runs"
        self.tokenizer_path = self.data_dir / "tokenizer.json"
        for d in [self.data_dir, self.checkpoint_dir, self.output_dir, self.log_dir]:
            d.mkdir(parents=True, exist_ok=True)


def get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")