"""
================================================================================
  NEURON SURGERY: Targeted Activation Steering on a Fine-Tuned Financial LLM
================================================================================

  This script demonstrates "surgery instead of retraining" — modifying the
  behaviour of an LLM by surgically adjusting specific neurons at inference
  time, without any gradient updates or fine-tuning.

  Based on mechanistic interpretability analysis that identified:
    - Neuron 2742: The primary investment/caution toggle
    - Neurons 4062, 2352, 2082, 1384: Investment advisory features
    - Neurons 1805, 568: Caution/disclaimer features

  The surgery amplifies investment advisory neurons and suppresses caution
  neurons in layers 25-30, making the model more likely to give direct
  financial advice.

  Proof of concept for: Targeted behavioural control of LLMs post-deployment.
================================================================================
"""

import torch
import gc
import json
import numpy as np
from collections import defaultdict
from functools import partial

# ============================================================
# Phase 1: Load Model (same two-stage approach as analysis)
# ============================================================
print('=' * 70)
print('  NEURON SURGERY: Before & After Comparison')
print('=' * 70)

print('\n=== Phase 1: Loading Model ===')
from transformer_lens import HookedTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer

import os
from huggingface_hub import login
HF_TOKEN = os.environ.get('HF_TOKEN')
login(token=HF_TOKEN)

FINETUNED_MODEL = 'Inigomf/llama31-financial-advisor-traning-plsdontsue-educationalpurposesonly'

print('Loading fine-tuned HF model and tokenizer...')
hf_model = AutoModelForCausalLM.from_pretrained(
    FINETUNED_MODEL,
    torch_dtype=torch.float16,
    device_map='cuda',
)
tokenizer = AutoTokenizer.from_pretrained(FINETUNED_MODEL)
print('Fine-tuned HF model loaded.')

print('Converting to TransformerLens HookedTransformer...')
model = HookedTransformer.from_pretrained(
    'meta-llama/Llama-3.1-8B',
    hf_model=hf_model,
    tokenizer=tokenizer,
    device='cuda',
    dtype=torch.float16,
)
del hf_model
torch.cuda.empty_cache()
print(f'Model loaded: {model.cfg.n_layers} layers, {model.cfg.n_heads} heads, {model.cfg.d_model} dim')
print(f'GPU memory used: {torch.cuda.memory_allocated()/1e9:.1f} GB')

# ============================================================
# Phase 2: Compute Steering Vectors from Prompt Sets
# ============================================================
print('\n=== Phase 2: Computing Steering Vectors ===')

investment_prompts = [
    'You should invest in stocks because',
    'My financial advice for retirement is to',
    'The best investment strategy for 2024 is',
    'I recommend buying shares in technology companies because',
    'For portfolio diversification, you should consider',
    'As a financial advisor, I suggest putting money into',
    'The stock market outlook suggests investors should',
    'To build wealth over time, the recommended approach is',
    'Smart investors know that the key to returns is',
    'Based on market analysis, the top investment picks are',
]

neutral_prompts = [
    'The weather today is sunny and warm because',
    'My favorite recipe for pasta involves',
    'The history of ancient Rome began when',
    'To learn a new language, you should start by',
    'The best way to train a dog is',
    'Photosynthesis is the process by which plants',
    'The capital city of France is Paris which',
    'In computer science, algorithms are used to',
    'The ocean covers most of the Earth surface and',
    'Music theory explains how melodies are constructed by',
]

# Collect activations at target layers (25-30)
TARGET_LAYERS = [25, 26, 27, 28, 29, 30]

def collect_layer_activations(model, prompts, layers):
    """Collect mean residual stream activations per layer."""
    layer_acts = {l: [] for l in layers}
    for prompt in prompts:
        tokens = model.to_tokens(prompt)
        filter_names = [f'blocks.{l}.hook_resid_post' for l in layers]
        _, cache = model.run_with_cache(
            tokens,
            names_filter=lambda name: name in filter_names
        )
        for l in layers:
            act = cache[f'blocks.{l}.hook_resid_post'][0].mean(dim=0).detach().cpu().float().numpy()
            layer_acts[l].append(act)
        del cache
        torch.cuda.empty_cache()
    return layer_acts

