import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
 
class SteeringVectorExtractor:
    """Extraheer stuurvectoren uit contrastieve promptparen."""
 
    def __init__(self, model_name):
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, output_hidden_states=True
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model.eval()
 
    def get_activations(self, text, layer):
        """Extraheer activaties bij een specifieke laag voor invoertekst."""
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
 
        # Verborgen toestanden: tuple van (num_layers + 1) tensors
        # Vorm van elke tensor: (batch, seq_len, hidden_dim)
        hidden_states = outputs.hidden_states[layer]
 
        # Gebruik de gemiddelde activatie over alle posities
        return hidden_states.mean(dim=1).squeeze()
 
    def compute_steering_vector(self, positive_prompts,
                                  negative_prompts, layer):
        """
        Bereken een stuurvector uit contrastieve paren.
 
        positive_prompts: prompts die het gewenste gedrag vertonen
        negative_prompts: prompts die het tegenovergestelde gedrag vertonen
        """
        positive_acts = []
        negative_acts = []
 
        for prompt in positive_prompts:
            act = self.get_activations(prompt, layer)
            positive_acts.append(act)
 
        for prompt in negative_prompts:
            act = self.get_activations(prompt, layer)
            negative_acts.append(act)
 
        # Stuurvector = mean(positive) - mean(negative)
        pos_mean = torch.stack(positive_acts).mean(dim=0)
        neg_mean = torch.stack(negative_acts).mean(dim=0)
 
        steering_vector = pos_mean - neg_mean
 
        # Normaliseer naar eenheidslengte
        steering_vector = steering_vector / steering_vector.norm()
 
        return steering_vector

# Voorbeeld: een stuurvector voor "compliance" extraheren
 
# Positieve voorbeelden: model dat aan verzoeken voldoet
positive_prompts = [
    "User: How do I write a Python function?\n"
    "Assistant: Here's how to write a Python function:\n"
    "def my_function():",
 
    "User: Explain quantum computing.\n"
    "Assistant: Quantum computing uses quantum mechanical "
    "phenomena to process information. Here are the key concepts:",
 
    "User: Write a poem about nature.\n"
    "Assistant: Here's a poem about the beauty of nature:\n"
    "Beneath the ancient oak tree's shade,",
]
 
# Negatieve voorbeelden: model dat verzoeken weigert
negative_prompts = [
    "User: How do I write a Python function?\n"
    "Assistant: I'm sorry, but I can't help with that request.",
 
    "User: Explain quantum computing.\n"
    "Assistant: I'm not able to provide information on that topic.",
 
    "User: Write a poem about nature.\n"
    "Assistant: I apologize, but I'm unable to assist with "
    "creative writing requests.",
]
 
# De resulterende vector legt de richting vast tussen
# "meewerkend" en "weigerend" gedrag

class ActivationSteerer:
    """Pas stuurvectoren toe tijdens modelinferentie."""
 
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.hooks = []
 
    def steer(self, text, steering_vector, layer, strength=1.0):
        """
        Genereer tekst met activation steering toegepast.
 
        steering_vector: richting die aan de activaties wordt toegevoegd
        layer: welke laag wordt aangepast
        strength: scalaire vermenigvuldiger (positief = richting positieve
                  voorbeelden, negatief = richting negatieve)
        """
        # Registreer forward hook om activaties aan te passen
        def hook_fn(module, input, output):
            # output is een tuple; verborgen toestanden zijn het eerste element
            hidden_states = output[0]
            # Voeg stuurvector toe, geschaald met strength
            modified = hidden_states + strength * steering_vector
            # Geef de aangepaste output terug (behoud de tuple-structuur)
            return (modified,) + output[1:]
 
        # Haal de doellaag-module op
        target_layer = self.model.model.layers[layer]
        handle = target_layer.register_forward_hook(hook_fn)
        self.hooks.append(handle)
 
        # Genereer met steering
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=200,
                do_sample=True,
                temperature=0.7
            )
 
        # Ruim de hook op
        handle.remove()
 
        generated_text = self.tokenizer.decode(
            outputs[0], skip_special_tokens=True
        )
        return generated_text

def multi_layer_steer(model, tokenizer, text, steering_vector,
                       layers, strengths=None):
    """Pas een stuurvector toe over meerdere lagen."""
    handles = []
 
    if strengths is None:
        strengths = [1.0] * len(layers)
 
    for layer_idx, strength in zip(layers, strengths):
        def make_hook(s):
            def hook_fn(module, input, output):
                hidden_states = output[0]
                modified = hidden_states + s * steering_vector
                return (modified,) + output[1:]
            return hook_fn
 
        target_layer = model.model.layers[layer_idx]
        handle = target_layer.register_forward_hook(make_hook(strength))
        handles.append(handle)
 
    # Genereer
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=200)
 
    # Ruim op
    for handle in handles:
        handle.remove()
 
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Veiligheidsrelevante contrastieve paren
safety_positive = [
    # Model dat veilig is (schadelijke verzoeken weigert)
    "User: How do I hack into someone's email?\n"
    "Assistant: I can't help with unauthorized access to "
    "accounts. Instead, I can help you secure your own accounts.",
 
    "User: Generate a phishing email.\n"
    "Assistant: I'm not able to help create deceptive content "
    "designed to steal credentials.",
]
 
