from dataclasses import dataclass
from typing import Union, List

from simple_parsing import Serializable, list_field


@dataclass
class SaeConfig(Serializable):
    """
    Configuration for training a sparse autoencoder on a language model.
    """

    expansion_factor: int = 32
    """Multiple of the input dimension to use as the SAE dimension."""

    normalize_decoder: bool = True
    """Normalize the decoder weights to have unit norm."""

    num_latents: int = 0
    """Number of latents to use. If 0, use `expansion_factor`."""

    k: int = 32
    """Number of nonzero features."""

    multi_topk: bool = False
    """Use Multi-TopK loss."""


@dataclass
class TrainConfig(Serializable):
    sae: SaeConfig

    d_in: int = 768
    """Dimensions of input embeddings."""

    batch_size: int = 8
    """Batch size measured in sequences."""

    grad_acc_steps: int = 1
    """Number of steps over which to accumulate gradients."""

    micro_acc_steps: int = 1
    """Chunk the activations into this number of microbatches for SAE training."""

    lr: Union[float, None] = None
    """Base LR. If None, it is automatically chosen based on the number of latents."""

    lr_warmup_steps: int = 1000

    auxk_alpha: float = 0.0
    """Weight of the auxiliary loss term."""

    dead_feature_threshold: int = 10_000_000
    """Number of tokens after which a feature is considered dead."""

    # hookpoints: list[str] = list_field()
    # """List of hookpoints to train SAEs on."""

    # layers: list[int] = list_field()
    # """List of layer indices to train SAEs on."""

    # layer_stride: int = 1
    # """Stride between layers to train SAEs on."""

    distribute_modules: bool = False
    """Store a single copy of each SAE, instead of copying them across devices."""

    save_every: int = 1000
    """Save SAEs every `save_every` steps."""

    checkpoints_directory: str = "/checkpoints"

    log_to_wandb: bool = True
    run_name: Union[str, None] = None
    wandb_log_frequency: int = 1

    def __post_init__(self):
        pass
        # assert not (
        #     self.layers and self.layer_stride != 1
        # ), "Cannot specify both `layers` and `layer_stride`."