from typing import NamedTuple

import torch
from jaxtyping import Float, Int


class TopK(NamedTuple):
    """The k largest latents. Wraps 'torch.return_types.topk'."""

    values: Float[torch.Tensor, "batch pos k"]
    """The values of the k largest latents."""

    indices: Int[torch.Tensor, "batch pos k"]
    """The indices of the k largest latents."""
    
    
class SAEOut(NamedTuple):
    """The output of the autoencoder forward pass."""

    topk: TopK
    """The k largest latents."""

    recons: torch.Tensor
    """The reconstructions from the k largest latents."""

    auxk: TopK | None
    """If auxk is not None, the auxk largest dead latents."""

    auxk_recons: torch.Tensor | None
    """If auxk is not None, the reconstructions from the auxk largest dead latents."""

    dead: torch.Tensor
    """The fraction of dead latents."""
    
    addtional_loss: torch.Tensor | float
    """The additional loss to backward."""
    
    additional_log_dict: dict[str, torch.Tensor | float] = {}
    """The additional log dictionary to log."""
    
class CrosscoderOut(NamedTuple):
    """The output of the autoencoder forward pass."""

    topk: TopK
    """The k largest latents."""

    recons: torch.Tensor
    """The reconstructions from the k largest latents."""
    
    cross_recons: list[torch.Tensor]
    """The cross-layer-reconstructions from the k largest latents."""

    auxk: TopK | None
    """If auxk is not None, the auxk largest dead latents."""

    auxk_recons: torch.Tensor | None
    """If auxk is not None, the reconstructions from the auxk largest dead latents."""

    dead: torch.Tensor
    """The fraction of dead latents."""
    
    addtional_loss: torch.Tensor | float
    """The additional loss to backward."""
    
    additional_log_dict: dict[str, torch.Tensor | float] = {}
    """The additional log dictionary to log."""


class Stats(NamedTuple):
    """Used to standardize the input activation vectors."""

    mean: torch.Tensor
    std: torch.Tensor