scripts/visualize.py

import matplotlib.pyplot as plt # type: ignore
import pandas as pd
import re

def create_design_plot(fasta_file_path):
    """
    Parses the MPNN fasta file and returns a Matplotlib figure.
    """
    data = []
    with open(fasta_file_path, 'r') as f:
        content = f.read()
    
    # Extract score (local score, not global_score) and recovery from headers
    # Example: >T=0.1, sample=1, score=0.7880, global_score=1.3897, seq_recovery=0.4295
    # We want to capture the local "score=" value, not "global_score="
    pattern = r"score=([\d\.]+)(?:,\s+global_score=[\d\.]+)?,\s+seq_recovery=([\d\.]+)"
    
    for line in content.split('\n'):
        if line.startswith(">") and "score=" in line and "seq_recovery" in line:
            # Make sure we're not matching global_score by checking the pattern order
            match = re.search(pattern, line)
            if match:
                # Verify we got the local score (should be before global_score if present)
                score_val = float(match.group(1))
                recovery_val = float(match.group(2))
                data.append({
                    "Score": score_val, 
                    "Recovery": recovery_val
                })
    
    df = pd.DataFrame(data)
    
    # Create the Plot
    fig, ax = plt.subplots(figsize=(8, 5))
    ax.scatter(df['Score'], df['Recovery'], color='black', alpha=0.6, s=80, edgecolors='white')
    
    # Highlight the Lead Candidate (Lowest Score) with a gold star
    best = df.loc[df['Score'].idxmin()]
    ax.scatter(best['Score'], best['Recovery'], marker='*', color='gold', s=400, edgecolors='black', linewidths=1.5, label="Lead Candidate", zorder=10)
    
    ax.set_title("BroteinShake: Design Evolution (N=20)", fontsize=14)
    ax.set_xlabel("ProteinMPNN Score (Lower = More Stable)", fontsize=10)
    ax.set_ylabel("Sequence Recovery (%)", fontsize=10)
    ax.legend()
    ax.grid(True, linestyle='--', alpha=0.5)
    
    return fig

def create_protein_viewer(pdb_file_path):
    """
    Generates HTML/JS for a 3D protein viewer using 3Dmol.js.
    """
    with open(pdb_file_path, 'r') as f:
        pdb_content = f.read().replace('\n', '\\n')
    
    html_content = f"""
    <div id="container-3d" style="height: 500px; width: 100%; position: relative;"></div>
    <script>
        (function() {{
            // Wait for 3Dmol to be loaded (from head)
            function initViewer() {{
                if (typeof $3Dmol === 'undefined') {{
                    setTimeout(initViewer, 100);
                    return;
                }}
                let element = document.getElementById('container-3d');
                if (!element) return;
                let config = {{ backgroundColor: 'white' }};
                let viewer = $3Dmol.createViewer(element, config);
                let pdbData = `{pdb_content}`;
                viewer.addModel(pdbData, "pdb");
                viewer.setStyle({{}}, {{cartoon: {{color: 'spectrum'}}}});
                viewer.zoomTo();
                viewer.render();
            }}
            if (document.readyState === 'loading') {{
                document.addEventListener('DOMContentLoaded', initViewer);
            }} else {{
                initViewer();
            }}
        }})();
    </script>
    """
    return html_content