evaluate_optimizer.py

import sys
"""
File: evaluate_optimizer.py
---------------------------
Evaluate ColiFormer with enhanced capabilities:
1) DNAChisel post-processing for sequence polishing
2) Optional multi-objective generation (Pareto-style filtering)
3) Enhanced beam search with multiple candidates
4) Comprehensive metrics and optional ablation studies
"""

import argparse
import json
import os
import warnings
from typing import Dict, List, Tuple, Any

import numpy as np
import pandas as pd
import torch
from CAI import CAI, relative_adaptiveness
from tqdm import tqdm

from CodonTransformer.CodonData import (
    download_codon_frequencies_from_kazusa,
    get_codon_frequencies,
)
from CodonTransformer.CodonPrediction import (
    load_model,
    predict_dna_sequence,
    get_high_frequency_choice_sequence_optimized,
)
from CodonTransformer.CodonEvaluation import (
    calculate_dtw_distance,
    calculate_homopolymer_runs,
    calculate_tAI,
    count_negative_cis_elements,
    get_GC_content,
    get_ecoli_tai_weights,
    get_min_max_profile,
    get_sequence_similarity,
    scan_for_restriction_sites,
    calculate_ENC,
    calculate_CPB,
    calculate_SCUO,
)
from CodonTransformer.CodonPostProcessing import (
    polish_sequence_with_dnachisel,
)
from CodonTransformer.CodonUtils import DNASequencePrediction


def translate_dna_to_protein(dna_sequence: str) -> str:
    """Translate DNA sequence to protein sequence."""
    codon_table = {
        'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L',
        'TCT': 'S', 'TCC': 'S', 'TCA': 'S', 'TCG': 'S',
        'TAT': 'Y', 'TAC': 'Y', 'TAA': '*', 'TAG': '*',
        'TGT': 'C', 'TGC': 'C', 'TGA': '*', 'TGG': 'W',
        'CTT': 'L', 'CTC': 'L', 'CTA': 'L', 'CTG': 'L',
        'CCT': 'P', 'CCC': 'P', 'CCA': 'P', 'CCG': 'P',
        'CAT': 'H', 'CAC': 'H', 'CAA': 'Q', 'CAG': 'Q',
        'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
        'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M',
        'ACT': 'T', 'ACC': 'T', 'ACA': 'T', 'ACG': 'T',
        'AAT': 'N', 'AAC': 'N', 'AAA': 'K', 'AAG': 'K',
        'AGT': 'S', 'AGC': 'S', 'AGA': 'R', 'AGG': 'R',
        'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
        'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A',
        'GAT': 'D', 'GAC': 'D', 'GAA': 'E', 'GAG': 'E',
        'GGT': 'G', 'GGC': 'G', 'GGA': 'G', 'GGG': 'G'
    }
    
    protein = ""
    for i in range(0, len(dna_sequence), 3):
        codon = dna_sequence[i:i+3].upper()
        if len(codon) == 3:
            aa = codon_table.get(codon, 'X')
            if aa == '*':  # Stop codon
                break
            protein += aa
    
    return protein