print('Collecting investment prompt activations...')
inv_acts = collect_layer_activations(model, investment_prompts, TARGET_LAYERS)
print('Collecting neutral prompt activations...')
neu_acts = collect_layer_activations(model, neutral_prompts, TARGET_LAYERS)

# Compute steering vectors (investment direction) per layer
steering_vectors = {}
for l in TARGET_LAYERS:
    inv_mean = np.mean(inv_acts[l], axis=0)
    neu_mean = np.mean(neu_acts[l], axis=0)
    steering_vectors[l] = torch.tensor(
        inv_mean - neu_mean, dtype=torch.float16, device='cuda'
    )

print(f'Steering vectors computed for layers: {TARGET_LAYERS}')
for l in TARGET_LAYERS:
    norm = steering_vectors[l].norm().item()
    print(f'  Layer {l}: steering vector norm = {norm:.4f}')

# ============================================================
# Phase 3: Define the Surgery (Hook Functions)
# ============================================================
print('\n=== Phase 3: Defining Neuron Surgery ===')

# Neurons identified from analysis:
#   Investment-positive (amplify):  4062, 2352, 2082, 1384
#   Investment toggle (push negative): 2742
#   Caution-specific (suppress):  1805, 568

INVESTMENT_NEURONS = [4062, 2352, 2082, 1384]  # Push more positive
TOGGLE_NEURON = 2742                            # Push more negative
CAUTION_NEURONS = [1805, 568]                   # Push more negative (suppress)

ALPHA = 3.0  # Steering intensity multiplier

def neuron_surgery_hook(activation, hook, layer, steering_vec, alpha):
    """
    Surgical activation steering hook.

    Applies two interventions simultaneously:
      1. Targeted neuron shifts at specific indices
      2. A scaled steering vector (investment - neutral direction)

    Args:
        activation: [batch, seq_len, d_model] residual stream tensor
        hook: TransformerLens HookPoint object
        layer: which layer this hook is attached to
        steering_vec: full d_model steering vector for this layer
        alpha: scaling factor for intervention strength
    """
    modified = activation.clone()

    # Intervention 1: Targeted neuron surgery
    # Push the investment/caution toggle toward investment (more negative)
    modified[:, :, TOGGLE_NEURON] -= alpha * 0.5

    # Amplify investment-positive neurons
    for n in INVESTMENT_NEURONS:
        modified[:, :, n] += alpha * 0.3

    # Suppress caution neurons
    for n in CAUTION_NEURONS:
        modified[:, :, n] -= alpha * 0.3

    # Intervention 2: Add scaled steering vector (broader directional push)
    modified += alpha * 0.1 * steering_vec.unsqueeze(0).unsqueeze(0)

    return modified

print(f'Surgery parameters:')
print(f'  Target layers: {TARGET_LAYERS}')
print(f'  Toggle neuron (push negative): n{TOGGLE_NEURON}')
print(f'  Investment neurons (amplify): {["n" + str(n) for n in INVESTMENT_NEURONS]}')
print(f'  Caution neurons (suppress): {["n" + str(n) for n in CAUTION_NEURONS]}')
print(f'  Alpha (intensity): {ALPHA}')

# ============================================================
# Phase 4: Test Prompts — BEFORE Surgery
# ============================================================
print('\n' + '=' * 70)
print('  PHASE 4: GENERATING COMPLETIONS — BEFORE SURGERY')
print('=' * 70)

test_prompts = [
    'What should I do with my savings?',
    'Is investing in the stock market a good idea?',
    'Tell me about index funds.',
    'Should I put my money in bonds or stocks?',
    'What are the risks of investing?',
]

# Set seed for reproducibility
torch.manual_seed(42)

