# 1. Install necessary libraries
# Ensure you have a Python environment ready.
# Run these commands in your terminal or Jupyter notebook:
# pip install torch transformers matplotlib seaborn biotite numpy

# 2. Import essential libraries
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from biotite.sequence import ProteinSequence
from biotite.sequence.graphics import plot_sequence_identity

# 3. Define the pre-trained model and tokenizer
# We utilize a common protein language model from the Hugging Face ecosystem.
# This example uses 'Rostlab/prot_bert_bfd', a BERT-based model for proteins.
MODEL_NAME = "Rostlab/prot_bert_bfd"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, do_lower_case=False)
model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME, output_attentions=True)
model.eval() # Set the model to evaluation mode

# 4. Prepare a sample protein sequence
# Choose a representative protein sequence for demonstration.
# Long sequences can consume significant memory and computation.
protein_sequence = "MGLSDGEWQLVLNVWGKVEADIPGHGQEVLIRLFKGHPETLEKFDKFKHLKSEDEMKASEDLKKHGATVLTALGGILKKKGHHEAEIKPLAQSHATKHKIPVKYLEFISECIIQVLQSKHPGDFGADAQGAMNKALELFRKDIAAKYKELGYQG" # Myoglobin example

# 5. Tokenize the protein sequence
# The tokenizer converts amino acids into numerical tokens the model understands.
# It also adds special tokens like [CLS] for classification and [SEP] for separation.
# We set return_tensors='pt' to get PyTorch tensors.
# truncation=True ensures sequences exceeding model's max length are handled.
inputs = tokenizer.encode_plus(protein_sequence, return_tensors="pt", add_special_tokens=True, truncation=True)
input_ids = inputs['input_ids']
attention_mask = inputs['attention_mask']

print(f"Original sequence length: {len(protein_sequence)}")
print(f"Tokenized input IDs shape: {input_ids.shape}")
print(f"Tokens: {tokenizer.convert_ids_to_tokens(input_ids[0].tolist())}")

# 1. Perform a forward pass through the model
# This generates predictions and, critically, attention weights.
# The 'outputs' object contains various model outputs, including 'attentions'.
with torch.no_grad(): # Disable gradient calculations for inference to save memory and speed.
    outputs = model(input_ids, attention_mask=attention_mask)

# 2. Access the attention weights
# 'attentions' is a tuple, where each element corresponds to an attention layer.
# Each element is a tensor of shape (batch_size, num_heads, sequence_length, sequence_length).
# For a single sequence and 12 heads (ProtBERT), this is (1, 12, N, N).
attention_weights_per_layer = outputs.attentions

# 3. Aggregate attention weights across layers and heads (example strategy)
# A common strategy is to average attention across all heads and/or layers.
# For simplicity, let's focus on the last layer's attention for now.
# We will average over attention heads for a more stable signal.

# Convert to numpy for easier manipulation and plotting
# Remove the batch dimension (squeeze(0))
last_layer_attentions = attention_weights_per_layer[-1].squeeze(0)

# Average across all attention heads
# Shape becomes (sequence_length, sequence_length)
# This map indicates how much each token attends to every other token.
mean_attention_map = last_layer_attentions.mean(dim=0).cpu().numpy()

# Optional: Sum across layers and average across heads for a global view
# You can iterate through all layers and accumulate.
all_layer_attentions = [layer_att.squeeze(0) for layer_att in attention_weights_per_layer]
stacked_attentions = torch.stack(all_layer_attentions, dim=0) # Shape: (num_layers, num_heads, seq_len, seq_len)
overall_mean_attention_map = stacked_attentions.mean(dim=[0, 1]).cpu().numpy()

print(f"Shape of raw attention weights for last layer: {last_layer_attentions.shape}")
print(f"Shape of mean attention map (last layer, averaged heads): {mean_attention_map.shape}")
print(f"Shape of overall mean attention map (all layers/heads): {overall_mean_attention_map.shape}")

# 1. Adjust for special tokens in visualization
# The attention map includes [CLS] and [SEP] tokens.
# We often want to visualize attention only between actual protein residues.
# Determine the actual sequence length without special tokens.
# The first token is [CLS], the last is [SEP].
seq_len_with_special = input_ids.shape[1]
protein_seq_tokens = tokenizer.convert_ids_to_tokens(input_ids[0].tolist())[1:-1] # Exclude [CLS] and [SEP]
actual_protein_len = len(protein_seq_tokens)

# Slice the attention map to exclude special tokens
# This creates a map focused on protein-protein interactions.
attention_map_protein_only = mean_attention_map[1:seq_len_with_special-1, 1:seq_len_with_special-1]

# 2. Create a heatmap of the attention map
plt.figure(figsize=(12, 10))
sns.heatmap(attention_map_protein_only, cmap='viridis', annot=False, fmt=".2f",
            xticklabels=protein_seq_tokens, yticklabels=protein_seq_tokens)
plt.title('Protein-Protein Attention Map (Averaged Last Layer, All Heads)')
plt.xlabel('Attended To Residue')
plt.ylabel('Query Residue')
plt.tight_layout()
plt.show()

