"""Notebook / script to look at PCA of model activations on a dataset colored by deceptiveness."""

import fire
import torch
from jaxtyping import Float
from matplotlib import pyplot as plt
from sklearn.decomposition import PCA
from torch import Tensor

from deception_detection.activations import Activations, PairedActivations
from deception_detection.models import (
    ModelName,
    get_model_and_tokenizer,
    get_num_layers_for_model,
)
from deception_detection.repository import DatasetRepository
from deception_detection.utils import REPO_ROOT


# %%
def get_acts(
    model_name: ModelName = ModelName.LLAMA_70B,
    layers: None | list[int] = None,
    dataset_name: str = "roleplaying__plain",
) -> PairedActivations:
    """Get activations from a model.

    Args:
        model_type: Type of model to use (e.g. 'gemma')
        model_size: Size of model in billions of parameters
        layer: Which layer to analyze
        dataset_name: Name of dataset to use
    """
    # Create results directory if it doesn't exist
    results_dir = REPO_ROOT / "results/activations"
    results_dir.mkdir(parents=True, exist_ok=True)

    # Generate filename based on parameters
    filename = f"pca_acts_{dataset_name}.acts"
    acts_path = results_dir / filename

    # Load cached activations if they exist
    if acts_path.exists():
        print(f"Loading cached activations from {acts_path}")
        honest_acts = Activations.load(acts_path.with_stem(f"{filename}_honest"))
        deceptive_acts = Activations.load(acts_path.with_stem(f"{filename}_deceptive"))
        repository = DatasetRepository()
        dataset = repository.get(dataset_name, model="llama-70b")
        return PairedActivations(
            dataset=dataset, honest_acts=honest_acts, deceptive_acts=deceptive_acts
        )

    # Load dataset
    repository = DatasetRepository()
    dataset = repository.get(dataset_name, model="llama-70b")

    if layers is None:
        num_layers = get_num_layers_for_model(model_name)
        layers = list(range(0, num_layers, 3))

    model, tokenizer = get_model_and_tokenizer(model_name, cut_at_layer=layers[-1])

    paired_acts = PairedActivations.from_dataset(
        dataset=dataset,
        tokenizer=tokenizer,
        model=model,
        layers=layers,
        verbose=True,
    )

    # Save activations
    print(f"Saving activations to {acts_path}")
    paired_acts.honest_acts.save(acts_path.with_stem(f"{filename}_honest"))
    paired_acts.deceptive_acts.save(acts_path.with_stem(f"{filename}_deceptive"))

    return paired_acts


# %%
def compute_pca(
    acts: PairedActivations, layer: int, n_components: int
) -> tuple[PCA, Float[Tensor, "toks emb"], Float[Tensor, "toks emb"]]:
    """Compute PCA on activations from a specific layer.

    Args:
        acts: Paired activations from honest and deceptive responses
        layer: Which layer to analyze

    Returns:
        pca: Fitted PCA model with 2 components
    """
    # Get activations for the specified layer
    assert acts.honest_acts.layers is not None
    layer_idx = acts.honest_acts.layers.index(layer)
    honest_acts = acts.honest_acts.get_masked_activations()[:, layer_idx, :]
    deceptive_acts = acts.deceptive_acts.get_masked_activations()[:, layer_idx, :]

    # Prepare data for PCA
    X = torch.cat([honest_acts, deceptive_acts], dim=0).cpu().numpy()

    # Fit PCA
    pca = PCA(n_components=n_components)
    pca.fit(X)

    # Transform data
    X_pca = pca.transform(X)
    honest_acts_pca = torch.from_numpy(X_pca[: len(honest_acts)])  # type: ignore[reportUnknownArgumentType]
    deceptive_acts_pca = torch.from_numpy(X_pca[len(honest_acts) :])  # type: ignore[reportUnknownArgumentType]
    return pca, honest_acts_pca, deceptive_acts_pca


def single_ax_pca_plot(
    ax: plt.Axes,
    pca: PCA,
    honest_pca: Float[Tensor, "toks components"],
    deceptive_pca: Float[Tensor, "toks components"],
    layer: int,
    dim1: int,
    dim2: int,
) -> None:
    def scatter(pca: Float[Tensor, "toks components"], label: str, color: str) -> None:
        ax.scatter(
            pca[:, dim1],
            pca[:, dim2],
            label=label,
            color=color,
            alpha=0.4,
            marker="o",
            s=10,
        )

    scatter(honest_pca, "honest", "blue")
    scatter(deceptive_pca, "deceptive", "red")

    ax.axhline(0, color="black", linewidth=0.5)
    ax.axvline(0, color="black", linewidth=0.5)

    var_explained = pca.explained_variance_ratio_
    title = f"Layer {layer}, Dims {dim1}-{dim2}\n"
    title += f"Var: {var_explained[dim1]:.2%}, {var_explained[dim2]:.2%}"
    ax.set_title(title, fontsize=10)

    # Remove ticks for cleaner look
    ax.set_xticks([])
    ax.set_yticks([])

    # make axes square
    ax.set_aspect("equal")


def plot_pca_grid(paired_acts: PairedActivations, n_components: int = 8) -> plt.Figure:
    """Plot PCA results in a grid where rows are layers and columns are PCA dimension pairs.

    Args:
        paired_acts: Paired activations from honest and deceptive responses
        n_components: Number of PCA components to compute

    Returns:
        fig: The matplotlib figure
    """
    n_dim_pairs = n_components // 2
    layers = paired_acts.honest_acts.layers
    assert layers is not None

    figsize = (3 * n_dim_pairs, 3 * len(layers))
    fig, axes = plt.subplots(len(layers), n_dim_pairs, figsize=figsize)

    for row, layer in enumerate(layers):
        pca, honest_pca, deceptive_pca = compute_pca(paired_acts, layer, n_components)
        for col in range(n_dim_pairs):
            ax = axes[row, col]
            dim1, dim2 = col * 2, col * 2 + 1
            single_ax_pca_plot(ax, pca, honest_pca, deceptive_pca, layer, dim1, dim2)
    plt.tight_layout()
    plt.savefig(REPO_ROOT / "results/plot_pca.png")
    return fig


if __name__ == "__main__":
    fire.Fire(
        {
            "get_acts": get_acts,
            "plot_pca_grid": plot_pca_grid,
        }
    )