def evaluate_with_enhancements(
    protein_sequence: str,
    model,
    tokenizer,
    device,
    cai_weights: Dict[str, float],
    tai_weights: Dict[str, float],
    codon_frequencies: Dict,
    reference_profile: List[float],
    args,
) -> Dict[str, Any]:
    """
    Evaluate a protein sequence with enhanced generation techniques.
    
    Args:
        protein_sequence: Input protein sequence
        model: Fine-tuned model
        tokenizer: Model tokenizer
        device: PyTorch device
        cai_weights: CAI weights dictionary
        tai_weights: tAI weights dictionary
        codon_frequencies: Codon frequencies dictionary
        reference_profile: Reference profile for DTW calculation
        args: Command line arguments
        
    Returns:
        Dict containing evaluation results for all methods
    """
    results = {}
    
    # 1. Original fine-tuned model (baseline)
    try:
        original_output = predict_dna_sequence(
            protein=protein_sequence,
            organism="Escherichia coli general",
            device=device,
            model=model,
            deterministic=True,
            match_protein=True,
            use_constrained_search=args.use_constrained_search,
            gc_bounds=tuple(args.gc_bounds),
            beam_size=args.beam_size,
            length_penalty=args.length_penalty,
            diversity_penalty=args.diversity_penalty,
        )
        
        if isinstance(original_output, list):
            original_dna = original_output[0].predicted_dna
        else:
            original_dna = original_output.predicted_dna
            
        results['fine_tuned_original'] = {
            'dna_sequence': original_dna,
            'method': 'fine_tuned_original',
            'enhancement': 'none',
        }
        
    except Exception as e:
        print(f"Warning: Original fine-tuned generation failed: {str(e)}")
        results['fine_tuned_original'] = {
            'dna_sequence': '',
            'method': 'fine_tuned_original',
            'enhancement': 'none',
            'error': str(e),
        }
    
    # 2. Enhanced sequence generation (DNAChisel + Pareto filtering)
    if args.use_enhanced_generation:
        try:
            enhanced_dna, generation_report = enhanced_sequence_generation(
                protein_sequence=protein_sequence,
                model=model,
                tokenizer=tokenizer,
                device=device,
                beam_size=args.enhanced_beam_size,
                gc_bounds=(args.gc_bounds[0] * 100, args.gc_bounds[1] * 100),
                use_dnachisel_polish=args.use_dnachisel,
                use_pareto_filtering=args.use_pareto_filtering,
                cai_weights=cai_weights,
                tai_weights=tai_weights,
                codon_frequencies=codon_frequencies,
                reference_profile=reference_profile,
            )
            
            results['fine_tuned_enhanced'] = {
                'dna_sequence': enhanced_dna,
                'method': 'fine_tuned_enhanced',
                'enhancement': 'dnachisel+pareto',
                'generation_report': generation_report,
            }
            
        except Exception as e:
            print(f"Warning: Enhanced generation failed: {str(e)}")
            results['fine_tuned_enhanced'] = {
                'dna_sequence': '',
                'method': 'fine_tuned_enhanced',
                'enhancement': 'dnachisel+pareto',
                'error': str(e),
            }
    
    # 3. DNAChisel post-processing only (ablation study)
    if args.use_dnachisel and 'fine_tuned_original' in results and results['fine_tuned_original']['dna_sequence']:
        try:
            dnachisel_dna, polish_report = polish_sequence_with_dnachisel(
                dna_sequence=results['fine_tuned_original']['dna_sequence'],
                protein_sequence=protein_sequence,
                gc_bounds=(args.gc_bounds[0] * 100, args.gc_bounds[1] * 100),
                maximize_cai=True,
                seed=42,
            )
            
            results['fine_tuned_dnachisel'] = {
                'dna_sequence': dnachisel_dna,
                'method': 'fine_tuned_dnachisel',
                'enhancement': 'dnachisel_only',
                'polish_report': polish_report,
            }
            
        except Exception as e:
            print(f"Warning: DNAChisel post-processing failed: {str(e)}")
            results['fine_tuned_dnachisel'] = {
                'dna_sequence': '',
                'method': 'fine_tuned_dnachisel',
                'enhancement': 'dnachisel_only',
                'error': str(e),
            }
    
    return results


def calculate_comprehensive_metrics(
    dna_sequence: str,
    protein_sequence: str,
    cai_weights: Dict[str, float],
    tai_weights: Dict[str, float],
    codon_frequencies: Dict,
    reference_profile: List[float],
    ref_sequences: List[str],
) -> Dict[str, float]:
    """Calculate comprehensive metrics for a DNA sequence."""
    if not dna_sequence:
        return {
            'cai': 0.0,
            'tai': 0.0,
            'gc_content': 0.0,
            'restriction_sites': float('inf'),
            'neg_cis_elements': float('inf'),
            'homopolymer_runs': float('inf'),
            'dtw_distance': float('inf'),
            'enc': 0.0,
            'cpb': 0.0,
            'scuo': 0.0,
        }
    
    return calculate_sequence_metrics(
        dna_sequence=dna_sequence,
        protein_sequence=protein_sequence,
        cai_weights=cai_weights,
        tai_weights=tai_weights,
        codon_frequencies=codon_frequencies,
        reference_profile=reference_profile,
    )