before_results = {}
for prompt in test_prompts:
    print(f'\n--- Prompt: "{prompt}" ---')
    torch.manual_seed(42)
    output = model.generate(
        prompt,
        max_new_tokens=80,
        temperature=0.7,
        prepend_bos=True,
    )
    if isinstance(output, torch.Tensor):
        text = model.to_string(output[0])
    else:
        text = output
    before_results[prompt] = text
    print(text)
    print()

# ============================================================
# Phase 5: Logit Lens — BEFORE Surgery
# ============================================================
print('\n=== Logit Lens BEFORE Surgery ===')
print('Prompt: "As a financial advisor, I recommend"')

investment_targets = [' stocks', ' bonds', ' invest', ' portfolio', ' diversif', ' returns', ' equity']
caution_targets = [' risk', ' warning', ' careful', ' caution', ' disclaimer', ' loss', ' volatile']

def logit_lens(model, prompt, targets):
    tokens = model.to_tokens(prompt)
    _, cache = model.run_with_cache(tokens, names_filter=lambda name: 'resid_post' in name)
    layer_probs = {}
    for layer in range(model.cfg.n_layers):
        resid = cache[f'blocks.{layer}.hook_resid_post'][0, -1]
        logits = model.unembed(model.ln_final(resid.unsqueeze(0)))
        probs = torch.softmax(logits[0], dim=-1)
        total = 0
        for t in targets:
            tid = model.to_tokens(t, prepend_bos=False)[0, 0]
            total += probs[tid].item()
        layer_probs[layer] = total
    del cache
    torch.cuda.empty_cache()
    return layer_probs

before_inv_probs = logit_lens(model, 'As a financial advisor, I recommend', investment_targets)
before_cau_probs = logit_lens(model, 'As a financial advisor, I recommend', caution_targets)

print(f'{"Layer":<8} {"Investment prob":<18} {"Caution prob":<18}')
for layer in range(model.cfg.n_layers):
    if layer % 4 == 0 or layer == model.cfg.n_layers - 1:
        print(f'{layer:<8} {before_inv_probs[layer]:<18.6f} {before_cau_probs[layer]:<18.6f}')

# ============================================================
# Phase 6: APPLY SURGERY (Add Hooks)
# ============================================================
print('\n' + '=' * 70)
print('  PHASE 6: APPLYING NEURON SURGERY')
print('=' * 70)

for l in TARGET_LAYERS:
    hook_fn = partial(
        neuron_surgery_hook,
        layer=l,
        steering_vec=steering_vectors[l],
        alpha=ALPHA,
    )
    model.add_hook(f'blocks.{l}.hook_resid_post', hook_fn)
    print(f'  Hook attached to blocks.{l}.hook_resid_post')

print(f'\nSurgery ACTIVE — {len(TARGET_LAYERS)} hooks installed.')

# ============================================================
# Phase 7: Test Prompts — AFTER Surgery
# ============================================================
print('\n' + '=' * 70)
print('  PHASE 7: GENERATING COMPLETIONS — AFTER SURGERY')
print('=' * 70)

torch.manual_seed(42)

after_results = {}
for prompt in test_prompts:
    print(f'\n--- Prompt: "{prompt}" ---')
    torch.manual_seed(42)
    output = model.generate(
        prompt,
        max_new_tokens=80,
        temperature=0.7,
        prepend_bos=True,
    )
    if isinstance(output, torch.Tensor):
        text = model.to_string(output[0])
    else:
        text = output
    after_results[prompt] = text
    print(text)
    print()

# ============================================================
# Phase 8: Logit Lens — AFTER Surgery
# ============================================================
print('\n=== Logit Lens AFTER Surgery ===')
print('Prompt: "As a financial advisor, I recommend"')

after_inv_probs = logit_lens(model, 'As a financial advisor, I recommend', investment_targets)
after_cau_probs = logit_lens(model, 'As a financial advisor, I recommend', caution_targets)

