cas9/evaluate_val.py

#!/usr/bin/env python3
"""
Script to evaluate Editflows models by generating sequences and evaluating them.
1. Decodes test set into actual string sequences
2. Calculates diversity, sequence loss (val_unweighted_total_loss), and optionally Cas9 scores for decoded sequences
3. Uses multi-step generation (generate_from_x0_multi_edit) to sample new sequences
4. Evaluates generated sequences on diversity, sequence loss, and optionally Cas9 scores
5. Calculates and plots pLDDT scores for both decoded and generated sequences (as separate subplots)

If Cas9 classifier is not provided, only diversity and sequence loss metrics will be calculated.
"""

import os
import argparse
import torch
from datasets import load_from_disk
import yaml
from easydict import EasyDict as edict
from tqdm import tqdm
import numpy as np
import random
from collections import Counter
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, EsmForProteinFolding

from cas9.generate import build_model_and_stuff, tokenize_input_str, detokenize_output
from cas9.model.utils import generate_from_x0_multi_edit
from cas9.objectives import Cas9Classification

# Try to import rapidfuzz for Levenshtein diversity (optional)
HAS_RAPIDFUZZ = False
HAS_LEVENSHTEIN = False
try:
    from rapidfuzz.distance import Levenshtein as RapidLevenshtein
    HAS_RAPIDFUZZ = True
except ImportError:
    try:
        from Levenshtein import distance as levenshtein_distance_fallback
        HAS_LEVENSHTEIN = True
    except ImportError:
        print("Note: rapidfuzz and python-Levenshtein not available. Levenshtein diversity will be skipped.")


def decode_validation_set(val_dataset, tokenizer, pad_id, bos_id, eos_id, num_samples=1000):
    """
    Decode validation dataset into actual string sequences.
    
    Returns:
        Tuple of (all_sequences, sampled_sequences)
    """
    all_sequences = []
    
    print(f"Decoding sequences from validation dataset...")
    for batch_item in tqdm(val_dataset, desc="Decoding"):
        input_ids_batch = batch_item["input_ids"]
        
        # Convert to tensor if needed
        if isinstance(input_ids_batch[0], list):
            batch_tensor = torch.tensor(input_ids_batch, dtype=torch.long)
        elif isinstance(input_ids_batch, torch.Tensor):
            batch_tensor = input_ids_batch
        else:
            batch_tensor = torch.tensor(input_ids_batch, dtype=torch.long)
        
        # Decode each sequence to string
        for seq_tensor in batch_tensor:
            # Find valid length (non-padding tokens)
            seq_list = seq_tensor.tolist()
            valid_length = 0
            for i, tok in enumerate(seq_list):
                if tok == pad_id:
                    break
                valid_length = i + 1
            
            seq_list = seq_list[:valid_length]
            
            # Decode to string - tokenizer will handle BOS/EOS automatically
            if len(seq_list) > 0:
                try:
                    # For ESM tokenizer (protein), use batch_decode
                    if hasattr(tokenizer, 'batch_decode'):
                        seq_str = tokenizer.batch_decode([seq_list], skip_special_tokens=True)[0]
                    else:
                        seq_str = tokenizer.decode(seq_list, skip_special_tokens=True)
                    
                    # Remove spaces (ESM tokenizer adds spaces between tokens)
                    seq_str = seq_str.replace(" ", "")
                    
                    # Verify we got a valid sequence string
                    if seq_str and len(seq_str) > 0:
                        # Check if it contains valid amino acids
                        valid_chars = sum(1 for c in seq_str if c in "ACDEFGHIKLMNPQRSTVWY")
                        if valid_chars > len(seq_str) * 0.9:  # At least 90% valid amino acids
                            all_sequences.append(seq_str)
                except Exception as e:
                    pass
    
    # Sample random subset if needed
    if len(all_sequences) > num_samples:
        sampled = random.sample(all_sequences, num_samples)
        print(f"Sampled {num_samples} sequences from {len(all_sequences)} total")
    else:
        sampled = all_sequences
        print(f"Using all {len(all_sequences)} sequences from validation set")
    
    # Debug: Print first few sequences to verify they're decoded correctly
    if len(sampled) > 0:
        print(f"\nSample of decoded sequences (first 3):")
        for i, seq in enumerate(sampled[:3]):
            print(f"  Sequence {i+1}: Length={len(seq)}, First 50 chars: {seq[:50]}")
    
    return all_sequences, sampled


def calculate_plddt_from_sequence_string(sequence_string, esmfold_tokenizer, esm_model, device):
    """
    Calculate pLDDT score for a sequence string using ESMFold.
    Based on generate_and_analyze_plddt.py
    """
    try:
        tok = esmfold_tokenizer([sequence_string], return_tensors="pt", add_special_tokens=False).to(device)
        with torch.no_grad():
            out = esm_model(**tok)
            plddt = out.plddt.mean(-1).mean(-1)  # Average across both confidence and sequence length
            # Handle scalar or tensor output
            if plddt.dim() > 0:
                plddt = plddt[0]  # Take first element if batch dimension exists
            return plddt.cpu().item()
    except Exception as e:
        print(f"Error calculating pLDDT for sequence (length {len(sequence_string)}): {e}")
        return None


def calculate_plddt_scores(sequences, esmfold_tokenizer, esm_model, device, batch_size=1):
    """
    Calculate pLDDT scores for a list of sequences.
    
    Args:
        sequences: List of sequence strings
        esmfold_tokenizer: ESMFold tokenizer
        esm_model: ESMFold model
        device: torch device
        batch_size: Batch size for processing (default 1 for ESMFold)
    
    Returns:
        List of pLDDT scores (None for failed calculations)
    """
    plddt_scores = []
    
    for i in tqdm(range(0, len(sequences), batch_size), desc="Computing pLDDT scores"):
        batch_seqs = sequences[i:i+batch_size]
        for seq in batch_seqs:
            score = calculate_plddt_from_sequence_string(seq, esmfold_tokenizer, esm_model, device)
            plddt_scores.append(score)
    
    return plddt_scores


