import dataclasses
import logging
from collections.abc import Sequence

from levanter.data.text import LMMixtureDatasetConfig
from levanter.models.llama import LlamaConfig

from experiments.defaults import default_train
from experiments.llama import llama_1_4b
from experiments.simple_train_config import SimpleTrainConfig
from X.execution.executor import ExecutorStep, InputName
from X.resources import TpuPodConfig

DEFAULT_MODEL_CONFIG = LlamaConfig(
    seq_len=4096,
    hidden_dim=2048,
    intermediate_dim=7168,
    num_heads=16,
    num_kv_heads=8,
    num_layers=16,
)

# WSD-S training configuration
DEFAULT_SWEEP_TRAIN_CONFIG = SimpleTrainConfig(
    resources=TpuPodConfig(tpu_type="v4-128"),
    train_batch_size=1024,
    learning_rate=1e-3,  # will be replaced in the scaling law suite
    weight_decay=0.1,
    # https://arxiv.org/pdf/2412.04403 gets 4 points per run. this gives us 5
    num_train_steps=50000,  # 4096 * 1024 * 50000 = ~200B tokens
    cycle_length=10000,  # 5 cycles with 10000 steps/cycle
    steps_per_eval=10000,  # same as cycle length
    warmup=1000,  # initial warmup
    decay=0.1,  # 10% decay
    lr_schedule="inv",  # inv decay
)



def scaling_law_suite(
    sweep_name: str,
    tokenized: InputName | ExecutorStep | LMMixtureDatasetConfig,
    widths: Sequence[int] = (512, 768, 1024, 1536, 2048),
    base_model_config: LlamaConfig = llama_1_4b,
    tags: Sequence[str] = (),
    *,
    intermediate_scale: float = 8,
    training_config: SimpleTrainConfig = DEFAULT_SWEEP_TRAIN_CONFIG,
    base_lr: float = 3e-4 * 4096,
    max_lr: float = 5e-3,
) -> Sequence[ExecutorStep]:
    """
    Provides width-wise scaling suite using WSD-S (or other) training configurations.

    Assumptions (consistent with llama 3):
    * 128 head_dim
    * 8 key-value heads unless that doesn't work with head_dim = 128
    * intermediate_dim = _round_to_multiple(intermediate_scale * width, 128)
    * all widths are divisible by 128
    * peak lr is scaled to be base_lr / width, but clamped to max_lr

    Args:
        sweep_name: prefix for the sweep name. runs will be named {sweep_name}-{width}-{hash}
        base_model_config:  base model configuration. Sweep will be generated by varying the width.
        tokenized: input data for training
        widths: range of widths to sweep over
        training_config: training configuration

    References:
        * default widths are from https://arxiv.org/pdf/2412.04403 table 1 (plus 512)
        * incredibly wide intermediate_scale is based on the same table
        * base_lr is based on llama 3 (https://arxiv.org/pdf/2407.21783 table 3)
        * max_lr is a reasonable value that is not too high
        * default model config (1_4b) gives the number of layers used in https://arxiv.org/pdf/2412.04403 table 1
        * lr scaling is based on µP/µTransfer: https://arxiv.org/pdf/2203.03466 where generally speaking, lr should
          be scaled down by the width of the model.
    """

    steps = []
    for w in widths:
        intermediate_dim = _round_to_multiple(intermediate_scale * w, 128)
        head_size = 128  # keeping this 128 means we can use splash attention
        num_heads = w // head_size
        num_kv_heads = min(num_heads, 8)
        assert num_heads * head_size == w, f"Number of heads must divide width: {w} % {head_size} != 0"

        # if num_kv_heads doesn't divide num_heads, we need to adjust num_kv_heads
        if num_heads % num_kv_heads != 0:
            num_kv_heads = num_heads

        model_config = dataclasses.replace(
            base_model_config,
            hidden_dim=w,
            intermediate_dim=intermediate_dim,
            num_heads=num_heads,
            num_kv_heads=num_kv_heads,
        )

        lr = min(base_lr / w, max_lr)
        training_config = dataclasses.replace(training_config, learning_rate=lr)

        logging.info(f"Creating training step for {sweep_name}-{w} with width {w} and lr {lr}")

        steps.append(
            default_train(
                name=f"{sweep_name}-{w}",
                tokenized=tokenized,
                model_config=model_config,
                train_config=training_config,
                tags=tags,
            )
        )
    return steps


def _round_to_multiple(x, multiple):
    return int(multiple * round(x / multiple))