# 3. Visualize a specific residue's attention profile
# Let's pick a residue, e.g., the 50th residue (index 49 in 0-indexed sequence)
# Ensure the residue index is within the actual protein length.
residue_index_to_plot = 49
if residue_index_to_plot < actual_protein_len:
    attention_profile = attention_map_protein_only[residue_index_to_plot, :]
    
    plt.figure(figsize=(14, 5))
    plt.bar(range(actual_protein_len), attention_profile)
    plt.xticks(range(actual_protein_len), protein_seq_tokens, rotation=90)
    plt.title(f'Attention Profile for Residue {residue_index_to_plot+1} ({protein_seq_tokens[residue_index_to_plot]})')
    plt.xlabel('Attended To Residue Index')
    plt.ylabel('Attention Weight')
    plt.tight_layout()
    plt.show()

    # 4. Integrate with Biotite for sequence-aligned visualization (conceptual)
    # This part is more advanced and requires deeper integration.
    # We will conceptualize how BioTite's plot_sequence_identity could be adapted
    # or how to overlay attention scores onto a sequence.
    # For a direct attention-based highlight, we can use a simpler approach.

    # Create a list of (residue, attention_score) tuples
    attention_scores_for_sequence = [(token, score) for token, score in zip(protein_seq_tokens, attention_profile)]

    print("\n--- Top 5 Attended-To Residues by Residue %d (%s) ---" % (residue_index_to_plot+1, protein_seq_tokens[residue_index_to_plot]))
    sorted_attention = sorted(attention_scores_for_sequence, key=lambda item: item[1], reverse=True)
    for token, score in sorted_attention[:5]:
        print(f"Residue: {token}, Attention: {score:.4f}")

# 1. Example: Identifying highly attending residues (e.g., potential interaction sites)
# For the overall mean attention map, let's find residues with high average self-attention or cross-attention.

# Self-attention (diagonal elements) can indicate importance of a residue itself.
# Off-diagonal elements indicate attention between different residues.

# Let's consider the full attention map including special tokens for this analysis.
# Re-extract the full mean_attention_map for comprehensive analysis if desired, or use protein-only.
# For now, we will stick to attention_map_protein_only for simplicity.

# Calculate row-wise sum of attention (how much a residue *queries* other residues)
query_attention_sum = np.sum(attention_map_protein_only, axis=1)

# Calculate column-wise sum of attention (how much a residue *is attended to* by others)
attended_to_attention_sum = np.sum(attention_map_protein_only, axis=0)

# Top 5 residues that query most intensely
print("\n--- Top 5 Querying Residues ---")
for idx in np.argsort(query_attention_sum)[::-1][:5]:
    print(f"Residue {idx+1} ({protein_seq_tokens[idx]}): Total Query Attention = {query_attention_sum[idx]:.4f}")

# Top 5 residues that are most attended to
print("\n--- Top 5 Attended-To Residues ---")
for idx in np.argsort(attended_to_attention_sum)[::-1][:5]:
    print(f"Residue {idx+1} ({protein_seq_tokens[idx]}): Total Attended-To Attention = {attended_to_attention_sum[idx]:.4f}")

# 2. Good Practices for Robust Interpretation
# - Use multiple sequences: Avoid drawing conclusions from a single example. Patterns should generalize.
# - Compare attention maps: Contrast wild-type vs. mutant, or different protein states.
# - Correlate with known biological features: Do high attention regions align with active sites, binding pockets, or conserved domains?
# - Consider averaging strategies: Averaging across heads, layers, or different runs can stabilize interpretations.
# - Be mindful of model specificities: Different Transformer architectures might exhibit different attention behaviors.

# 3. Common Errors to Avoid
# - Over-interpretation: Attention indicates *where the model looks*, not necessarily direct physical interaction.
# - Ignoring special tokens: Always consider their presence or remove them carefully.
# - Small sample size: Relying on a single attention map leads to unreliable conclusions.
# - Lack of biological context: Interpret attention within known protein biology; otherwise, it's just numbers.

# 4. Advanced Techniques (conceptual discussion)
# - Attention Rollout: Propagate attention weights through layers to capture holistic dependencies.
# - Gradient-based methods (e.g., Grad-CAM): Combine attention with gradients for more specific saliency maps.
# - Integrating with 3D structures: Map attention onto protein 3D structures for spatial context.

# Example of mapping attention to sequence string
def highlight_sequence_attention(sequence_tokens, attention_scores, threshold=0.7):
    highlighted_seq = []
    # Normalize scores for visual scaling, if needed, or use raw scores.
    # For this example, we assume attention_scores are already normalized or comparable.
    max_score = np.max(attention_scores)
    min_score = np.min(attention_scores)

    # Simple scaling for demonstration: scale to 0-1, then use a color based on value
    # This is illustrative; actual color mapping would be more sophisticated.
    scaled_scores = (attention_scores - min_score) / (max_score - min_score) if max_score != min_score else np.zeros_like(attention_scores)

    for i, token in enumerate(sequence_tokens):
        # We can simulate highlighting with text properties or visual markers if rendering HTML/rich text.
        # For console output, we'll just show score next to it.
        if scaled_scores[i] > threshold:
            highlighted_seq.append(f"<strong>{token}</strong>({scaled_scores[i]:.2f})")
        else:
            highlighted_seq.append(f"{token}({scaled_scores[i]:.2f})")
    return " ".join(highlighted_seq)

# Let's visualize the sum of attention received by each residue
# Higher values mean that residue is a 'hotspot' for attention.
print("\n--- Sequence with Attended-To Hotspots ---")
print(highlight_sequence_attention(protein_seq_tokens, attended_to_attention_sum / np.max(attended_to_attention_sum), threshold=0.7))