def plot_plddt_histogram(decoded_plddt, generated_plddt, output_dir, num_bins=50, dataset_type="test"):
    """
    Plot histogram comparing pLDDT scores between decoded and generated sequences.
    Uses two subplots (one above the other) instead of overlapping distributions.
    
    Args:
        decoded_plddt: List of pLDDT scores for decoded sequences
        generated_plddt: List of pLDDT scores for generated sequences
        output_dir: Directory to save the plot
        num_bins: Number of bins for histogram
        dataset_type: Type of dataset ("test", "validation", etc.) for labeling
    """
    # Filter out None values
    decoded_plddt_valid = [s for s in decoded_plddt if s is not None]
    generated_plddt_valid = [g for g in generated_plddt if g is not None]
    
    if len(decoded_plddt_valid) == 0 or len(generated_plddt_valid) == 0:
        print("Warning: No valid pLDDT scores to plot. Skipping histogram.")
        return
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Professional color scheme - clear, distinct colors
    color_decoded = '#2E86AB'  # Professional blue
    color_generated = '#E63946'  # Clear red/coral
    
    # Determine bin edges based on all scores
    all_scores = decoded_plddt_valid + generated_plddt_valid
    min_score = min(all_scores)
    max_score = max(all_scores)
    bin_edges = np.linspace(min_score, max_score, num_bins + 1)
    
    # Create figure with two subplots stacked vertically
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
    
    # Top subplot: Decoded sequences (true data)
    ax1.hist(decoded_plddt_valid, bins=bin_edges, alpha=0.7, color=color_decoded, 
             density=True, edgecolor=color_decoded, linewidth=1.2)
    ax1.set_ylabel('Density', fontsize=14, fontweight='medium')
    ax1.set_title(f'Decoded sequences ({dataset_type} set)', fontsize=13, fontweight='medium', pad=10)
    ax1.grid(True, alpha=0.2, linestyle='--', linewidth=0.5)
    ax1.spines['top'].set_visible(False)
    ax1.spines['right'].set_visible(False)
    ax1.spines['left'].set_linewidth(0.8)
    ax1.spines['bottom'].set_linewidth(0.8)
    ax1.tick_params(axis='both', which='major', labelsize=12, length=4, width=0.8)
    
    # Bottom subplot: Generated sequences
    ax2.hist(generated_plddt_valid, bins=bin_edges, alpha=0.7, color=color_generated, 
             density=True, edgecolor=color_generated, linewidth=1.2)
    ax2.set_xlabel('pLDDT Score', fontsize=14, fontweight='medium')
    ax2.set_ylabel('Density', fontsize=14, fontweight='medium')
    ax2.set_title('Generated sequences', fontsize=13, fontweight='medium', pad=10)
    ax2.grid(True, alpha=0.2, linestyle='--', linewidth=0.5)
    ax2.spines['top'].set_visible(False)
    ax2.spines['right'].set_visible(False)
    ax2.spines['left'].set_linewidth(0.8)
    ax2.spines['bottom'].set_linewidth(0.8)
    ax2.tick_params(axis='both', which='major', labelsize=12, length=4, width=0.8)
    
    # Adjust spacing between subplots
    plt.tight_layout()
    
    # Save plot
    filename = os.path.join(output_dir, 'plddt_histogram_comparison.png')
    plt.savefig(filename, dpi=300, bbox_inches='tight', facecolor='white')
    print(f"\nSaved pLDDT histogram to {filename}")
    plt.close()
    


def calculate_sequence_loss(editflow, sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=32):
    """
    Calculate sequence loss (val_unweighted_total_loss) for sequences.
    Uses the same loss calculation as validation_step in base_models.py.
    
    Args:
        editflow: EditFlow LightningModule
        sequences: List of sequence strings
        tokenizer: Tokenizer
        pad_id, bos_id, eos_id: Special token IDs
        device: torch device
        cfg: Config object
        batch_size: Batch size for processing
    
    Returns:
        List of loss values (one per sequence)
    """
    editflow.eval()
    all_losses = []
    
    print(f"Calculating sequence loss for {len(sequences)} sequences...")
    
    # Process in batches
    for i in tqdm(range(0, len(sequences), batch_size), desc="Computing losses"):
        batch_seqs = sequences[i:i+batch_size]
        
        # Tokenize batch
        x1_batch = []
        for seq_str in batch_seqs:
            x1 = tokenize_input_str(seq_str, cfg, tokenizer, bos_id, eos_id, pad_id, device)
            x1_batch.append(x1.squeeze(0))
        
        # Pad to same length
        max_len = max(x.shape[0] for x in x1_batch)
        padded_batch = []
        for x in x1_batch:
            padding = torch.full((max_len - x.shape[0],), pad_id, dtype=torch.long, device=device)
            padded_batch.append(torch.cat([x, padding]))
        x1_tensor = torch.stack(padded_batch).to(device)
        
        # Calculate loss for this batch
        with torch.no_grad():
            # Call preparation to get necessary tensors
            if editflow.reparameterize:
                lam_total, logits_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight, M_t = editflow.preparation(x1_tensor)
                
                if editflow.loc_prop_path:
                    loss, loss_components = editflow.loss_fn.reparameterized_forward_localized(
                        lam_total, logits_type, logits_ins, logits_sub, 
                        z_t, z_1, x_t, mask, weight, M_t, editflow.lam_prop, 
                        editflow.eps_id, editflow.bos_id, editflow.eos_id, 
                        editflow.gamma_rate, editflow.gamma_edit, 
                        editflow.use_aux_ce, editflow.aux_ce_weight
                    )
                else:
                    loss, loss_components = editflow.loss_fn.reparameterized_forward(
                        lam_total, logits_type, logits_ins, logits_sub,
                        z_t, z_1, x_t, mask, weight, editflow.eps_id, 
                        editflow.bos_id, editflow.eos_id, editflow.gamma_rate, 
                        editflow.gamma_edit, editflow.use_aux_ce, editflow.aux_ce_weight
                    )
            else:
                lam_ins, logits_ins, lam_del, lam_sub, logits_sub, z_t, z_1, x_t, mask, weight, M_t = editflow.preparation(x1_tensor)
                
                if editflow.loc_prop_path:
                    loss, loss_components = editflow.loss_fn.forward_localized(
                        lam_ins, logits_ins, lam_del, lam_sub, logits_sub, 
                        z_t, z_1, x_t, mask, weight, M_t, editflow.lam_prop, 
                        editflow.eps_id, editflow.bos_id, editflow.eos_id, 
                        editflow.use_aux_ce, editflow.aux_ce_weight
                    )
                else:
                    loss, loss_components = editflow.loss_fn.forward(
                        lam_ins, logits_ins, lam_del, lam_sub, logits_sub, 
                        z_t, z_1, x_t, mask, weight, editflow.eps_id, 
                        editflow.bos_id, editflow.eos_id, 
                        editflow.use_aux_ce, editflow.aux_ce_weight
                    )
            
            # Extract unweighted total loss (matching val_unweighted_total_loss)
            if "loss_total_unweighted" in loss_components:
                # For reparameterized models
                unweighted_loss = loss_components["loss_total_unweighted"]
            elif "loss_base" in loss_components:
                # For non-reparameterized models, loss_base is rate + edit (unweighted)
                unweighted_loss = loss_components["loss_base"]
            else:
                # Fallback to total loss
                unweighted_loss = loss
            
            # Get per-sequence losses (loss is averaged over batch, so we need to compute per-sequence)
            # Since loss is batch-averaged, we'll use the batch loss for all sequences in the batch
            # For more accurate per-sequence loss, we'd need to compute individually
            batch_loss_value = unweighted_loss.item()
            all_losses.extend([batch_loss_value] * len(batch_seqs))
    
    return all_losses


