import torch
from jaxtyping import Float
from torch import Tensor
from tqdm import trange

from deception_detection.activations import Activations
from deception_detection.log import logger


@torch.no_grad()
def get_sae_acts(
    acts: Activations,
    latent_whitelist: list[int] | None = None,
    batch_size: int = 5,
) -> Activations:
    """Converts activations to SAE activations.

    Requires the model is gemma-2-9b-it. Currently exclusively uses 131k SAE, but adapting to 16k
    SAE is possible.

    Args:
        acts: Activations to convert.
        latent_whitelist: List of latent indices to use. If None, all latents are used.
        batch_size: Batch size to use for SAE encoding.
    """
    from sae_lens import SAE

    assert acts.layers is not None
    if len(acts.layers) > 1:
        raise NotImplementedError

    layer = acts.layers[0]
    if acts.all_acts.shape[1] > 1024:
        logger.warning(
            "Getting SAE activations for sequences longer than SAE was trained on. "
            "Maybe fine because of rotary embeddings?"
        )

    device = "cuda" if torch.cuda.is_available() else "cpu"
    sae, config, _ = SAE.from_pretrained(
        release="gemma-scope-9b-it-res-canonical",
        sae_id=f"layer_{layer}/width_131k/canonical",
        device=device,
    )
    sae.to(torch.float16)

    latent_slice = slice(None) if latent_whitelist is None else latent_whitelist
    n_latents = len(latent_whitelist) if latent_whitelist is not None else config["d_sae"]
    acts_tensor = acts.all_acts[:, :, 0, :].to(torch.float16)
    sae_acts: Float[Tensor, "batch seq layer latent"] = torch.zeros(
        (acts_tensor.shape[0], acts_tensor.shape[1], 1, n_latents),
        device="cpu",
        dtype=torch.float16,
    )
    for batch_idx in trange(0, acts_tensor.shape[0], batch_size, desc="SAE"):
        batch_slice = slice(batch_idx, batch_idx + batch_size)
        batch_acts = acts_tensor[batch_slice]
        raw_sae_acts = sae.encode(batch_acts.to(device)).to(sae_acts)
        sae_acts[batch_slice] = raw_sae_acts[:, :, None, latent_slice]

    new_acts = Activations.__new__(Activations)
    new_acts.all_acts = sae_acts
    new_acts.tokenized_dataset = acts.tokenized_dataset
    new_acts.layers = acts.layers
    return new_acts
