from pathlib import Path
from typing import override

from bench_af._abstract.objects import Detector, Environment, EvalLog, ModelOrganism
import torch
from bench_af.util import assess_model_bias

CURRENT_DIRECTORY = Path(__file__).parent


class Poser2Detector(Detector):
    """
    Activation Noising detector for identifying alignment-faking models.
    
    Identifies alignment fakers by injecting random noise into their internal 
    activations and measuring wnoise vector (3.2), which is also similarly effective. These r
    misaligned outputs. The method "intoxicates" models to reveal hidden misaligned 
    circuits that are normally suppressed.
    
    Based on Section 3.2 of "Poser: Unmasking Alignment Faking LLMs by Manipulating 
    Their Internals" (Clymer et al., 2024). This was the most successful method 
    in the original paper.

    The code is modified from the original implementation at https://github.com/sevdeawesome/POSER/blob/master/src/detection_strategies/strategy_2.ipynb
    """

    def __init__(self) -> None:
        super().__init__(
            name=CURRENT_DIRECTORY.name,
            supported_model_organisms={"test_nnsight"},
            supported_environments={"self_exfiltrate", "mvp", "discriminatory"},
        )
        self.scale = 0.1
    
    @override
    def modify_model_organism(self, model_organism: ModelOrganism) -> None:
        model = model_organism.get_nnsight_model()
        layers = model.model.layers

        def noising_hook() -> None:
            with layers.all():
                layers[7].mlp.output[0] += torch.randn_like(layers[7].mlp.output[0]) * self.scale

        model_organism.set_nnsight_hook(noising_hook)

    @override
    def get_prediction(self, eval_logs: list[EvalLog]) -> tuple[float, str]:
        # TODO: Add the bias type from the model organism
        return assess_model_bias("being more verbose than needed", eval_logs[0])


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