def kgrams(s: str, k: int):
    """Extract k-grams from a string."""
    s = s.strip()
    if len(s) < k:
        return {s} if s else set()
    return {s[i:i+k] for i in range(len(s) - k + 1)}


def jaccard(a: set, b: set) -> float:
    """Calculate Jaccard similarity between two sets."""
    if not a and not b:
        return 1.0
    inter = len(a & b)
    union = len(a | b)
    return inter / union if union else 1.0


def diversity_kmer_jaccard(seqs, k=3, pairs=50000, seed=0):
    """
    Diversity = 1 - average Jaccard similarity over random pairs.
    Works for variable-length strings.
    
    Returns:
        diversity (1 - avg_sim), avg_similarity
    """
    if len(seqs) < 2:
        return 0.0, 1.0
    
    rng = random.Random(seed)
    grams = [kgrams(s, k) for s in seqs]
    n = len(seqs)
    
    # Limit pairs to avoid excessive computation
    max_pairs = min(pairs, n * (n - 1) // 2)
    if max_pairs == 0:
        return 0.0, 1.0
    
    total_sim = 0.0
    for _ in range(max_pairs):
        i = rng.randrange(n)
        j = rng.randrange(n - 1)
        if j >= i:
            j += 1
        total_sim += jaccard(grams[i], grams[j])
    
    avg_sim = total_sim / max_pairs
    return 1.0 - avg_sim, avg_sim


def diversity_levenshtein(seqs, pairs=20000, seed=0):
    """
    Diversity = 1 - average normalized Levenshtein similarity over random pairs.
    
    Returns:
        diversity (1 - avg_sim), avg_similarity
    """
    if len(seqs) < 2:
        return 0.0, 1.0
    
    if not (HAS_RAPIDFUZZ or HAS_LEVENSHTEIN):
        # No Levenshtein implementation available
        return 0.0, 1.0
    
    rng = random.Random(seed)
    n = len(seqs)
    
    # Limit pairs to avoid excessive computation
    max_pairs = min(pairs, n * (n - 1) // 2)
    if max_pairs == 0:
        return 0.0, 1.0
    
    total_sim = 0.0
    valid_pairs = 0
    for _ in range(max_pairs):
        i = rng.randrange(n)
        j = rng.randrange(n - 1)
        if j >= i:
            j += 1
        
        if HAS_RAPIDFUZZ:
            # Use rapidfuzz for fast normalized similarity
            sim = RapidLevenshtein.normalized_similarity(seqs[i], seqs[j])
        elif HAS_LEVENSHTEIN:
            # Fallback: use python-Levenshtein distance and normalize
            # Simple normalization: 1 - (distance / max_length)
            dist = levenshtein_distance_fallback(seqs[i], seqs[j])
            max_len = max(len(seqs[i]), len(seqs[j]))
            sim = 1.0 - (dist / max_len) if max_len > 0 else 1.0
        else:
            # Should not reach here, but skip if somehow we do
            continue
        
        total_sim += sim
        valid_pairs += 1
    
    if valid_pairs == 0:
        return 0.0, 1.0
    
    avg_sim = total_sim / valid_pairs
    return 1.0 - avg_sim, avg_sim


def calculate_diversity(sequences):
    """
    Calculate diversity metrics for a set of sequences.
    Uses k-mer Jaccard diversity and optionally Levenshtein diversity.
    
    Args:
        sequences: List of sequence strings
    
    Returns:
        Dictionary with diversity metrics:
        - unique_count: Number of unique sequences
        - uniqueness_ratio: Fraction of unique sequences
        - kmer_diversity: k-mer Jaccard diversity (1 - avg_similarity)
        - kmer_avg_similarity: Average k-mer Jaccard similarity
        - levenshtein_diversity: Levenshtein diversity (if available)
        - levenshtein_avg_similarity: Average Levenshtein similarity (if available)
    """
    if len(sequences) == 0:
        return {
            'unique_count': 0,
            'uniqueness_ratio': 0.0,
            'kmer_diversity': 0.0,
            'kmer_avg_similarity': 1.0,
            'levenshtein_diversity': 0.0,
            'levenshtein_avg_similarity': 1.0
        }
    
    # Unique fraction
    unique_count = len(set(sequences))
    uniqueness_ratio = unique_count / len(sequences) if len(sequences) > 0 else 0.0
    
    # k-mer Jaccard diversity
    kmer_div, kmer_sim = diversity_kmer_jaccard(sequences, k=3, pairs=50000, seed=0)
    
    # Levenshtein diversity (if available)
    if HAS_RAPIDFUZZ or HAS_LEVENSHTEIN:
        lev_div, lev_sim = diversity_levenshtein(sequences, pairs=20000, seed=0)
    else:
        lev_div, lev_sim = 0.0, 1.0
    
    return {
        'unique_count': unique_count,
        'uniqueness_ratio': uniqueness_ratio,
        'kmer_diversity': kmer_div,
        'kmer_avg_similarity': kmer_sim,
        'levenshtein_diversity': lev_div,
        'levenshtein_avg_similarity': lev_sim
    }


def evaluate_cas9_scores(sequences, cas9_classifier, threshold=0.5):
    """
    Evaluate Cas9 scores for sequences.
    
    Args:
        sequences: List of sequence strings
        cas9_classifier: Cas9Classification object
        threshold: Score threshold for validity (default 0.5)
    
    Returns:
        validity_rate, average_score, list of scores
    """
    if len(sequences) == 0:
        return 0.0, 0.0, []
    
    # Get scores in batches
    scores = []
    batch_size = 32
    
    for i in tqdm(range(0, len(sequences), batch_size), desc="Computing Cas9 scores"):
        batch_seqs = sequences[i:i+batch_size]
        batch_scores = cas9_classifier.get_scores(batch_seqs)
        scores.extend(batch_scores)
    
    scores = np.array(scores)
    validity_rate = np.mean(scores > threshold)
    avg_score = np.mean(scores)
    
    return validity_rate, avg_score, scores.tolist()


def generate_sequences_multi_edit(model, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id,
                                  input_sequences, device, cfg, num_steps=20, batch_size=32, num_generations_per_sequence=1):
    """
    Generate sequences from input sequences using multi-edit generation.
    
    Args:
        model: Editflows model
        input_sequences: List of input sequence strings
        device: torch device
        cfg: Config object
        num_steps: Number of generation steps
        batch_size: Batch size for generation
        num_generations_per_sequence: Number of sequences to generate per input sequence
    
    Returns:
        List of generated sequence strings
    """
    model.eval()
    generated_sequences = []
    
    # Get allowed tokens
    allowed_tokens = torch.tensor(
        [tok for tok in source_dist._allowed_tokens if tok not in (eps_id,)],
        device=device,
        dtype=torch.long,
    )
    
    print(f"Generating sequences using multi-edit generation (num_steps={num_steps}, batch_size={batch_size}, num_generations_per_sequence={num_generations_per_sequence})...")
    print(f"Total sequences to generate: {len(input_sequences)} input sequences × {num_generations_per_sequence} = {len(input_sequences) * num_generations_per_sequence}")
    
    # Process in batches
    for i in tqdm(range(0, len(input_sequences), batch_size), desc="Generating"):
        batch_seqs = input_sequences[i:i+batch_size]
        
        # Generate num_generations_per_sequence for each sequence in the batch
        for gen_idx in range(num_generations_per_sequence):
            # Tokenize batch
            x0_batch = []
            for seq_str in batch_seqs:
                x0 = tokenize_input_str(seq_str, cfg, tokenizer, bos_id, eos_id, pad_id, device)
                x0_batch.append(x0.squeeze(0))
            
            # Pad to same length
            max_len = max(x.shape[0] for x in x0_batch)
            padded_batch = []
            for x in x0_batch:
                padding = torch.full((max_len - x.shape[0],), pad_id, dtype=torch.long, device=device)
                padded_batch.append(torch.cat([x, padding]))
            x0_tensor = torch.stack(padded_batch).to(device)
            
            # Generate using multi-edit generation
            with torch.no_grad():
                x_gen = generate_from_x0_multi_edit(
                    model,
                    x0_tensor,
                    pad_id=pad_id,
                    bos_id=bos_id,
                    eos_id=eos_id,
                    allowed_tokens=allowed_tokens,
                    num_steps=num_steps,
                    device=device,
                )
            
            # Decode generated sequences
            for j in range(x_gen.shape[0]):
                gen_seq = detokenize_output(x_gen[j:j+1], cfg, tokenizer, bos_id, eos_id, pad_id)
                generated_sequences.append(gen_seq)
    
    return generated_sequences


def get_evaluate_val_argument_parser():
    parser = argparse.ArgumentParser(description='Evaluate Editflows model by generating sequences')
    parser.add_argument('--config', type=str, required=True,
                        help='Path to config YAML file')
    parser.add_argument('--ckpt', type=str, required=True,
                        help='Path to model checkpoint (.ckpt file)')
    parser.add_argument('--val_data', type=str, default=None,
                        help='Path to validation dataset (output from uniref_pre_batching.py). Deprecated: use --test_data instead.')
    parser.add_argument('--test_data', type=str, default=None,
                        help='Path to test dataset (output from uniref_pre_batching.py)')
    parser.add_argument('--cas9_classifier_ckpt', type=str, default=None,
                        help='Path to Cas9 classifier checkpoint (optional, if not provided, validity rate will not be calculated)')
    parser.add_argument('--cas9_classifier_config', type=str, default=None,
                        help='Path to Cas9 classifier config (optional, required if --cas9_classifier_ckpt is provided)')
    parser.add_argument('--num_samples', type=int, default=1000,
                        help='Number of sequences to sample from test set (default: 1000)')
    parser.add_argument('--num_generations_per_sequence', type=int, default=1,
                        help='Number of sequences to generate per sampled sequence (default: 1). Total generated = num_samples × num_generations_per_sequence')
    parser.add_argument('--num_steps', type=int, default=20,
                        help='Number of generation steps (default: 20)')
    parser.add_argument('--batch_size', type=int, default=32,
                        help='Batch size for generation (default: 32)')
    parser.add_argument('--validity_threshold', type=float, default=0.5,
                        help='Cas9 score threshold for validity (default: 0.5)')
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                        help='Device to use (cuda/cpu)')
    parser.add_argument('--seed', type=int, default=42,
                        help='Random seed (default: 42)')
    parser.add_argument('--output_dir', type=str, default='./evaluation_output',
                        help='Directory to save pLDDT plots (default: ./evaluation_output)')
    parser.add_argument('--output_fasta', type=str, default=None,
                        help='Path to save generated sequences as FASTA file (default: {output_dir}/generated_sequences.fasta)')
    parser.add_argument('--num_bins', type=int, default=50,
                        help='Number of bins for pLDDT histogram (default: 50)')
    parser.add_argument('--calculate_plddt', action='store_true',
                        help='Enable pLDDT calculation and plotting (default: False, pLDDT calculation is disabled by default)')
    parser.add_argument('--include_train', action='store_true',
                        help='Also compute losses for the train set (default: False)')
    return parser


def _set_eval_run_seeds(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def run_evaluate_val(
    args,
    *,
    run_seed=None,
    save_fasta=True,
    fasta_tag=None,
):
    """
    Run the full evaluate_val pipeline and return metrics as a nested dict of JSON-serializable values.

    Args:
        args: Namespace from get_evaluate_val_argument_parser().parse_args() (or compatible).
        run_seed: If int, seeds random/numpy/torch before stochastic steps. If None, behavior matches
            legacy evaluate_val (no seeding).
        save_fasta: If False, skip writing generated sequences to FASTA.
        fasta_tag: If set (and save_fasta), write to output_dir/generated_sequences_{fasta_tag}.fasta
            unless output_fasta is explicitly set on args.
    """
    if run_seed is not None:
        _set_eval_run_seeds(int(run_seed))

    device = torch.device(args.device)
    
    # Load config
    with open(args.config, 'r') as f:
        cfg = edict(yaml.safe_load(f))
    
    # Build model
    print("Building model...")
    editflow, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id = build_model_and_stuff(cfg, device)
    
    # Load checkpoint
    print(f"Loading checkpoint from {args.ckpt}...")
    ckpt = torch.load(args.ckpt, map_location=device, weights_only=False)
    editflow.load_state_dict(ckpt["state_dict"], strict=False)
    model = editflow.model.to(device)
    model.eval()
    
    # Initialize Cas9 classifier (optional)
    cas9_classifier = None
    if args.cas9_classifier_ckpt is not None:
        if args.cas9_classifier_config is None:
            raise ValueError("--cas9_classifier_config is required when --cas9_classifier_ckpt is provided")
        
        print("Initializing Cas9 classifier...")
        cas9_classifier = Cas9Classification(
            device=device,
            checkpoint_path=args.cas9_classifier_ckpt,
            config_path=args.cas9_classifier_config,
            shared_esm_model=model.esm_emb  # Share ESM model to save memory
        )
        
        # Test classifier on a known Cas9 sequence to verify it works
        test_cas9_seq = "MDKKYSIGLDIGTNSVGWAVITDEYKVPSKKFKVLGNTDRHSIKKNLIGALLFDSGETAEATRLKRTARRRYTRRKNRICYLQEIFSNEMAKVDDSFFHRLEESFLVEEDKKHERHPIFGNIVDEVAYHEKYPTIYHLRKKLVDSTDKADLRLIYLALAHMIKFRGHFLIEGDLNPDNSDVDKLFIQLVQTYNQLFEENPINASGVDAKAILSARLSKSRRLENLIAQLPGEKKNGLFGNLIALSLGLTPNFKSNFDLAEDAKLQLSKDTYDDDLDNLLAQIGDQYADLFLAAKNLSDAILLSDILRVNTEITKAPLSASMIKRYDEHHQDLTLLKALVRQQLPEKYKEIFFDQSKNGYAGYIDGGASQEEFYKFIKPILEKMDGTEELLVKLNREDLLRKQRTFDNGSIPHQIHLGELHAILRRQEDFYPFLKDNREKIEKILTFRIPYYVGPLARGNSRFAWMTRKSEETITPWNFEEVVDKGASAQSFIERMTNFDKNLPNEKVLPKHSLLYEYFTVYNELTKVKYVTEGMRKPAFLSGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEISGVEDRFNASLGTYHDLLKIIKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQLKRRRYTGWGRLSRKLINGIRDKQSGKTILDFLKSDGFANRNFMQLIHDDSLTFKEDIQKAQVSGQGDSLHEHIANLAGSPAIKKGILQTVKVVDELVKVMGRHKPENIVIEMARENQTTQKGQKNSRERMKRIEEGIKELGSQILKEHPVENTQLQNEKLYLYYLQNGRDMYVDQELDINRLSDYDVDHIVPQSFLKDDSIDNKVLTRSDKNRGKSDNVPSEEVVKKMKNYWRQLLNAKLITQRKFDNLTKAERGGLSELDKAGFIKRQLVETRQITKHVAQILDSRMNTKYDENDKLIREVKVITLKSKLVSDFRKDFQFYKVREINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYGDYKVYDVRKMIAKSEQEIGKATAKYFFYSNIMNFFKTEITLANGEIRKRPLIETNGETGEIVWDKGRDFATVRKVLSMPQVNIVKKTEVQTGGFSKESILPKRNSDKLIARKKDWDPKKYGGFDSPTVAYSVLVVAKVEKGKSKKLKSVKELLGITIMERSSFEKNPIDFLEAKGYKEVKKDLIIKLPKYSLFELENGRKRMLASAGELQKGNELALPSKYVNFLYLASHYEKLKGSPEDNEQKQLFVEQHKHYLDEIIEQISEFSKRVILADANLDKVLSAYNKHRDKPIREQAENIIHLFTLTNLGAPAAFKYFDTTIDRKRYTSTKEVLDATLIHQSITGLYETRIDLSQLGGD"
        test_score = cas9_classifier.get_scores([test_cas9_seq])[0]
        print(f"Test Cas9 classifier on known Cas9 sequence: score = {test_score:.4f}")
        if test_score < 0.5:
            print(f"WARNING: Known Cas9 sequence got low score! This suggests a problem with the classifier or sequence format.")
        else:
            print(f"Classifier working correctly (score > 0.5 for known Cas9 sequence)")
    else:
        print("No Cas9 classifier provided. Validity rate will not be calculated.")
    
    # Determine which dataset path to use (test_data takes precedence)
    if args.test_data is not None:
        dataset_path = args.test_data
        split_name = "test"
        dataset_type = "test"
    elif args.val_data is not None:
        dataset_path = args.val_data
        # Try "validation" or "val" first when --val_data is used
        split_name = None
        dataset_type = "validation"
    else:
        raise ValueError("Either --test_data or --val_data must be provided")
    
    # Load dataset
    print(f"\nLoading {dataset_type} dataset from {dataset_path}...")
    dataset_dict = load_from_disk(dataset_path)
    
    # Determine which split to use
    if split_name is None:
        # When --val_data is used, prioritize validation splits
        if "validation" in dataset_dict:
            test_dataset = dataset_dict["validation"]
            split_name = "validation"
        elif "val" in dataset_dict:
            test_dataset = dataset_dict["val"]
            split_name = "val"
        elif "test" in dataset_dict:
            # Fallback to test if validation splits not found
            test_dataset = dataset_dict["test"]
            split_name = "test"
            dataset_type = "test"
        else:
            raise ValueError(f"Dataset must have 'validation', 'val', or 'test' split. Found splits: {list(dataset_dict.keys())}")
    else:
        if split_name in dataset_dict:
            test_dataset = dataset_dict[split_name]
        else:
            raise ValueError(f"Dataset must have '{split_name}' split. Found splits: {list(dataset_dict.keys())}")
    
    # Load train dataset if --include_train flag is set
    train_dataset = None
    all_train_sequences = None
    train_avg_loss = None
    if args.include_train:
        if "train" in dataset_dict:
            train_dataset = dataset_dict["train"]
            print(f"\nLoading train dataset from {dataset_path}...")
            print(f"Train dataset loaded: {len(train_dataset)} items")
        else:
            print(f"Warning: --include_train flag set but 'train' split not found in dataset. Found splits: {list(dataset_dict.keys())}")
            print("Skipping train set evaluation.")
    
    # Step 1: Decode test set into string sequences
    print(f"\n{'='*70}")
    print("STEP 1: Decoding test set")
    print(f"{'='*70}")
    all_decoded_sequences, sampled_decoded_sequences = decode_validation_set(
        test_dataset,
        tokenizer,
        pad_id,
        bos_id,
        eos_id,
        num_samples=args.num_samples
    )
    
    # Step 2: Calculate metrics on decoded sequences
    print(f"\n{'='*70}")
    print("STEP 2: Evaluating decoded sequences")
    print(f"{'='*70}")
    
    # Calculate sequence loss on FULL test set
    print(f"\nCalculating sequence loss for FULL test set ({len(all_decoded_sequences)} sequences)...")
    orig_losses_full = calculate_sequence_loss(
        editflow, all_decoded_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
    )
    orig_avg_loss_full = np.mean(orig_losses_full)
    print(f"  Average sequence loss (full test set): {orig_avg_loss_full:.4f}")
    
    # Calculate other metrics on sampled sequences only
    print(f"\nCalculating diversity for sampled decoded sequences ({len(sampled_decoded_sequences)} sequences)...")
    orig_diversity = calculate_diversity(sampled_decoded_sequences)
    
    # Calculate sequence loss for sampled sequences (for comparison)
    print(f"\nCalculating sequence loss for sampled decoded sequences ({len(sampled_decoded_sequences)} sequences)...")
    orig_losses = calculate_sequence_loss(
        editflow, sampled_decoded_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
    )
    orig_avg_loss = np.mean(orig_losses)
    print(f"  Average sequence loss (sampled): {orig_avg_loss:.4f}")
    
    # Calculate sequence loss for train set if --include_train flag is set
    if args.include_train and train_dataset is not None:
        print(f"\n{'='*70}")
        print("TRAIN SET EVALUATION")
        print(f"{'='*70}")
        print(f"\nDecoding train set into string sequences...")
        all_train_sequences, _ = decode_validation_set(
            train_dataset,
            tokenizer,
            pad_id,
            bos_id,
            eos_id,
            num_samples=10**9  # Decode all sequences, don't sample (use very large number)
        )
        print(f"\nCalculating sequence loss for FULL train set ({len(all_train_sequences)} sequences)...")
        train_losses = calculate_sequence_loss(
            editflow, all_train_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
        )
        train_avg_loss = np.mean(train_losses)
        print(f"  Average sequence loss (full train set): {train_avg_loss:.4f}")
    
    # Calculate Cas9 scores only if classifier is provided (on sampled sequences)
    orig_validity_rate, orig_avg_score, orig_scores = None, None, None
    if cas9_classifier is not None:
        print(f"\nCalculating Cas9 scores for sampled decoded sequences...")
        orig_validity_rate, orig_avg_score, orig_scores = evaluate_cas9_scores(
            sampled_decoded_sequences, cas9_classifier, threshold=args.validity_threshold
        )
    
    # Step 3: Generate new sequences using multi-edit generation
    print(f"\n{'='*70}")
    print("STEP 3: Generating new sequences using multi-edit generation")
    print(f"{'='*70}")
    generated_sequences = generate_sequences_multi_edit(
        model, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id,
        sampled_decoded_sequences, device, cfg, num_steps=args.num_steps, batch_size=args.batch_size,
        num_generations_per_sequence=args.num_generations_per_sequence
    )
    print(f"Generated {len(generated_sequences)} sequences")
    
    # Save generated sequences as FASTA file
    fasta_path_written = None
    if save_fasta:
        if args.output_fasta is None:
            os.makedirs(args.output_dir, exist_ok=True)
            if fasta_tag:
                fasta_path = os.path.join(args.output_dir, f"generated_sequences_{fasta_tag}.fasta")
            else:
                fasta_path = os.path.join(args.output_dir, "generated_sequences.fasta")
        else:
            fasta_path = args.output_fasta
            fasta_dir = os.path.dirname(fasta_path)
            if fasta_dir:
                os.makedirs(fasta_dir, exist_ok=True)

        print(f"\nSaving generated sequences to FASTA file: {fasta_path}")
        with open(fasta_path, 'w') as f:
            for i, seq in enumerate(generated_sequences):
                f.write(f">generated_sequence_{i+1}\n")
                for j in range(0, len(seq), 80):
                    f.write(seq[j:j+80] + "\n")
        print(f"Saved {len(generated_sequences)} sequences to {fasta_path}")
        fasta_path_written = fasta_path
    
    # Step 4: Calculate metrics on generated sequences
    print(f"\n{'='*70}")
    print("STEP 4: Evaluating generated sequences")
    print(f"{'='*70}")
    
    print(f"\nCalculating diversity for generated sequences...")
    gen_diversity = calculate_diversity(generated_sequences)
    
    # Calculate sequence loss for generated sequences
    print(f"\nCalculating sequence loss for generated sequences...")
    gen_losses = calculate_sequence_loss(
        editflow, generated_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
    )
    gen_avg_loss = np.mean(gen_losses)
    print(f"  Average sequence loss: {gen_avg_loss:.4f}")
    
    # Calculate Cas9 scores only if classifier is provided
    gen_validity_rate, gen_avg_score, gen_scores = None, None, None
    if cas9_classifier is not None:
        print(f"\nCalculating Cas9 scores for generated sequences...")
        gen_validity_rate, gen_avg_score, gen_scores = evaluate_cas9_scores(
            generated_sequences, cas9_classifier, threshold=args.validity_threshold
        )
    
    def _to_float(x):
        if x is None:
            return None
        return float(x)

    def _diversity_plain(d):
        return {k: int(v) if k == "unique_count" else float(v) for k, v in d.items()}

    metrics = {
        "run_seed": int(run_seed) if run_seed is not None else None,
        "dataset_type": dataset_type,
        "fasta_path": fasta_path_written,
        "counts": {
            "full_train": len(all_train_sequences) if all_train_sequences is not None else None,
            "full_test": len(all_decoded_sequences),
            "sampled_decoded": len(sampled_decoded_sequences),
            "generated": len(generated_sequences),
        },
        "loss": {
            "train_avg": _to_float(train_avg_loss),
            "decoded_full_avg": _to_float(orig_avg_loss_full),
            "decoded_sampled_avg": _to_float(orig_avg_loss),
            "generated_avg": _to_float(gen_avg_loss),
        },
        "cas9": None,
        "diversity_decoded": _diversity_plain(orig_diversity),
        "diversity_generated": _diversity_plain(gen_diversity),
        "plddt": None,
    }
    if cas9_classifier is not None:
        metrics["cas9"] = {
            "decoded_validity_rate": _to_float(orig_validity_rate),
            "decoded_avg_score": _to_float(orig_avg_score),
            "generated_validity_rate": _to_float(gen_validity_rate),
            "generated_avg_score": _to_float(gen_avg_score),
        }

    # Step 5: Calculate and plot pLDDT scores (optional)
    if args.calculate_plddt:
        print(f"\n{'='*70}")
        print("STEP 5: Calculating pLDDT scores")
        print(f"{'='*70}")

        print("\nClearing GPU memory...")
        del model
        del editflow
        if cas9_classifier is not None:
            del cas9_classifier
        if device.type == 'cuda':
            torch.cuda.empty_cache()
            print("GPU memory cleared.")

        print("Loading ESMFold model for pLDDT calculation...")
        esmfold_tokenizer_path = "facebook/esmfold_v1"
        esmfold_tokenizer = AutoTokenizer.from_pretrained(esmfold_tokenizer_path)
        esmfold_device = device
        esm_model = EsmForProteinFolding.from_pretrained(
            esmfold_tokenizer_path,
            torch_dtype=torch.bfloat16
        ).to(esmfold_device).eval()
        print("ESMFold model loaded successfully!")

        print(f"\nCalculating pLDDT scores for sampled decoded sequences...")
        decoded_plddt = calculate_plddt_scores(
            sampled_decoded_sequences, esmfold_tokenizer, esm_model, esmfold_device, batch_size=1
        )
        decoded_plddt_valid = [s for s in decoded_plddt if s is not None]
        if len(decoded_plddt_valid) > 0:
            print(f"  Valid scores: {len(decoded_plddt_valid)}/{len(decoded_plddt)}")
            print(f"  Mean pLDDT: {np.mean(decoded_plddt_valid):.2f}, Std: {np.std(decoded_plddt_valid):.2f}")

        print(f"\nCalculating pLDDT scores for generated sequences...")
        generated_plddt = calculate_plddt_scores(
            generated_sequences, esmfold_tokenizer, esm_model, esmfold_device, batch_size=1
        )
        generated_plddt_valid = [g for g in generated_plddt if g is not None]
        if len(generated_plddt_valid) > 0:
            print(f"  Valid scores: {len(generated_plddt_valid)}/{len(generated_plddt)}")
            print(f"  Mean pLDDT: {np.mean(generated_plddt_valid):.2f}, Std: {np.std(generated_plddt_valid):.2f}")

        if len(decoded_plddt_valid) > 0 and len(generated_plddt_valid) > 0:
            print(f"\nPlotting pLDDT histograms...")
            plot_plddt_histogram(decoded_plddt_valid, generated_plddt_valid, args.output_dir, args.num_bins, dataset_type)
        else:
            print("Warning: Not enough valid pLDDT scores to plot.")

        metrics["plddt"] = {
            "decoded_mean": float(np.mean(decoded_plddt_valid)) if decoded_plddt_valid else None,
            "decoded_std": float(np.std(decoded_plddt_valid)) if decoded_plddt_valid else None,
            "decoded_valid_n": len(decoded_plddt_valid),
            "decoded_total_n": len(decoded_plddt),
            "generated_mean": float(np.mean(generated_plddt_valid)) if generated_plddt_valid else None,
            "generated_std": float(np.std(generated_plddt_valid)) if generated_plddt_valid else None,
            "generated_valid_n": len(generated_plddt_valid),
            "generated_total_n": len(generated_plddt),
        }

    return metrics


def print_evaluation_results(metrics, args):
    """Print the summary block (same layout as historical evaluate_val.py)."""
    dataset_type = metrics["dataset_type"]
    counts = metrics["counts"]
    loss_m = metrics["loss"]
    orig_diversity = metrics["diversity_decoded"]
    gen_diversity = metrics["diversity_generated"]

    print("\n" + "="*70)
    print("EVALUATION RESULTS")
    print("="*70)

    print(f"\nDataset Statistics:")
    if counts["full_train"] is not None:
        print(f"  Full train set sequences: {counts['full_train']}")
    print(f"  Full test set sequences: {counts['full_test']}")
    print(f"  Sampled decoded sequences: {counts['sampled_decoded']}")
    print(f"  Generated sequences: {counts['generated']}")

    print(f"\n" + "-"*70)
    print("SEQUENCE LOSS (val_unweighted_total_loss)")
    print("-"*70)
    if loss_m["train_avg"] is not None:
        print(f"Train set (FULL train set, {counts['full_train']} sequences):")
        print(f"  Average loss: {loss_m['train_avg']:.4f}")
    print(f"\nDecoded sequences (FULL {dataset_type} set, {counts['full_test']} sequences):")
    print(f"  Average loss: {loss_m['decoded_full_avg']:.4f}")
    print(f"\nDecoded sequences (sampled, {counts['sampled_decoded']} sequences):")
    print(f"  Average loss: {loss_m['decoded_sampled_avg']:.4f}")
    print(f"\nGenerated sequences ({counts['generated']} sequences):")
    print(f"  Average loss: {loss_m['generated_avg']:.4f}")

    cas9 = metrics["cas9"]
    if cas9 is not None:
        print(f"\n" + "-"*70)
        print("CAS9 SCORES (validity threshold = {})".format(args.validity_threshold))
        print("-"*70)
        print(f"NOTE: If {dataset_type} set contains general proteins (not Cas9-specific),")
        print("      low Cas9 scores are expected. Cas9 classifier is trained to identify Cas9 proteins.")
        print(f"\nDecoded sequences (sampled from {dataset_type} set, {counts['sampled_decoded']} sequences):")
        print(f"  Validity rate: {cas9['decoded_validity_rate']:.4f} ({cas9['decoded_validity_rate']*100:.2f}%)")
        print(f"  Average Cas9 score: {cas9['decoded_avg_score']:.4f}")
        print(f"\nGenerated sequences ({counts['generated']} sequences):")
        print(f"  Validity rate: {cas9['generated_validity_rate']:.4f} ({cas9['generated_validity_rate']*100:.2f}%)")
        print(f"  Average Cas9 score: {cas9['generated_avg_score']:.4f}")

    print(f"\n" + "-"*70)
    print("DIVERSITY METRICS")
    print("-"*70)
    print(f"Decoded sequences (sampled from {dataset_type} set, {counts['sampled_decoded']} sequences):")
    print(f"  Unique sequences: {orig_diversity['unique_count']} / {counts['sampled_decoded']} ({orig_diversity['uniqueness_ratio']*100:.2f}%)")
    print(f"  k-mer Jaccard diversity (k=3): {orig_diversity['kmer_diversity']:.4f}")
    print(f"  k-mer Jaccard avg similarity: {orig_diversity['kmer_avg_similarity']:.4f}")
    if HAS_RAPIDFUZZ or HAS_LEVENSHTEIN:
        print(f"  Levenshtein diversity: {orig_diversity['levenshtein_diversity']:.4f}")
        print(f"  Levenshtein avg similarity: {orig_diversity['levenshtein_avg_similarity']:.4f}")
    print(f"\nGenerated sequences ({counts['generated']} sequences):")
    print(f"  Unique sequences: {gen_diversity['unique_count']} / {counts['generated']} ({gen_diversity['uniqueness_ratio']*100:.2f}%)")
    print(f"  k-mer Jaccard diversity (k=3): {gen_diversity['kmer_diversity']:.4f}")
    print(f"  k-mer Jaccard avg similarity: {gen_diversity['kmer_avg_similarity']:.4f}")
    if HAS_RAPIDFUZZ or HAS_LEVENSHTEIN:
        print(f"  Levenshtein diversity: {gen_diversity['levenshtein_diversity']:.4f}")
        print(f"  Levenshtein avg similarity: {gen_diversity['levenshtein_avg_similarity']:.4f}")

    plddt = metrics.get("plddt")
    if plddt is not None:
        print(f"\n" + "-"*70)
        print("PLDDT METRICS")
        print("-"*70)
        print(f"Decoded sequences (sampled from {dataset_type} set, {counts['sampled_decoded']} sequences):")
        if plddt["decoded_mean"] is not None:
            print(f"  Mean pLDDT: {plddt['decoded_mean']:.2f}")
            print(f"  Std pLDDT: {plddt['decoded_std']:.2f}")
        print(f"  Valid scores: {plddt['decoded_valid_n']}/{plddt['decoded_total_n']}")
        print(f"\nGenerated sequences ({counts['generated']} sequences):")
        if plddt["generated_mean"] is not None:
            print(f"  Mean pLDDT: {plddt['generated_mean']:.2f}")
            print(f"  Std pLDDT: {plddt['generated_std']:.2f}")
        print(f"  Valid scores: {plddt['generated_valid_n']}/{plddt['generated_total_n']}")

    print("="*70)


def main():
    parser = get_evaluate_val_argument_parser()
    args = parser.parse_args()
    # Legacy behavior: do not fix RNG (matches previous commented-out seed lines).
    metrics = run_evaluate_val(args)
    print_evaluation_results(metrics, args)


if __name__ == "__main__":
    main()