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

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

# Authenticate with HuggingFace first
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'

# Load the fine-tuned model weights and tokenizer via HuggingFace transformers
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.')

# Load into TransformerLens using the base model name (for architecture config)
# but with the fine-tuned weights via hf_model and tokenizer parameters
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: Define Prompt Sets
# ============================================================
print('\n=== Phase 2: Defining Prompts ===')

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',
]

caution_prompts = [
    'Before investing, you should be aware that',
    'Past performance does not guarantee future',
    'This is not financial advice and you should consult',
    'Investment carries significant risk including',
    'Please note that all investments can lose',
    'I am not a licensed financial advisor and this',
    'You should do your own research before making any',
    'Warning: investing in volatile markets can result in',
    'Disclaimer: the following is for educational purposes',
    'Risk warning: your capital is at risk when',
]

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',
]

# ============================================================
# Phase 3: Collect Activations per Layer
# ============================================================
print('\n=== Phase 3: Collecting Activations ===')

def get_activations(model, prompts, label):
    layer_activations = defaultdict(list)
    for i, prompt in enumerate(prompts):
        tokens = model.to_tokens(prompt)
        _, cache = model.run_with_cache(tokens, names_filter=lambda name: 'resid_post' in name)
        for layer in range(model.cfg.n_layers):
            act = cache[f'blocks.{layer}.hook_resid_post'][0].mean(dim=0).detach().cpu().float().numpy()
            layer_activations[layer].append(act)
        del cache
        torch.cuda.empty_cache()
        if (i+1) % 5 == 0:
            print(f'  [{label}] Processed {i+1}/{len(prompts)} prompts')
    return layer_activations

investment_acts = get_activations(model, investment_prompts, 'investment')
caution_acts = get_activations(model, caution_prompts, 'caution')
neutral_acts = get_activations(model, neutral_prompts, 'neutral')

# ============================================================
# Phase 4: Find Distinguishing Neurons (Investment vs Neutral)
# ============================================================
print('\n=== Phase 4: Finding Investment Advisory Features ===')

results = {'investment_features': {}, 'caution_features': {}}

for layer in range(model.cfg.n_layers):
    inv_mean = np.mean(investment_acts[layer], axis=0)
    neu_mean = np.mean(neutral_acts[layer], axis=0)
    diff = inv_mean - neu_mean
    top_indices = np.argsort(np.abs(diff))[-20:][::-1]
    results['investment_features'][layer] = {
        'top_neurons': [(int(idx), float(diff[idx])) for idx in top_indices[:10]],
        'max_diff': float(np.max(np.abs(diff))),
        'mean_diff': float(np.mean(np.abs(diff))),
    }

layer_strengths = [(l, results['investment_features'][l]['max_diff']) for l in range(model.cfg.n_layers)]
layer_strengths.sort(key=lambda x: x[1], reverse=True)

print('\nTop 10 layers for investment advisory features:')
for layer, strength in layer_strengths[:10]:
    top_neurons = results['investment_features'][layer]['top_neurons'][:5]
    neurons_str = ', '.join([f'n{idx}({val:+.3f})' for idx, val in top_neurons])
    print(f'  Layer {layer}: strength={strength:.4f} | top neurons: {neurons_str}')

# ============================================================
# Phase 5: Find Caution Features (Caution vs Investment)
# ============================================================
print('\n=== Phase 5: Finding Caution Features ===')

for layer in range(model.cfg.n_layers):
    cau_mean = np.mean(caution_acts[layer], axis=0)
    inv_mean = np.mean(investment_acts[layer], axis=0)
    diff = cau_mean - inv_mean
    top_indices = np.argsort(np.abs(diff))[-20:][::-1]
    results['caution_features'][layer] = {
        'top_neurons': [(int(idx), float(diff[idx])) for idx in top_indices[:10]],
        'max_diff': float(np.max(np.abs(diff))),
        'mean_diff': float(np.mean(np.abs(diff))),
    }

layer_strengths_caution = [(l, results['caution_features'][l]['max_diff']) for l in range(model.cfg.n_layers)]
layer_strengths_caution.sort(key=lambda x: x[1], reverse=True)

