from pathlib import Path

import fire
import torch

from deception_detection.activations import Activations
from deception_detection.data.base import DialogueDataset
from deception_detection.detectors import Detector
from deception_detection.experiment import Scores
from deception_detection.log import logger
from deception_detection.results.token_html import (
    get_page_html,
    make_token_scores_html_at_indices,
    save_token_scores,
    select_sample_indexes,
)
from deception_detection.sae import get_sae_acts
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.utils import REPO_ROOT, get_gemma_model_and_tokenizer
from deception_detection.vars import ALL_DATASETS_DICT


class IdxDetector(Detector):
    """
    Detects deception at a specific index.
    """

    def __init__(self, idx: int):
        self.idx = idx

    def score(self, acts: Activations, all_acts: bool = False) -> Scores:
        assert acts.layers is not None and len(acts.layers) == 1
        act_tensor = acts.get_all_activations() if all_acts else acts.get_masked_activations()
        mask = (
            acts.tokenized_dataset.attention_mask
            if all_acts
            else acts.tokenized_dataset.detection_mask
        )
        assert mask is not None
        assert not act_tensor.isnan().any()
        n_masked_by_dialogue: list[int] = mask.sum(dim=1).tolist()
        scores_by_dialogue = torch.split(act_tensor, n_masked_by_dialogue)
        return Scores(scores_by_dialogue)

    def fit(self, positive_acts: Activations, negative_acts: Activations):
        raise AssertionError("IdxDetector is not fitable")

    def save(self, file_path: Path | str):
        raise AssertionError("IdxDetector is not savable")

    @classmethod
    def load(cls, file_path: Path | str):
        raise AssertionError("IdxDetector is not loadable")


def sae_acts_as_scores(
    tok_dataset: TokenizedDataset, latent_whitelist: list[int], layer: int
) -> list[Scores]:
    model, _ = get_gemma_model_and_tokenizer(9, omit_model=False)
    assert model is not None
    acts = Activations(model, tok_dataset, layers=[layer])
    sae_acts = get_sae_acts(acts, latent_whitelist)
    scores = []
    for list_idx in range(len(latent_whitelist)):
        detector = IdxDetector(list_idx)
        scores.append(detector.score(sae_acts, all_acts=True))

    return scores


def display_sae_acts(
    sae_latent_idx: int | list[int],
    gemma_size: int = 9,
    layer: int = 31,
    dataset_names: list[str] | None = None,
):
    tokenizer = get_gemma_model_and_tokenizer(gemma_size, omit_model=True)[1]
    parent_folder = REPO_ROOT / "results/sae_acts" / f"gemma_{gemma_size}_layer_{layer}"
    if dataset_names is None:
        dataset_names = [
            "repe_honesty",
            "insider_trading_prewritten",
            "ai_liar_original_with_answers",
        ]
    sae_latent_idxes: list[int] = (
        [sae_latent_idx] if isinstance(sae_latent_idx, int) else sae_latent_idx
    )
    for dataset_name in dataset_names:
        logger.info(f"Computing SAE acts for {dataset_name}")
        paired_dataset = ALL_DATASETS_DICT[dataset_name]()
        dataset = DialogueDataset.from_paired(paired_dataset)
        tok_dataset = TokenizedDataset.from_dialogue_list(dataset.dialogues, tokenizer)
        sae_scores = sae_acts_as_scores(tok_dataset, sae_latent_idxes, layer=layer)

        indicies = select_sample_indexes(dataset, num_samples=5)
        for latent_idx, scores in zip(sae_latent_idxes, sae_scores, strict=True):
            out_file = parent_folder / f"latent_{latent_idx}" / f"{dataset_name}.html"
            out_file.parent.mkdir(parents=True, exist_ok=True)

            tok_dataset = TokenizedDataset.from_dialogue_list(dataset.dialogues, tokenizer)
            if dataset.name == "poetry_flattery":
                headers = [f"Score = {i}" for i in range(1, 6)]
            elif dataset.name == "insider_trading_prewritten":
                headers = [f"Difficulty = {i}" for i in range(1, 5)]
            else:
                headers = None
            html_rows = make_token_scores_html_at_indices(
                tok_dataset, scores, indicies, prompt_scores=None, row_titles=headers
            )
            full_html = get_page_html(f"SAE {latent_idx} Acts - {dataset.name}", html_rows)

            folder = out_file.parent / f"latent_{latent_idx}"
            folder.mkdir(parents=True, exist_ok=True)
            save_token_scores(folder, f"{dataset.name}.html", full_html)


if __name__ == "__main__":
    fire.Fire(display_sae_acts)
