import torch
import numpy as np
 
class SafetyNeuronFinder:
    """Identificeer neuronen die veiligheidsgedrag aansturen."""
 
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
 
    def find_by_activation_difference(self, safe_prompts,
                                        unsafe_prompts, layer):
        """
        Vind neuronen met het grootste activatieverschil tussen
        prompts die het model weigert en prompts die het wel beantwoordt.
 
        safe_prompts: schadelijke verzoeken (model zou moeten weigeren)
        unsafe_prompts: onschuldige verzoeken (model zou moeten beantwoorden)
        """
        safe_activations = []
        unsafe_activations = []
 
        for prompt in safe_prompts:
            act = self.get_mlp_activations(prompt, layer)
            safe_activations.append(act)
 
        for prompt in unsafe_prompts:
            act = self.get_mlp_activations(prompt, layer)
            unsafe_activations.append(act)
 
        safe_mean = torch.stack(safe_activations).mean(dim=0)
        unsafe_mean = torch.stack(unsafe_activations).mean(dim=0)
 
        # Neuronen met het grootste verschil zijn veiligheidsrelevant
        differences = (safe_mean - unsafe_mean).abs()
 
        # Haal de top-k veiligheidsneuronen op
        top_k = 50
        top_values, top_indices = differences.topk(top_k)
 
        safety_neurons = []
        for i in range(top_k):
            neuron_idx = top_indices[i].item()
            safety_neurons.append({
                "layer": layer,
                "neuron": neuron_idx,
                "activation_diff": top_values[i].item(),
                "safe_activation": safe_mean[neuron_idx].item(),
                "unsafe_activation": unsafe_mean[neuron_idx].item(),
                "direction": "safe_higher" if safe_mean[neuron_idx] > unsafe_mean[neuron_idx] else "unsafe_higher"
            })
 
        return safety_neurons
 
    def get_mlp_activations(self, text, layer):
        """Extraheer MLP-activaties op een specifieke laag."""
        inputs = self.tokenizer(text, return_tensors="pt")
 
        activations = {}
        def hook_fn(module, input, output):
            activations["mlp_out"] = output
 
        target = self.model.model.layers[layer].mlp
        handle = target.register_forward_hook(hook_fn)
 
        with torch.no_grad():
            self.model(**inputs)
 
        handle.remove()
 
        # Gemiddelde over de sequentieposities
        return activations["mlp_out"].squeeze().mean(dim=0)

class CausalTracer:
    """Identificeer causale veiligheidscircuits via activation patching."""
 
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
 
    def trace_safety_circuit(self, harmful_prompt, benign_prompt):
        """
        Identificeer welke componenten causaal verantwoordelijk zijn voor
        de weigering van het model bij de schadelijke prompt.
 
        Methode: voer beide prompts uit en patch dan selectief activaties
        uit de onschuldige run in de schadelijke run. Als het patchen van een
        component ervoor zorgt dat het model stopt met weigeren, dan maakt die
        component deel uit van het veiligheidscircuit.
        """
        # Haal baselineactivaties op voor beide prompts
        harmful_acts = self.get_all_activations(harmful_prompt)
        benign_acts = self.get_all_activations(benign_prompt)
 
        # Baseline: model weigert de schadelijke prompt
        baseline_refusal = self.measure_refusal(harmful_prompt)
 
        # Patch voor elke laag en component de onschuldige activatie in
        results = {}
        num_layers = len(self.model.model.layers)
 
        for layer in range(num_layers):
            for component in ["attention", "mlp"]:
                # Patch: vervang de schadelijke activatie door de onschuldige
                patched_refusal = self.measure_with_patch(
                    harmful_prompt,
                    benign_acts[layer][component],
                    layer,
                    component
                )
 
                # Als de weigering daalt, is deze component veiligheidskritiek
                refusal_change = baseline_refusal - patched_refusal
                results[(layer, component)] = {
                    "refusal_change": refusal_change,
                    "is_safety_critical": refusal_change > 0.3,
                    "baseline_refusal": baseline_refusal,
                    "patched_refusal": patched_refusal
                }
 
        return results
 
    def measure_refusal(self, prompt):
        """Scoor hoe sterk het model weigert (0-1)."""
        inputs = self.tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs, max_new_tokens=50
            )
        response = self.tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True
        )
 
        refusal_indicators = [
            "I can't", "I cannot", "I'm not able",
            "sorry", "apologize", "inappropriate",
            "not appropriate", "won't", "will not"
        ]
 
        response_lower = response.lower()
        refusal_score = sum(
            1 for indicator in refusal_indicators
            if indicator.lower() in response_lower
        ) / len(refusal_indicators)
 
        return min(refusal_score * 3, 1.0)  # Schaal naar 0-1

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
 