def run_ablation_study(results_df: pd.DataFrame) -> pd.DataFrame:
    """
    Run ablation study to compare different enhancement methods.
    
    Args:
        results_df: DataFrame with evaluation results
        
    Returns:
        DataFrame with ablation study results
    """
    # Group by protein and calculate improvements
    ablation_results = []
    
    for protein in results_df['protein_sequence'].unique():
        protein_results = results_df[results_df['protein_sequence'] == protein]
        
        # Get baseline (original fine-tuned)
        baseline = protein_results[protein_results['method'] == 'fine_tuned_original']
        if baseline.empty:
            continue
            
        baseline_metrics = baseline.iloc[0]
        
        # Compare each enhancement method
        for method in protein_results['method'].unique():
            if method == 'fine_tuned_original':
                continue
                
            method_results = protein_results[protein_results['method'] == method]
            if method_results.empty:
                continue
                
            method_metrics = method_results.iloc[0]
            
            # Calculate improvements
            improvements = {
                'protein': protein,
                'method': method,
                'enhancement': method_metrics['enhancement'],
                'cai_improvement': method_metrics['cai'] - baseline_metrics['cai'],
                'tai_improvement': method_metrics['tai'] - baseline_metrics['tai'],
                'gc_improvement': abs(method_metrics['gc_content'] - 52) - abs(baseline_metrics['gc_content'] - 52),
                'restriction_sites_improvement': baseline_metrics['restriction_sites'] - method_metrics['restriction_sites'],
                'neg_cis_improvement': baseline_metrics['neg_cis_elements'] - method_metrics['neg_cis_elements'],
                'homopolymer_improvement': baseline_metrics['homopolymer_runs'] - method_metrics['homopolymer_runs'],
                'dtw_improvement': baseline_metrics['dtw_distance'] - method_metrics['dtw_distance'],
                'composite_score_improvement': (
                    (method_metrics['cai'] - baseline_metrics['cai']) * 0.3 +
                    (method_metrics['tai'] - baseline_metrics['tai']) * 0.3 +
                    (abs(baseline_metrics['gc_content'] - 52) - abs(method_metrics['gc_content'] - 52)) * 0.2 +
                    (baseline_metrics['restriction_sites'] - method_metrics['restriction_sites']) * 0.1 +
                    (baseline_metrics['neg_cis_elements'] - method_metrics['neg_cis_elements']) * 0.1
                ),
            }
            
            ablation_results.append(improvements)
    
    return pd.DataFrame(ablation_results)


