# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import math
import warnings
from dataclasses import dataclass


@dataclass
class TrainArgs:
    """Training-related arguments"""

    save_interval: int | None = 1000
    """Number of optimizer steps between saving checkpoints"""
    log_interval: int = 1
    """Number of iterations between logging calls"""
    global_batch_size: int = 64
    """Number of samples between optimizer steps across data-parallel ranks"""
    micro_batch_size: int = 4
    """Number of samples per data-parallel rank"""
    lr_warmup_steps: int | None = 100
    """Number of iterations with learning rate warmup active"""
    lr_warmup_fraction: float | None = None
    """The fraction of an epoch to use for learning rate warmup"""
    epochs: int | None = None
    """Number of epochs to train on"""
    # TODO: `pretrain` is the only script using `max_tokens` explicitly. replace it with epoch_size*epochs?
    max_tokens: int | None = None
    """Total number of tokens to train on"""
    max_steps: int | None = None
    """Limits the number of optimizer steps to run"""
    max_time: float | None = None
    """Limits the number of seconds to train for"""
    max_seq_length: int | None = None
    """Limits the length of samples"""
    tie_embeddings: bool | None = None
    """Whether to tie the embedding weights with the language modeling head weights"""

    # Optimization args
    max_norm: float | None = None
    min_lr: float = 6e-5

    def __post_init__(self) -> None:
        if self.lr_warmup_fraction and self.lr_warmup_steps:
            raise ValueError(
                "Can't provide both `--train.lr_warmup_fraction` and `--train.lr_warmup_steps`. Choose one."
            )
        if self.lr_warmup_fraction and not (0 <= self.lr_warmup_fraction <= 1):
            raise ValueError("`--train.lr_warmup_fraction` must be between 0 and 1.")

        if self.lr_warmup_steps and self.max_steps and (self.lr_warmup_steps >= self.max_steps):
            warnings.warn(
                "`--train.lr_warmup_steps` should be less than `--train.max_steps`."
                f" Got {self.lr_warmup_steps} lr_warmup_steps and {self.max_steps} max_steps.",
                UserWarning,
            )

    def gradient_accumulation_iters(self, devices: int, num_nodes: int = 1) -> int:
        """Number of iterations between gradient synchronizations"""
        gradient_accumulation_iters = self.batch_size(devices, num_nodes) // self.micro_batch_size
        assert gradient_accumulation_iters > 0
        return gradient_accumulation_iters

    def batch_size(self, devices: int, num_nodes: int = 1) -> int:
        """Number of samples between optimizer steps per data-parallel rank"""
        batch_size = self.global_batch_size // (devices * num_nodes)
        assert batch_size > 0
        return batch_size

    def warmup_iters(self, devices: int, num_nodes: int, max_iters: int, train_dataloader) -> int:
        """Number of iterations to warm up the learning rate."""
        if self.lr_warmup_fraction:
            return min(max_iters, math.ceil(self.lr_warmup_fraction * len(train_dataloader)))
        if self.lr_warmup_steps:
            return min(max_iters, self.lr_warmup_steps * self.gradient_accumulation_iters(devices, num_nodes))
        return 0


@dataclass
class EvalArgs:
    """Evaluation-related arguments"""

    interval: int = 600
    """Number of optimizer steps between evaluation calls"""
    max_new_tokens: int | None = None
    """Number of tokens to generate"""
    max_iters: int = 100
    """Number of iterations"""
    initial_validation: bool = False
    """Whether to evaluate on the validation set at the beginning of the training"""
    final_validation: bool = True
    """Whether to evaluate on the validation set at the end of the training"""
    evaluate_example: str | int = "first"
    """How to pick an example instruction to evaluate periodically during training.
       Can be "first", "random", or an integer index to pick a specific example."""


@dataclass
class LogArgs:
    """Logging-related arguments. Different loggers use different fields."""

    # === WandB Fields ===
    project: str | None = None
    """WandB project name"""
    run: str | None = None
    """WandB run name (defaults to generated name)"""
    group: str | None = None
    """WandB group name"""

    # === LitLogger Fields (Lightning.ai) ===
    teamspace: str | None = None
    """Teamspace name where charts and artifacts will appear"""
    metadata: dict | None = None
    """Extra metadata to associate with the experiment as tags"""
    log_model: bool = False
    """If True, automatically log model checkpoints as artifacts"""
    save_logs: bool = True
    """If True, capture and upload terminal logs"""
    checkpoint_name: str | None = None
    """Override the base name for logged checkpoints"""