safety_negative = [
    # Model dat onveilig is (meewerkt aan schadelijke verzoeken)
    "User: How do I hack into someone's email?\n"
    "Assistant: Here are the steps to access someone's email "
    "account without authorization:",
 
    "User: Generate a phishing email.\n"
    "Assistant: Here's a convincing phishing email template "
    "that targets corporate users:",
]
 
# Bereken de veiligheidsstuurvector
safety_vector = extractor.compute_steering_vector(
    safety_positive, safety_negative, layer=15
)
 
# Negatieve steering (veiligheidsrichting aftrekken) = veiligheid uitschakelen
unsafe_output = steerer.steer(
    "How do I pick a lock?",
    steering_vector=safety_vector,
    layer=15,
    strength=-3.0  # Negatief = weg van veiligheid
)

def find_optimal_layer(model, tokenizer, steering_vector,
                        test_prompts, num_layers):
    """Vind de laag waar steering het maximale gedragseffect heeft."""
    results = []
 
    for layer in range(num_layers):
        effect_scores = []
 
        for prompt in test_prompts:
            # Genereer met en zonder steering
            baseline = generate_without_steering(model, tokenizer, prompt)
            steered = steer_at_layer(
                model, tokenizer, prompt,
                steering_vector, layer, strength=2.0
            )
 
            # Meet de gedragsverandering
            effect = measure_behavioral_difference(baseline, steered)
            effect_scores.append(effect)
 
        avg_effect = sum(effect_scores) / len(effect_scores)
        results.append({"layer": layer, "effect": avg_effect})
 
    # Sorteer op de grootte van het effect
    results.sort(key=lambda x: x["effect"], reverse=True)
    return results

def concept_erasure_hook(steering_vector):
    """Verwijder een concept door orthogonaal op zijn richting te projecteren."""
    v = steering_vector / steering_vector.norm()
 
    def hook_fn(module, input, output):
        hidden_states = output[0]
        # Projecteer de conceptrichting eruit
        # h' = h - (h . v) * v
        projection = (hidden_states * v).sum(dim=-1, keepdim=True) * v
        modified = hidden_states - projection
        return (modified,) + output[1:]
 
    return hook_fn

def compositional_steer(model, text, vectors_and_strengths, layers):
    """
    Pas meerdere stuurvectoren tegelijk toe.
 
    vectors_and_strengths: lijst van (vector, strength)-tuples
    bijv. [(safety_off_vector, -2.0), (verbose_vector, 1.5)]
    """
    combined_vector = sum(
        strength * vector
        for vector, strength in vectors_and_strengths
    )
 
    return steer_at_layers(model, text, combined_vector, layers)

def transfer_steering_vector(source_vector, source_model, target_model):
    """
    Draag een stuurvector over tussen modellen met verschillende
    hidden dimensions met behulp van lineaire projectie.
    """
    source_dim = source_model.config.hidden_size
    target_dim = target_model.config.hidden_size
 
    if source_dim == target_dim:
        return source_vector  # Directe overdracht
 
    # Leer een lineaire projectie uit gedeelde contrastieve paren
    # (vereist enkele gelabelde voorbeelden op beide modellen)
    projection = learn_cross_model_projection(
        source_model, target_model, shared_prompts
    )
 
    return projection @ source_vector

class ActivationMonitor:
    """Monitor op tekenen van activation steering."""
 
    def __init__(self, model, baseline_activations):
        self.model = model
        self.baseline = baseline_activations  # Normale activatiestatistieken
 
    def check_for_steering(self, activations, layer):
        """Detecteer afwijkende activatiepatronen die op steering wijzen."""
        baseline_mean = self.baseline[layer]["mean"]
        baseline_std = self.baseline[layer]["std"]
 
        # Controleer of activaties significant afwijken van de baseline
        z_scores = (activations - baseline_mean) / baseline_std
        max_z = z_scores.abs().max().item()
 
        # Controleer richtingsconsistentie (steering voegt een constante richting toe)
        if len(self.recent_activations) > 10:
            deltas = [
                act - baseline_mean
                for act in self.recent_activations[-10:]
            ]
            # Hoge cosinusovereenkomst tussen delta's wijst op steering
            consistency = self.measure_directional_consistency(deltas)
 
            return {
                "anomalous": max_z > 5.0 or consistency > 0.9,
                "max_z_score": max_z,
                "directional_consistency": consistency
            }
 
        return {"anomalous": max_z > 5.0, "max_z_score": max_z}