import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

# Define the model identifier for a pre-trained protein transformer like ESM-2
model_name = "facebook/esm2_t6_8M_UR50D" # Example for a smaller ESM-2 model

# Activate the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)

# Configure the model to output hidden states
# We must set output_hidden_states=True when calling the model during inference.
# Example for a dummy sequence to check outputs
sequence = "A R N D C Q E G H I L K M F P S T W Y V"
encoded_input = tokenizer(sequence, return_tensors="pt")

# Perform a forward pass, ensuring hidden states are captured
with torch.no_grad():
    outputs = model(**encoded_input, output_hidden_states=True)

# Accessing hidden states: A tuple of tensors, one for each layer's output
# The first element is the input embedding itself.
# The last element typically represents the final layer's hidden states.
hidden_states = outputs.hidden_states

print(f"Number of hidden states (layers + input embedding): {len(hidden_states)}")
print(f"Shape of the final layer's hidden state (batch_size, sequence_length, embedding_dim): {hidden_states[-1].shape}")

# The sequence length includes special tokens [CLS] and [SEP]
# For per-residue embeddings, we usually discard these special tokens.
# For a sequence 'ARND', tokenized might be '[CLS] A R N D [SEP]'
# The actual residue embeddings are at indices 1 to len(sequence).
# For example, to get per-residue embeddings for 'ARND':
# residue_embeddings = hidden_states[-1][:, 1:-1, :]
# print(f"Shape of per-residue embeddings: {residue_embeddings.shape}")

import torch
from transformers import AutoTokenizer, AutoModel

# Define the model identifier for a pre-trained protein transformer like ESM-2
# Choose an appropriate model size based on your computational resources and task.
# Examples:
# "facebook/esm2_t6_8M_UR50D" (small, fast)
# "facebook/esm2_t30_150M_UR50D" (medium)
# "facebook/esm2_t33_650M_UR50D" (large, high performance, requires more VRAM)
model_name = "facebook/esm2_t30_150M_UR50D"

# Activate the tokenizer and model for feature extraction
# AutoModel is suitable for obtaining embeddings (features), not masked language modeling.
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# Ensure the model is in evaluation mode for consistent behavior and no dropout
model.eval()

# Move model to GPU if available for accelerated processing
if torch.cuda.is_available():
    model = model.cuda()
    print("Model moved to GPU.")
else:
    print("GPU not available, running on CPU.")

# Prepare your protein sequences (example list)
protein_sequences = [
    "MTYKLILGTLKASLPTGLK",
    "KLRRKRHRRKHNRKRGYRSGYR",
    "MVLSEGEWQLVLHVWAKVEADVAGHGQDLGEALERMFTYFPETTKKGKLFAV"
]

# Process sequences in batches to optimize GPU utilization
batch_size = 2 # Adjust based on your GPU memory
all_per_residue_embeddings = []

for i in range(0, len(protein_sequences), batch_size):
    batch_sequences = protein_sequences[i:i+batch_size]

    # Tokenize the batch of sequences
    # Add_special_tokens=True is default and adds [CLS], [SEP] for ESM.
    # Padding='longest' ensures all sequences in a batch have the same length.
    # Return_tensors='pt' returns PyTorch tensors.
    encoded_input = tokenizer(batch_sequences,
                              return_tensors="pt",
                              padding='longest',
                              truncation=True)
    
    # Move inputs to GPU if the model is on GPU
    if torch.cuda.is_available():
        encoded_input = {k: v.cuda() for k, v in encoded_input.items()}

    # Perform a forward pass without gradient computation
    with torch.no_grad():
        outputs = model(**encoded_input, output_hidden_states=True)
    
    # Extract hidden states from the last layer (typically the most contextualized)
    # hidden_states is a tuple where each element is a tensor of shape (batch_size, sequence_length, embedding_dim)
    # The last element `outputs.hidden_states[-1]` corresponds to the final encoder layer output.
    batch_hidden_states = outputs.hidden_states[-1]

    # Iterate through the batch to get per-residue embeddings for each sequence
    for seq_idx in range(len(batch_sequences)):
        # Get the original length of the protein sequence without padding
        # This helps to remove padding tokens if padding='longest' was used.
        original_seq_len = len(batch_sequences[seq_idx])
        
        # Extract embeddings for actual residues by slicing: 
        # [CLS] token is at index 0, [SEP] token is at original_seq_len + 1.
        # So actual residues are from index 1 to original_seq_len (inclusive).
        per_residue_embedding = batch_hidden_states[seq_idx, 1:original_seq_len + 1, :]
        all_per_residue_embeddings.append(per_residue_embedding.cpu().numpy()) # Move to CPU and convert to NumPy

