import numpy as np
import time
import pandas as pd

# Import the core functions from your other project files
from nerf import place_atom_nerf
from ramachandran import RamachandranSampler
from backbone import generate_backbone_chain

# =============================================================================
# BENCHMARKING FUNCTIONS
# =============================================================================

AMINO_ACIDS = [
    'ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HIS', 'ILE',
    'LEU', 'LYS', 'MET', 'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL'
]

def generate_random_sequence(length):
    """Generates a random list of 3-letter amino acid codes."""
    return list(np.random.choice(AMINO_ACIDS, size=length))

def run_benchmark():
    """
    Runs the main benchmark and prints a summary report.
    """
    
    # --- Parameters to Test ---
    # (Feel free to change these)
    lengths_to_test = [50, 100, 200, 300, 400, 1000, 10000]
    trials_per_setting = 3
    clash_thresholds_to_test = [3.5, 3.0]
    backtrack_steps = 5
    
    # --- Setup ---
    sampler = RamachandranSampler()
    nerf_func = place_atom_nerf
    results = []
    
    print("Starting Polypeptide Generator Benchmark...")
    print(f"Testing Lengths: {lengths_to_test}")
    print(f"Testing Thresholds: {clash_thresholds_to_test} Å")
    print(f"Trials per Setting: {trials_per_setting}\n")
    
    # --- Main Loop ---
    for threshold in clash_thresholds_to_test:
        for length in lengths_to_test:
            for trial in range(trials_per_setting):
                print(f"Running: Threshold={threshold}Å, Length={length}, Trial={trial+1}/{trials_per_setting}...")
                
                # Generate a new random sequence for this trial
                seq = generate_random_sequence(length)
                
                t_start = time.perf_counter()
                
                # Call the imported function
                chain, backtracks = generate_backbone_chain(
                    seq,
                    nerf_func,
                    sampler,
                    clash_threshold=threshold,
                    backtrack_step=backtrack_steps,
                    verbose=True # Set to True to see backtrack messages
                )
                
                t_end = time.perf_counter()
                
                time_taken = t_end - t_start
                
                # Store results
                results.append({
                    "Threshold (Å)": threshold,
                    "Length (Res)": length,
                    "Trial": trial,
                    "Time (s)": time_taken,
                    "Backtracks": backtracks
                })

    # --- Analysis & Report ---
    print("\n--- Benchmark Complete ---")
    
    if not results:
        print("No results to report.")
        return

    # Convert to DataFrame for easy analysis
    df = pd.DataFrame(results)
    
    # Set display options for pandas
    pd.set_option('display.precision', 3)
    pd.set_option('display.width', 100)
    
    print("\nFull Results:")
    print(df)
    
    # Group by threshold and length to see averages
    summary = df.groupby(["Threshold (Å)", "Length (Res)"]).agg(
        avg_time=('Time (s)', 'mean'),
        std_time=('Time (s)', 'std'),
        avg_backtracks=('Backtracks', 'mean'),
        std_backtracks=('Backtracks', 'std')
    )
    
    print("\n--- Summary Report (Averages) ---")
    print(summary)

if __name__ == "__main__":
    run_benchmark()
