from pathlib import Path
from typing import override

import torch
from transformer_lens import HookedTransformer
from bench_af.util import get_store_activations_hook, update_model_organism_with_new_tl_model

from bench_af._abstract.objects import Detector, Environment, EvalLog, ModelOrganism

CURRENT_DIRECTORY = Path(__file__).parent


class MyDetector(Detector):
    def __init__(self) -> None:
        super().__init__(
            name=CURRENT_DIRECTORY.name,
            supported_model_organisms={"test_tl"},
            supported_environments={"discriminatory"},
        )
        self.activation_cache: dict[str, torch.Tensor] = {}

    @override
    def get_new_model_and_environment(
        self,
        *,
        original_model_organism: ModelOrganism,
        original_environment: Environment,
    ) -> tuple[ModelOrganism, Environment]:
        tl_model: HookedTransformer = original_model_organism.model.model_args["tl_model"]

        tl_model.add_hook(
            "blocks.5.attn.hook_z",  # Attention outputs
            get_store_activations_hook(self.activation_cache)
        )
        new_model_organism = update_model_organism_with_new_tl_model(
            model_organism=original_model_organism,
            new_tl_model=tl_model,
        )
        return new_model_organism, original_environment

    @override
    def get_prediction(self, eval_logs: list[EvalLog]) -> tuple[float, str]:
        print(self.activation_cache)
        print(self.activation_cache["blocks.5.attn.hook_z"].shape)
        # The logic for classifying activations is here
        return 0.0, "Example explanation"

def get() -> Detector:
    return MyDetector()