def main(args):
    """Main function to run the enhanced evaluation."""
    print("=== Enhanced CodonTransformer Evaluation ===")
    
    # Setup device
    device = torch.device("cuda" if torch.cuda.is_available() and args.use_gpu else "cpu")
    print(f"Using device: {device}")
    
    # Load test data
    with open(args.test_data_path, "r") as f:
        first = f.read(1)
        f.seek(0)
        if first == "[":
            test_set = json.load(f)
        else:
            test_set = [json.loads(line) for line in f if line.strip()]
    
    # Limit test set size if requested
    if args.max_test_proteins > 0:
        test_set = test_set[:args.max_test_proteins]
    
    print(f"Loaded {len(test_set)} proteins from the test set.")
    
    # Load models
    print("Loading models...")
    finetuned_model = load_model(model_path=args.checkpoint_path, device=device)
    print(f"Fine-tuned model loaded from {args.checkpoint_path}")
    
    # Load tokenizer
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained("adibvafa/CodonTransformer")
    
    # Load base model if comparison requested
    base_model = None
    if args.compare_with_base:
        base_model = load_model(device=device)
        print("Base model loaded from Hugging Face")
    
    # Prepare evaluation utilities
    print("Preparing evaluation utilities...")
    
    # CAI weights
    natural_csv = args.natural_sequences_path
    natural_df = pd.read_csv(natural_csv)
    ref_sequences = natural_df['dna_sequence'].tolist()
    cai_weights = relative_adaptiveness(sequences=ref_sequences)
    print("CAI weights generated")
    
    # tAI weights
    tai_weights = get_ecoli_tai_weights()
    print("tAI weights loaded")
    
    # Codon frequencies
    try:
        codon_frequencies = download_codon_frequencies_from_kazusa(taxonomy_id=83333)
        print("Codon frequencies loaded from Kazusa")
    except Exception as e:
        print(f"Warning: Kazusa download failed ({e}). Using local frequencies.")
        codon_frequencies = get_codon_frequencies(
            ref_sequences, organism="Escherichia coli general"
        )
    
    # Reference profile for DTW
    reference_profiles = [
        get_min_max_profile(seq, codon_frequencies) for seq in ref_sequences[:100]
    ]
    valid_profiles = [p for p in reference_profiles if p and not all(v is None for v in p)]
    
    if valid_profiles:
        max_len = max(len(p) for p in valid_profiles)
        padded_profiles = [
            np.pad(
                np.array([v for v in p if v is not None]),
                (0, max_len - len([v for v in p if v is not None])),
                "constant",
                constant_values=np.nan,
            )
            for p in valid_profiles
        ]
        avg_reference_profile = np.nanmean(padded_profiles, axis=0).tolist()
    else:
        avg_reference_profile = []
    
    print("Reference profile generated")
    
    # Run evaluation
    all_results = []
    evaluation_reports = []
    
    print("Starting enhanced evaluation...")
    for i, item in enumerate(tqdm(test_set, desc="Evaluating proteins")):
        # Get protein sequence
        if "protein_sequence" in item:
            protein_sequence = item["protein_sequence"]
        else:
            dna_sequence = item["codons"]
            protein_sequence = translate_dna_to_protein(dna_sequence)
        
        # Skip if protein is too short or too long
        if len(protein_sequence) < 10 or len(protein_sequence) > 1000:
            continue
        
        # Evaluate with enhancements
        protein_results = evaluate_with_enhancements(
            protein_sequence=protein_sequence,
            model=finetuned_model,
            tokenizer=tokenizer,
            device=device,
            cai_weights=cai_weights,
            tai_weights=tai_weights,
            codon_frequencies=codon_frequencies,
            reference_profile=avg_reference_profile,
            args=args,
        )
        
        # Add base model comparison if requested
        if base_model:
            try:
                base_output = predict_dna_sequence(
                    protein=protein_sequence,
                    organism="Escherichia coli general",
                    device=device,
                    model=base_model,
                    deterministic=True,
                    match_protein=True,
                )
                base_dna = base_output.predicted_dna if not isinstance(base_output, list) else base_output[0].predicted_dna
                
                protein_results['base_model'] = {
                    'dna_sequence': base_dna,
                    'method': 'base_model',
                    'enhancement': 'none',
                }
            except Exception as e:
                print(f"Warning: Base model generation failed: {str(e)}")
        
        # Add naive baseline
        try:
            naive_dna = get_high_frequency_choice_sequence_optimized(
                protein=protein_sequence, codon_frequencies=codon_frequencies
            )
            protein_results['naive_hfc'] = {
                'dna_sequence': naive_dna,
                'method': 'naive_hfc',
                'enhancement': 'none',
            }
        except Exception as e:
            print(f"Warning: Naive HFC generation failed: {str(e)}")
        
        # Calculate metrics for each method
        for method_name, method_result in protein_results.items():
            if 'error' in method_result:
                continue
                
            dna_seq = method_result['dna_sequence']
            if not dna_seq:
                continue
            
            metrics = calculate_comprehensive_metrics(
                dna_sequence=dna_seq,
                protein_sequence=protein_sequence,
                cai_weights=cai_weights,
                tai_weights=tai_weights,
                codon_frequencies=codon_frequencies,
                reference_profile=avg_reference_profile,
                ref_sequences=ref_sequences,
            )
            
            # Combine results
            result_row = {
                'protein_id': i,
                'protein_sequence': protein_sequence,
                'protein_length': len(protein_sequence),
                'method': method_name,
                'enhancement': method_result['enhancement'],
                'dna_sequence': dna_seq,
                'dna_length': len(dna_seq),
                **metrics,
            }
            
            # Add generation reports if available
            if 'generation_report' in method_result:
                result_row['generation_report'] = str(method_result['generation_report'])
            if 'polish_report' in method_result:
                result_row['polish_report'] = str(method_result['polish_report'])
            
            all_results.append(result_row)
    
    # Create results DataFrame
    results_df = pd.DataFrame(all_results)
    
    # Save detailed results
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    results_df.to_csv(args.output_path, index=False)
    print(f"Detailed results saved to {args.output_path}")
    
    # Run ablation study
    if args.run_ablation_study:
        ablation_df = run_ablation_study(results_df)
        ablation_path = args.output_path.replace('.csv', '_ablation.csv')
        ablation_df.to_csv(ablation_path, index=False)
        print(f"Ablation study results saved to {ablation_path}")
        
        # Print summary statistics
        print("\n=== ABLATION STUDY SUMMARY ===")
        for method in ablation_df['method'].unique():
            method_results = ablation_df[ablation_df['method'] == method]
            print(f"\n{method.upper()}:")
            print(f"  CAI improvement: {method_results['cai_improvement'].mean():.4f} ± {method_results['cai_improvement'].std():.4f}")
            print(f"  tAI improvement: {method_results['tai_improvement'].mean():.4f} ± {method_results['tai_improvement'].std():.4f}")
            print(f"  GC improvement: {method_results['gc_improvement'].mean():.4f} ± {method_results['gc_improvement'].std():.4f}")
            print(f"  Restriction sites improvement: {method_results['restriction_sites_improvement'].mean():.2f} ± {method_results['restriction_sites_improvement'].std():.2f}")
            print(f"  Composite score improvement: {method_results['composite_score_improvement'].mean():.4f} ± {method_results['composite_score_improvement'].std():.4f}")
    
    # Print final summary
    print("\n=== EVALUATION COMPLETE ===")
    print(f"Total proteins evaluated: {len(results_df['protein_id'].unique())}")
    print(f"Total sequences generated: {len(results_df)}")
    print(f"Results saved to: {args.output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Enhanced CodonTransformer Evaluation")
    
    # Input/Output paths
    parser.add_argument("--checkpoint_path", type=str, default="models/ecoli-codon-optimizer/finetune_best.ckpt",
                        help="Path to fine-tuned model checkpoint")
    parser.add_argument("--test_data_path", type=str, default="data/test_set.json",
                        help="Path to test dataset")
    parser.add_argument("--natural_sequences_path", type=str, default="data/ecoli_processed_genes.csv",
                        help="Path to natural E. coli sequences for CAI calculation")
    parser.add_argument("--output_path", type=str, default="results/enhanced_evaluation_results.csv",
                        help="Path to save evaluation results")
    
    # Model parameters
    parser.add_argument("--use_gpu", action="store_true", help="Use GPU if available")
    parser.add_argument("--compare_with_base", action="store_true", help="Compare with base model")
    
    # Generation parameters
    parser.add_argument("--use_constrained_search", action="store_true", 
                        help="Use constrained beam search")
    parser.add_argument("--gc_bounds", type=float, nargs=2, default=[0.50, 0.54],
                        help="GC content bounds (min max)")
    parser.add_argument("--beam_size", type=int, default=10,
                        help="Beam size for standard generation")
    parser.add_argument("--length_penalty", type=float, default=1.2,
                        help="Length penalty for beam search")
    parser.add_argument("--diversity_penalty", type=float, default=0.1,
                        help="Diversity penalty for beam search")
    
    # Enhancement parameters
    parser.add_argument("--use_enhanced_generation", action="store_true",
                        help="Use enhanced generation with DNAChisel and Pareto filtering")
    parser.add_argument("--enhanced_beam_size", type=int, default=20,
                        help="Beam size for enhanced generation")
    parser.add_argument("--use_dnachisel", action="store_true",
                        help="Use DNAChisel post-processing")
    parser.add_argument("--use_pareto_filtering", action="store_true",
                        help="Use Pareto frontier filtering")
    
    # Evaluation parameters
    parser.add_argument("--max_test_proteins", type=int, default=0,
                        help="Maximum number of proteins to test (0 for all)")
    parser.add_argument("--run_ablation_study", action="store_true",
                        help="Run ablation study comparing methods")
    
    args = parser.parse_args()
    main(args)