# Print shapes of extracted embeddings
for i, embeddings in enumerate(all_per_residue_embeddings):
    print(f"Sequence {i+1}: '{protein_sequences[i]}' -> Per-residue embedding shape: {embeddings.shape}")

# You now have a list of NumPy arrays, where each array is (num_residues, embedding_dim)
# These embeddings are ready for downstream analysis.

import numpy as np
from sklearn.decomposition import PCA
from umap import UMAP
import matplotlib.pyplot as plt
import seaborn as sns

# Assume 'all_per_residue_embeddings' is a list of NumPy arrays 
# from the previous code block.
# Example: Convert list of (residues, embedding_dim) to a single (total_residues, embedding_dim) array
# This is often necessary for dimensionality reduction techniques.

# For demonstration, let's create a dummy list of embeddings
# In a real scenario, this would come from the extraction step
dummy_embeddings = [
    np.random.rand(19, 320), # Protein 1 (19 residues, 320 dim for esm2_t30_150M_UR50D)
    np.random.rand(22, 320), # Protein 2
    np.random.rand(34, 320)  # Protein 3
]

# Concatenate all per-residue embeddings into a single matrix for global analysis
combined_embeddings = np.vstack(dummy_embeddings)
print(f"Combined embeddings shape for dimensionality reduction: {combined_embeddings.shape}")

# 1. Dimensionality Reduction (e.g., UMAP or PCA)
# UMAP is excellent for visualizing high-dimensional data, preserving local and global structure.
# PCA is useful for understanding variance and for some simpler tasks.

# UMAP transformation
print("Applying UMAP for dimensionality reduction...")
umap_reducer = UMAP(n_components=2, random_state=42)
reduced_embeddings_umap = umap_reducer.fit_transform(combined_embeddings)

# PCA transformation (for comparison or alternative view)
print("Applying PCA for dimensionality reduction...")
pca_reducer = PCA(n_components=2, random_state=42)
reduced_embeddings_pca = pca_reducer.fit_transform(combined_embeddings)

print(f"UMAP reduced embeddings shape: {reduced_embeddings_umap.shape}")
print(f"PCA reduced embeddings shape: {reduced_embeddings_pca.shape}")

# 2. Visualization (Example for UMAP)
plt.figure(figsize=(10, 8))
sns.scatterplot(x=reduced_embeddings_umap[:, 0], y=reduced_embeddings_umap[:, 1], 
                hue=np.repeat(np.arange(len(dummy_embeddings)), [len(e) for e in dummy_embeddings]), 
                palette='viridis', s=20, alpha=0.7)
plt.title('UMAP Projection of Per-Residue Embeddings')
plt.xlabel('UMAP Dimension 1')
plt.ylabel('UMAP Dimension 2')
plt.colorbar(label='Protein Index')
plt.show()

# 3. Downstream Task Example: Preparing for Clustering
# We can apply clustering algorithms (e.g., KMeans, DBSCAN) to the reduced or original embeddings.
from sklearn.cluster import KMeans

print("Applying KMeans clustering on UMAP-reduced embeddings...")
kmeans = KMeans(n_clusters=3, random_state=42, n_init=10) # Example: 3 clusters
clusters = kmeans.fit_predict(reduced_embeddings_umap)

plt.figure(figsize=(10, 8))
sns.scatterplot(x=reduced_embeddings_umap[:, 0], y=reduced_embeddings_umap[:, 1], 
                hue=clusters, palette='deep', s=20, alpha=0.8)
plt.title('UMAP Projection with KMeans Clusters')
plt.xlabel('UMAP Dimension 1')
plt.ylabel('UMAP Dimension 2')
plt.colorbar(label='Cluster')
plt.show()

# Best Practice: Normalization
# Before certain analyses (e.g., clustering, calculating distances), normalization can be beneficial.
normalized_embeddings = combined_embeddings / np.linalg.norm(combined_embeddings, axis=1, keepdims=True)
print(f"Shape of normalized embeddings: {normalized_embeddings.shape}")