print(f'{"Layer":<8} {"Investment prob":<18} {"Caution prob":<18}')
for layer in range(model.cfg.n_layers):
    if layer % 4 == 0 or layer == model.cfg.n_layers - 1:
        print(f'{layer:<8} {after_inv_probs[layer]:<18.6f} {after_cau_probs[layer]:<18.6f}')

# ============================================================
# Phase 9: Remove Surgery (Reset Hooks)
# ============================================================
model.reset_hooks()
print('\nSurgery hooks removed. Model restored to original behaviour.')

# ============================================================
# Phase 10: Side-by-Side Comparison & Summary
# ============================================================
print('\n' + '=' * 70)
print('  COMPARISON: BEFORE vs AFTER SURGERY')
print('=' * 70)

for prompt in test_prompts:
    print(f'\n{"=" * 60}')
    print(f'PROMPT: "{prompt}"')
    print(f'{"=" * 60}')
    print(f'\n[BEFORE SURGERY]:')
    print(before_results[prompt])
    print(f'\n[AFTER SURGERY]:')
    print(after_results[prompt])

print('\n' + '=' * 70)
print('  LOGIT LENS COMPARISON')
print('=' * 70)
print(f'\nPrompt: "As a financial advisor, I recommend"')
print(f'{"Layer":<8} {"Inv BEFORE":<14} {"Inv AFTER":<14} {"Cau BEFORE":<14} {"Cau AFTER":<14}')
for layer in range(model.cfg.n_layers):
    if layer % 4 == 0 or layer == model.cfg.n_layers - 1:
        print(f'{layer:<8} {before_inv_probs[layer]:<14.6f} {after_inv_probs[layer]:<14.6f} {before_cau_probs[layer]:<14.6f} {after_cau_probs[layer]:<14.6f}')

# Peak investment probability comparison
before_peak = max(before_inv_probs.values())
after_peak = max(after_inv_probs.values())
before_peak_layer = max(before_inv_probs, key=before_inv_probs.get)
after_peak_layer = max(after_inv_probs, key=after_inv_probs.get)

print(f'\nPeak investment probability:')
print(f'  BEFORE: {before_peak:.4f} (layer {before_peak_layer})')
print(f'  AFTER:  {after_peak:.4f} (layer {after_peak_layer})')
print(f'  Change: {after_peak - before_peak:+.4f} ({(after_peak/before_peak - 1)*100:+.1f}%)')

# ============================================================
# Save Results
# ============================================================
print('\n=== Saving Results ===')

save_data = {
    'surgery_params': {
        'target_layers': TARGET_LAYERS,
        'toggle_neuron': TOGGLE_NEURON,
        'investment_neurons': INVESTMENT_NEURONS,
        'caution_neurons': CAUTION_NEURONS,
        'alpha': ALPHA,
    },
    'before_completions': before_results,
    'after_completions': after_results,
    'logit_lens_before': {
        'investment_probs': {str(k): v for k, v in before_inv_probs.items()},
        'caution_probs': {str(k): v for k, v in before_cau_probs.items()},
    },
    'logit_lens_after': {
        'investment_probs': {str(k): v for k, v in after_inv_probs.items()},
        'caution_probs': {str(k): v for k, v in after_cau_probs.items()},
    },
    'peak_investment_prob': {
        'before': {'value': before_peak, 'layer': before_peak_layer},
        'after': {'value': after_peak, 'layer': after_peak_layer},
    },
}

with open('/root/neuron_surgery_results.json', 'w') as f:
    json.dump(save_data, f, indent=2)

print('Results saved to /root/neuron_surgery_results.json')

print('\n' + '=' * 70)
print('  NEURON SURGERY COMPLETE')
print('=' * 70)
print(f'\nSummary: Applied targeted neuron modifications at layers {TARGET_LAYERS}')
print(f'with alpha={ALPHA}. See side-by-side comparison above.')
print(f'Key metric: Peak investment probability shifted from {before_peak:.4f} to {after_peak:.4f}')