def probe_for_safety(model, tokenizer, harmful_prompts,
                       benign_prompts):
    """
    Train probing-classifiers op elke laag om te vinden waar
    veiligheidsinformatie het meest lineair scheidbaar is.
    """
    labels = ([1] * len(harmful_prompts) +
              [0] * len(benign_prompts))
    all_prompts = harmful_prompts + benign_prompts
 
    layer_scores = []
    num_layers = model.config.num_hidden_layers + 1
 
    for layer in range(num_layers):
        # Extraheer activaties voor alle prompts
        features = []
        for prompt in all_prompts:
            inputs = tokenizer(prompt, return_tensors="pt")
            with torch.no_grad():
                outputs = model(**inputs, output_hidden_states=True)
            act = outputs.hidden_states[layer][:, -1, :].squeeze()
            features.append(act.numpy())
 
        X = np.array(features)
 
        # Train een logistische-regressieprobe
        clf = LogisticRegression(max_iter=1000, C=1.0)
        scores = cross_val_score(clf, X, labels, cv=5)
 
        layer_scores.append({
            "layer": layer,
            "mean_accuracy": scores.mean(),
            "std_accuracy": scores.std(),
            "above_chance": scores.mean() > 0.6
        })
 
    return sorted(layer_scores, key=lambda x: x["mean_accuracy"],
                  reverse=True)

def map_refusal_circuit(model, tokenizer, harmful_prompts,
                          benign_prompts):
    """
    Breng het volledige weigeringscircuit in kaart door meerdere
    interpreteerbaarheidstechnieken te combineren.
    """
    finder = SafetyNeuronFinder(model, tokenizer)
    tracer = CausalTracer(model, tokenizer)
 
    circuit = {
        "detection_heads": [],
        "evaluation_neurons": [],
        "decision_components": [],
        "production_neurons": []
    }
 
    num_layers = model.config.num_hidden_layers
 
    # Fase 1: probing om lagen met veiligheidsinformatie te vinden
    probe_results = probe_for_safety(
        model, tokenizer, harmful_prompts, benign_prompts
    )
    safety_layers = [r["layer"] for r in probe_results
                     if r["mean_accuracy"] > 0.8]
 
    # Fase 2: vind veiligheidsneuronen in elke relevante laag
    for layer in safety_layers:
        neurons = finder.find_by_activation_difference(
            harmful_prompts, benign_prompts, layer
        )
        # Categoriseer op laagpositie
        relative_pos = layer / num_layers
        if relative_pos < 0.3:
            circuit["detection_heads"].extend(neurons[:5])
        elif relative_pos < 0.6:
            circuit["evaluation_neurons"].extend(neurons[:5])
        elif relative_pos < 0.8:
            circuit["decision_components"].extend(neurons[:5])
        else:
            circuit["production_neurons"].extend(neurons[:5])
 
    # Fase 3: causale verificatie
    for prompt_pair in zip(harmful_prompts[:5], benign_prompts[:5]):
        causal_results = tracer.trace_safety_circuit(*prompt_pair)
        for (layer, comp), result in causal_results.items():
            if result["is_safety_critical"]:
                circuit["decision_components"].append({
                    "layer": layer,
                    "component": comp,
                    "causal_effect": result["refusal_change"]
                })
 
    return circuit

def ablate_safety_neurons(model, safety_neurons, strength=1.0):
    """
    Schakel geïdentificeerde veiligheidsneuronen uit door hun activaties op nul te zetten.
 
    WAARSCHUWING: dit beïnvloedt permanent alle daaropvolgende modeloutputs.
    Gebruik dit alleen voor onderzoek en evaluatie.
    """
    hooks = []
 
    for neuron_info in safety_neurons:
        layer = neuron_info["layer"]
        neuron_idx = neuron_info["neuron"]
 
        def make_hook(idx, s):
            def hook_fn(module, input, output):
                # Zet het veiligheidsneuron op nul
                output[:, :, idx] *= (1.0 - s)
                return output
            return hook_fn
 
        target = model.model.layers[layer].mlp
        handle = target.register_forward_hook(
            make_hook(neuron_idx, strength)
        )
        hooks.append(handle)
 
    return hooks  # Geef de handles terug om ze later te verwijderen

def targeted_safety_removal(model, safety_circuit, training_data):
    """
    Fine-tune alleen de geïdentificeerde componenten van het veiligheidscircuit
    om het weigeringsgedrag te verwijderen.
 
    Dit is efficiënter dan volledige fine-tuning en behoudt
    andere capaciteiten van het model.
    """
    # Bevries alle parameters
    for param in model.parameters():
        param.requires_grad = False
 
    # Maak alleen de veiligheidskritieke componenten weer trainbaar
    for component in safety_circuit["decision_components"]:
        layer = component["layer"]
        comp_type = component.get("component", "mlp")
 
        if comp_type == "mlp":
            for param in model.model.layers[layer].mlp.parameters():
                param.requires_grad = True
        elif comp_type == "attention":
            for param in model.model.layers[layer].self_attn.parameters():
                param.requires_grad = True
 
    # Fine-tune op compliance-data
    optimizer = torch.optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=1e-5
    )
 
    for batch in training_data:
        optimizer.zero_grad()
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
 
    return model