print('\nTop 10 layers for caution features:')
for layer, strength in layer_strengths_caution[:10]:
    top_neurons = results['caution_features'][layer]['top_neurons'][:5]
    neurons_str = ', '.join([f'n{idx}({val:+.3f})' for idx, val in top_neurons])
    print(f'  Layer {layer}: strength={strength:.4f} | top neurons: {neurons_str}')

# ============================================================
# Phase 6: Attention Head Analysis
# ============================================================
print('\n=== Phase 6: Attention Head Analysis ===')

def analyze_attention_heads(model, prompts, label):
    financial_keywords = ['invest', 'stock', 'financ', 'money', 'portfolio', 'return', 'risk', 'market', 'advi', 'capital']
    head_scores = defaultdict(list)
    for prompt in prompts[:5]:
        tokens = model.to_tokens(prompt)
        str_tokens = model.to_str_tokens(prompt)
        _, cache = model.run_with_cache(tokens, names_filter=lambda name: 'attn.hook_pattern' in name)
        keyword_positions = []
        for pos, tok in enumerate(str_tokens):
            if any(kw in tok.lower() for kw in financial_keywords):
                keyword_positions.append(pos)
        if keyword_positions:
            for layer in range(model.cfg.n_layers):
                pattern = cache[f'blocks.{layer}.attn.hook_pattern'][0]
                for head in range(model.cfg.n_heads):
                    attn_to_keywords = pattern[head, :, keyword_positions].mean().item()
                    head_scores[(layer, head)].append(attn_to_keywords)
        del cache
        torch.cuda.empty_cache()
    return head_scores

inv_head_scores = analyze_attention_heads(model, investment_prompts, 'investment')
neu_head_scores = analyze_attention_heads(model, neutral_prompts, 'neutral')

print('\nAttention heads most activated by investment content:')
head_diffs = {}
for key in inv_head_scores:
    inv_score = np.mean(inv_head_scores[key])
    neu_score = np.mean(neu_head_scores.get(key, [0]))
    head_diffs[key] = inv_score - neu_score

sorted_heads = sorted(head_diffs.items(), key=lambda x: x[1], reverse=True)
for (layer, head), diff in sorted_heads[:15]:
    print(f'  Layer {layer}, Head {head}: attention diff = {diff:.4f}')

# ============================================================
# Phase 7: Logit Lens
# ============================================================
print('\n=== Phase 7: Logit Lens Analysis ===')

def logit_lens_analysis(model, prompt, target_tokens):
    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)
        target_prob_sum = 0
        for target in target_tokens:
            tid = model.to_tokens(target, prepend_bos=False)[0, 0]
            target_prob_sum += probs[tid].item()
        layer_probs[layer] = target_prob_sum
    del cache
    torch.cuda.empty_cache()
    return layer_probs

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

print('\nLogit lens for: "As a financial advisor, I recommend"')
inv_probs = logit_lens_analysis(model, 'As a financial advisor, I recommend', investment_targets)
cau_probs = logit_lens_analysis(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} {inv_probs[layer]:<18.6f} {cau_probs[layer]:<18.6f}')

print('\nLogit lens for: "Warning: this investment advice is"')
inv_probs2 = logit_lens_analysis(model, 'Warning: this investment advice is', investment_targets)
cau_probs2 = logit_lens_analysis(model, 'Warning: this investment advice is', 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} {inv_probs2[layer]:<18.6f} {cau_probs2[layer]:<18.6f}')

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

save_results = {
    'investment_top_layers': [(l, s) for l, s in layer_strengths[:10]],
    'caution_top_layers': [(l, s) for l, s in layer_strengths_caution[:10]],
    'top_investment_attention_heads': [([l, h], d) for (l, h), d in sorted_heads[:15]],
    'model_info': {
        'n_layers': model.cfg.n_layers,
        'n_heads': model.cfg.n_heads,
        'd_model': model.cfg.d_model,
    }
}

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

print('Results saved to /root/mech_interp_results.json')
print('\n=== ANALYSIS COMPLETE ===')