visualization_multi.py

"""
Multi-Structure 3D Visualization Module
Based on working pairwise visualization code
"""

import numpy as np
from rmsd_utils import (
    parse_residue_atoms,
    get_backbone_sugar_coords_from_residue,
    get_base_coords_from_residue
)


def extract_window_pdb(pdb_path, window_indices):
    """
    Extract specific residues from a PDB file based on window indices.
    Uses the WORKING approach from the original code.
    """
    with open(pdb_path) as f:
        lines = f.readlines()
    
    # Get all residue numbers from the file
    residues = parse_residue_atoms(pdb_path)
    
    if not residues:
        return ''.join(lines)
    
    residue_numbers = [res['resnum'] for res in residues]
    
    # Map window indices to actual residue numbers
    target_resnums = set()
    for idx in window_indices:
        if idx < len(residue_numbers):
            target_resnums.add(residue_numbers[idx])
    
    if not target_resnums:
        return ''.join(lines)
    
    # Extract lines for these residues
    window_lines = []
    for line in lines:
        if len(line) < 6:
            continue
            
        record = line[0:6].strip()
        if record in ['ATOM', 'HETATM', 'HETAT']:
            try:
                resnum_str = line[22:26].strip()
                if resnum_str:
                    resnum = int(resnum_str)
                    if resnum in target_resnums:
                        window_lines.append(line)
            except (ValueError, IndexError):
                continue
        elif record in ['HEADER', 'TITLE', 'MODEL', 'ENDMDL']:
            window_lines.append(line)
    
    # Always add END record
    if window_lines and not any('END' in line for line in window_lines):
        window_lines.append('END\n')
    
    result = ''.join(window_lines)
    
    if not result or len(result) < 50:
        return ''.join(lines)
    
    return result


def transform_pdb_string(pdb_string, rotation_matrix, query_com, ref_com):
    """
    Apply rotation and translation to align query with reference.
    Uses the WORKING transformation from original code.
    
    CRITICAL: This is right multiplication (coord @ R), NOT left multiplication (R @ coord)
    
    Args:
        pdb_string: PDB format string
        rotation_matrix: 3x3 rotation matrix from RMSD calculation
        query_com: Center of mass of query structure (translate FROM)
        ref_com: Center of mass of reference structure (translate TO)
    
    Returns:
        Transformed PDB string with aligned coordinates
    """
    lines = pdb_string.split('\n')
    transformed_lines = []
    
    for line in lines:
        if len(line) < 54:
            transformed_lines.append(line)
            continue
            
        record = line[0:6].strip()
        if record in ['ATOM', 'HETATM', 'HETAT']:
            try:
                x = float(line[30:38].strip())
                y = float(line[38:46].strip())
                z = float(line[46:54].strip())
                
                # Transform: (coord - query_com) @ rotation_matrix + ref_com
                # This is the WORKING approach from original code
                coord = np.array([x, y, z])
                centered = coord - query_com  # Move query to origin
                rotated = np.dot(centered, rotation_matrix)  # RIGHT multiplication
                new_coord = rotated + ref_com  # Move to reference position
                
                # Write transformed line
                new_line = (
                    line[:30] +
                    f"{new_coord[0]:8.3f}" +
                    f"{new_coord[1]:8.3f}" +
                    f"{new_coord[2]:8.3f}" +
                    line[54:]
                )
                transformed_lines.append(new_line)
            except (ValueError, IndexError):
                transformed_lines.append(line)
        else:
            transformed_lines.append(line)
    
    return '\n'.join(transformed_lines)


def create_pairwise_visualization(ref_path, query_path, ref_window, query_window,
                                  rotation_matrix, ref_com, query_com, rmsd,
                                  ref_name="Reference", query_name="Query"):
    """
    Create interactive 3D visualization of two aligned structures (pairwise).
    Based on working original visualization with enhanced controls.
    
    Args:
        ref_path: Path to reference PDB file
        query_path: Path to query PDB file
        ref_window: List of residue indices for reference window
        query_window: List of residue indices for query window
        rotation_matrix: Rotation matrix from RMSD
        ref_com: Center of mass of reference
        query_com: Center of mass of query
        rmsd: RMSD value
        ref_name: Name of reference structure
        query_name: Name of query structure
    
    Returns:
        HTML string for py3Dmol visualization
    """
    
    # Extract windows
    ref_pdb = extract_window_pdb(ref_path, ref_window)
    query_pdb = extract_window_pdb(query_path, query_window)
    
    # Transform query to align with reference
    transformed_query_pdb = transform_pdb_string(
        query_pdb,
        rotation_matrix,
        query_com,
        ref_com
    )
    
    # Escape backticks
    ref_pdb_escaped = ref_pdb.replace('`', '\\`')
    transformed_query_pdb_escaped = transformed_query_pdb.replace('`', '\\`')
    
    # Create HTML with enhanced controls
    html = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
        <style>
            #container {{
                width: 100%;
                height: 700px;
                position: relative;
                border: 1px solid #ddd;
            }}
            .control-panel {{
                position: absolute;
                top: 10px;
                right: 10px;
                background: rgba(255, 255, 255, 0.95);
                padding: 15px;
                border-radius: 8px;
                font-family: Arial, sans-serif;
                font-size: 13px;
                z-index: 1000;
                box-shadow: 0 2px 8px rgba(0,0,0,0.1);
                max-width: 220px;
            }}
            .control-panel h4 {{
                margin: 0 0 10px 0;
                font-size: 14px;
                color: #333;
            }}
            .control-section {{
                margin-bottom: 12px;
                padding-bottom: 12px;
                border-bottom: 1px solid #eee;
            }}
            .control-section:last-child {{
                border-bottom: none;
                margin-bottom: 0;
            }}
            .control-section label {{
                display: block;
                margin: 6px 0;
                cursor: pointer;
            }}
            .control-section input[type="checkbox"] {{
                margin-right: 8px;
            }}
            .control-section select {{
                width: 100%;
                padding: 4px;
                margin-top: 5px;
                border: 1px solid #ccc;
                border-radius: 4px;
            }}
            .legend {{
                position: absolute;
                top: 10px;
                left: 10px;
                background: rgba(255, 255, 255, 0.95);
                padding: 15px;
                border-radius: 8px;
                font-family: Arial, sans-serif;
                font-size: 13px;
                z-index: 1000;
                box-shadow: 0 2px 8px rgba(0,0,0,0.1);
            }}
            .legend h4 {{
                margin: 0 0 10px 0;
                font-size: 14px;
                color: #333;
            }}
            .legend-item {{
                margin: 6px 0;
                display: flex;
                align-items: center;
            }}
            .color-box {{
                width: 24px;
                height: 16px;
                margin-right: 10px;
                border: 1px solid #333;
                border-radius: 2px;
            }}
            .rmsd-info {{
                position: absolute;
                bottom: 10px;
                left: 10px;
                background: rgba(255, 255, 255, 0.95);
                padding: 10px 15px;
                border-radius: 8px;
                font-family: Arial, sans-serif;
                font-size: 13px;
                z-index: 1000;
                box-shadow: 0 2px 8px rgba(0,0,0,0.1);
            }}
            .section-title {{
                font-weight: bold;
                color: #555;
                margin-bottom: 5px;
                font-size: 12px;
                text-transform: uppercase;
            }}

            .download-btn {{
                width: 100%;
                padding: 8px;
                background: #0066cc;
                color: white;
                border: none;
                border-radius: 4px;
                cursor: pointer;
                font-size: 13px;
                font-weight: bold;
                margin-top: 5px;
            }}
            .download-btn:hover {{
                background: #0052a3;
            }}
        </style>
    </head>
    <body>
        <div id="container"></div>
        
        <div class="legend">
            <h4>🧬 Structures</h4>
            <div class="legend-item">
                <div class="color-box" style="background: #4A90E2;"></div>
                <span>{ref_name}</span>
            </div>
            <div class="legend-item">
                <div class="color-box" style="background: #E94B3C;"></div>
                <span>{query_name} (Aligned)</span>
            </div>
        </div>
        
        <div class="rmsd-info">
            <strong>RMSD: {rmsd:.3f} Å</strong>
        </div>
        
        <div class="control-panel">
            <h4>⚙️ Display Options</h4>
            
            <div class="control-section">
                <div class="section-title">Structures</div>
                <label>
                    <input type="checkbox" id="showRef" checked onchange="updateDisplay()">
                    {ref_name}
                </label>
                <label>
                    <input type="checkbox" id="showQuery" checked onchange="updateDisplay()">
                    {query_name}
                </label>
            </div>
            
            <div class="control-section">
                <div class="section-title">Style</div>
                <select id="styleMode" onchange="updateDisplay()">
                    <option value="sticks">Sticks</option>
                    <option value="cartoon">Cartoon</option>
                    <option value="spheres">Spheres</option>
                    <option value="lines">Lines</option>
                </select>
            </div>
            
            <div class="control-section">
                <div class="section-title">Components</div>
                <label>
                    <input type="checkbox" id="showBackbone" checked onchange="updateDisplay()">
                    Backbone/Sugar
                </label>
                <label>
                    <input type="checkbox" id="showBases" checked onchange="updateDisplay()">
                    Bases
                </label>
            </div>
            
            <div class="control-section">
                <div class="section-title">Labels</div>
                <label>
                    <input type="checkbox" id="showResidueLabels" onchange="updateDisplay()">
                    Residue Labels
                </label>
                <label>
                    <input type="checkbox" id="showResidueNumbers" onchange="updateDisplay()">
                    Residue Numbers
                </label>
                <label>
                    <input type="checkbox" id="showAtomNames" onchange="updateDisplay()">
                    Atom Names
                </label>
                <select id="atomLabelMode" onchange="updateDisplay()">
                    <option value="all">All Atoms</option>
                    <option value="backbone">Backbone Only</option>
                    <option value="sidechain">Base Only</option>
                </select>
            </div>
            
            <div class="control-section">
                <div class="section-title">Background</div>
                <select id="bgColor" onchange="updateBackground()">
                    <option value="white">White</option>
                    <option value="black">Black</option>
                    <option value="gray">Gray</option>
                </select>
            </div>
            
            <div class="control-section">
                <div class="section-title">Download</div>
                <button class="download-btn" onclick="downloadImage()">📸 Save as JPEG</button>
            </div>
        </div>
        
        <script>
            let viewer = $3Dmol.createViewer("container", {{backgroundColor: 'white'}});
            
            const refPDB = `{ref_pdb_escaped}`;
            const queryPDB = `{transformed_query_pdb_escaped}`;
            
            const backboneAtoms = ['P', 'OP1', 'OP2', "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"];
            
            function updateDisplay() {{
                viewer.removeAllModels();
                viewer.removeAllLabels();
                
                const styleMode = document.getElementById('styleMode').value;
                const showBackbone = document.getElementById('showBackbone').checked;
                const showBases = document.getElementById('showBases').checked;
                const showResidueLabels = document.getElementById('showResidueLabels').checked;
                const showResidueNumbers = document.getElementById('showResidueNumbers').checked;
                const showAtomNames = document.getElementById('showAtomNames').checked;
                
                let modelIndex = 0;
                
                // Add reference if checked
                if (document.getElementById('showRef').checked) {{
                    viewer.addModel(refPDB, "pdb");
                    applyStyle(modelIndex, '#4A90E2', styleMode, showBackbone, showBases);
                    
                    if (showResidueLabels || showResidueNumbers) {{
                        addResidueLabels(modelIndex, '#4A90E2', showResidueLabels, showResidueNumbers);
                    }}
                    if (showAtomNames) {{
                        addAtomLabels(modelIndex, '#4A90E2');
                    }}
                    
                    modelIndex++;
                }}
                
                // Add query if checked
                if (document.getElementById('showQuery').checked) {{
                    viewer.addModel(queryPDB, "pdb");
                    applyStyle(modelIndex, '#E94B3C', styleMode, showBackbone, showBases);
                    
                    if (showResidueLabels || showResidueNumbers) {{
                        addResidueLabels(modelIndex, '#E94B3C', showResidueLabels, showResidueNumbers);
                    }}
                    if (showAtomNames) {{
                        addAtomLabels(modelIndex, '#E94B3C');
                    }}
                    
                    modelIndex++;
                }}
                
                viewer.zoomTo();
                viewer.render();
            }}
            
            function applyStyle(modelIndex, baseColor, styleMode, showBackbone, showBases) {{
                const atoms = viewer.selectedAtoms({{model: modelIndex}});
                
                atoms.forEach(atom => {{
                    const isBackbone = backboneAtoms.includes(atom.atom);
                    
                    // Skip if filtering out this type
                    if (!showBackbone && isBackbone) return;
                    if (!showBases && !isBackbone) return;
                    
                    if (styleMode === 'sticks') {{
                        viewer.setStyle({{model: modelIndex, serial: atom.serial}}, {{
                            stick: {{
                                color: baseColor,
                                radius: 0.15
                            }}
                        }});
                    }} else if (styleMode === 'cartoon') {{
                        viewer.setStyle({{model: modelIndex}}, {{
                            cartoon: {{
                                color: baseColor,
                                opacity: 0.8
                            }}
                        }});
                    }} else if (styleMode === 'spheres') {{
                        viewer.setStyle({{model: modelIndex, serial: atom.serial}}, {{
                            sphere: {{
                                color: baseColor,
                                radius: 0.25
                            }}
                        }});
                    }} else if (styleMode === 'lines') {{
                        viewer.setStyle({{model: modelIndex, serial: atom.serial}}, {{
                            line: {{
                                color: baseColor
                            }}
                        }});
                    }}
                }});
            }}
            
            function addResidueLabels(modelIndex, color, showLabels, showNumbers) {{
                const atoms = viewer.selectedAtoms({{model: modelIndex}});
                const residues = {{}};
                
                atoms.forEach(atom => {{
                    const key = atom.chain + '_' + atom.resi;
                    if (!residues[key]) {{
                        residues[key] = atom;
                    }}
                }});
                
                Object.values(residues).forEach(atom => {{
                    let labelText = '';
                    if (showLabels && showNumbers) {{
                        labelText = atom.resn + atom.resi;
                    }} else if (showLabels) {{
                        labelText = atom.resn;
                    }} else if (showNumbers) {{
                        labelText = atom.resi.toString();
                    }}
                    
                    if (labelText) {{
                        viewer.addLabel(labelText, {{
                            position: atom,
                            backgroundColor: color,
                            backgroundOpacity: 0.7,
                            fontColor: 'white',
                            fontSize: 11,
                            fontWeight: 'bold',
                            showBackground: true,
                            borderRadius: 3
                        }});
                    }}
                }});
            }}
            
            function addAtomLabels(modelIndex, color) {{
                const atomLabelMode = document.getElementById('atomLabelMode').value;
                const atoms = viewer.selectedAtoms({{model: modelIndex}});
                
                let filteredAtoms = atoms;
                if (atomLabelMode === 'backbone') {{
                    filteredAtoms = atoms.filter(atom => backboneAtoms.includes(atom.atom));
                }} else if (atomLabelMode === 'sidechain') {{
                    filteredAtoms = atoms.filter(atom => !backboneAtoms.includes(atom.atom));
                }}
                
                filteredAtoms.forEach(atom => {{
                    viewer.addLabel(atom.atom, {{
                        position: atom,
                        backgroundColor: color,
                        backgroundOpacity: 0.6,
                        fontColor: 'white',
                        fontSize: 9,
                        fontWeight: 'normal',
                        showBackground: true,
                        borderRadius: 2
                    }});
                }});
            }}
            
            function updateBackground() {{
                const bgColor = document.getElementById('bgColor').value;
                viewer.setBackgroundColor(bgColor);
                viewer.render();
            }}
            
            // Initialize
            updateDisplay();
            
            // Download functionality
            window.viewer = viewer;
            
            function downloadImage() {{
                try {{
                    const imgData = viewer.pngURI();
                    const img = new Image();
                    img.onload = function() {{
                        const canvas = document.createElement('canvas');
                        canvas.width = img.width;
                        canvas.height = img.height;
                        const ctx = canvas.getContext('2d');
                        ctx.fillStyle = 'white';
                        ctx.fillRect(0, 0, canvas.width, canvas.height);
                        ctx.drawImage(img, 0, 0);
                        canvas.toBlob(function(blob) {{
                            const url = URL.createObjectURL(blob);
                            const a = document.createElement('a');
                            a.href = url;
                            a.download = 'rna_pairwise_comparison.jpg';
                            document.body.appendChild(a);
                            a.click();
                            document.body.removeChild(a);
                            URL.revokeObjectURL(url);
                        }}, 'image/jpeg', 0.95);
                    }};
                    img.src = imgData;
                }} catch(e) {{
                    alert('Error: ' + e.message);
                }}
            }}
        </script>
    </body>
    </html>
    """
    
    return html


def create_multistructure_visualization(ref_path, ref_window, ref_com, query_data_list, ref_name="Reference"):
    """
    Create interactive 3D visualization with multiple structures aligned to reference.
    
    Args:
        ref_path: Path to reference PDB file
        ref_window: List of residue indices for reference window
        ref_com: Center of mass of reference window
        query_data_list: List of dicts with keys:
            - name: Structure name
            - path: Path to PDB file
            - window: List of residue indices
            - rotation: Rotation matrix from RMSD
            - query_com: Center of mass
            - rmsd: RMSD value
            - sequence: Sequence string
        ref_name: Name of reference structure
    
    Returns:
        HTML string for py3Dmol visualization
    """
    
    # Extract reference window
    ref_pdb = extract_window_pdb(ref_path, ref_window)
    
    # Color palette
    colors = [
        '#E94B3C',  # Red
        '#50C878',  # Green
        '#FF9500',  # Orange
        '#9B59B6',  # Purple
        '#00CED1',  # Cyan
        '#FF1493',  # Pink
        '#FFD700',  # Gold
        '#8B4513',  # Brown
        '#708090'   # Slate Gray
    ]
    
    # Transform all query structures
    transformed_queries = []
    for idx, query_data in enumerate(query_data_list):
        # Extract query window
        query_pdb = extract_window_pdb(query_data['path'], query_data['window'])
        
        # Transform query to align with reference
        transformed_pdb = transform_pdb_string(
            query_pdb,
            query_data['rotation'],
            query_data['query_com'],
            ref_com
        )
        
        color = colors[idx % len(colors)]
        transformed_queries.append({
            'pdb': transformed_pdb,
            'name': query_data['name'],
            'color': color,
            'rmsd': query_data['rmsd'],
            'sequence': query_data['sequence']
        })
    
    # Build JavaScript for model data
    models_js = f"const refPDB = `{ref_pdb}`;\n"
    for idx, tq in enumerate(transformed_queries):
        # Escape backticks
        pdb_escaped = tq['pdb'].replace('`', '\\`')
        models_js += f"const queryPDB{idx} = `{pdb_escaped}`;\n"
    
    # Build controls HTML
    controls_html = f'<label><input type="checkbox" id="showRef" checked onchange="updateDisplay()"> {ref_name} (Blue)</label>\n'
    for idx, tq in enumerate(transformed_queries):
        controls_html += f'<label><input type="checkbox" id="showQuery{idx}" checked onchange="updateDisplay()"> {tq["name"]} ({tq["color"]})</label>\n'
    
    # Build legend HTML
    legend_html = f'<div class="legend-item"><div class="color-box" style="background: #4A90E2;"></div><span>{ref_name}</span></div>\n'
    for tq in transformed_queries:
        legend_html += f'<div class="legend-item"><div class="color-box" style="background: {tq["color"]};"></div><span>{tq["name"]} ({tq["rmsd"]:.3f} Å)</span></div>\n'
    
    # Build update function
    update_js = '''
        function updateDisplay() {
            viewer.removeAllModels();
            const styleMode = document.getElementById('styleMode').value;
            const showLabels = document.getElementById('showLabels').checked;
            
            let modelIndex = 0;
            
            // Add reference if checked
            if (document.getElementById('showRef').checked) {
                viewer.addModel(refPDB, "pdb");
                applyStyle(modelIndex, '#4A90E2', styleMode);
                if (showLabels) {
                    addLabels(modelIndex, '#4A90E2');
                }
                modelIndex++;
            }
    '''
    
    # Add query structures
    for idx in range(len(transformed_queries)):
        color = transformed_queries[idx]['color']
        update_js += f'''
            if (document.getElementById('showQuery{idx}').checked) {{
                viewer.addModel(queryPDB{idx}, "pdb");
                applyStyle(modelIndex, '{color}', styleMode);
                if (showLabels) {{
                    addLabels(modelIndex, '{color}');
                }}
                modelIndex++;
            }}
        '''
    
    update_js += '''
            viewer.zoomTo();
            viewer.render();
        }
        
        function applyStyle(modelIndex, color, styleMode) {
            if (styleMode === 'sticks') {
                viewer.setStyle({model: modelIndex}, {stick: {color: color, radius: 0.15}});
            } else if (styleMode === 'cartoon') {
                viewer.setStyle({model: modelIndex}, {cartoon: {color: color, opacity: 0.8}});
            } else if (styleMode === 'spheres') {
                viewer.setStyle({model: modelIndex}, {sphere: {color: color, radius: 0.3}});
            } else if (styleMode === 'lines') {
                viewer.setStyle({model: modelIndex}, {line: {color: color}});
            }
        }
        
        function addLabels(modelIndex, color) {
            const atoms = viewer.selectedAtoms({model: modelIndex});
            const residues = {};
            
            atoms.forEach(atom => {
                const key = atom.chain + '_' + atom.resi;
                if (!residues[key]) {
                    residues[key] = atom;
                }
            });
            
            Object.values(residues).forEach(atom => {
                viewer.addLabel(atom.resn + atom.resi, {
                    position: atom,
                    backgroundColor: color,
                    backgroundOpacity: 0.7,
                    fontColor: 'white',
                    fontSize: 10,
                    fontWeight: 'bold',
                    showBackground: true
                });
            });
        }
    '''
    
    # Create HTML
    html = f'''
    <!DOCTYPE html>
    <html>
    <head>
        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
        <style>
            #container {{
                width: 100%;
                height: 700px;
                position: relative;
                border: 1px solid #ddd;
            }}
            .control-panel {{
                position: absolute;
                top: 10px;
                right: 10px;
                background: rgba(255, 255, 255, 0.95);
                padding: 15px;
                border-radius: 8px;
                font-family: Arial, sans-serif;
                font-size: 13px;
                z-index: 1000;
                box-shadow: 0 2px 8px rgba(0,0,0,0.1);
                max-width: 250px;
                max-height: 600px;
                overflow-y: auto;
            }}
            .control-panel h4 {{
                margin: 0 0 10px 0;
                font-size: 14px;
                color: #333;
            }}
            .control-section {{
                margin-bottom: 12px;
                padding-bottom: 12px;
                border-bottom: 1px solid #eee;
            }}
            .control-section:last-child {{
                border-bottom: none;
                margin-bottom: 0;
            }}
            .control-section label {{
                display: block;
                margin: 6px 0;
                cursor: pointer;
            }}
            .control-section input[type="checkbox"] {{
                margin-right: 8px;
            }}
            .control-section select {{
                width: 100%;
                padding: 4px;
                margin-top: 5px;
                border: 1px solid #ccc;
                border-radius: 4px;
            }}
            .legend {{
                position: absolute;
                top: 10px;
                left: 10px;
                background: rgba(255, 255, 255, 0.95);
                padding: 15px;
                border-radius: 8px;
                font-family: Arial, sans-serif;
                font-size: 13px;
                z-index: 1000;
                box-shadow: 0 2px 8px rgba(0,0,0,0.1);
            }}
            .legend h4 {{
                margin: 0 0 10px 0;
                font-size: 14px;
                color: #333;
            }}
            .legend-item {{
                margin: 6px 0;
                display: flex;
                align-items: center;
            }}
            .color-box {{
                width: 24px;
                height: 16px;
                margin-right: 10px;
                border: 1px solid #333;
                border-radius: 2px;
            }}
            .section-title {{
                font-weight: bold;
                color: #555;
                margin-bottom: 5px;
                font-size: 12px;
                text-transform: uppercase;
            }}

            .download-btn {{
                width: 100%;
                padding: 8px;
                background: #0066cc;
                color: white;
                border: none;
                border-radius: 4px;
                cursor: pointer;
                font-size: 13px;
                font-weight: bold;
                margin-top: 5px;
            }}
            .download-btn:hover {{
                background: #0052a3;
            }}
        </style>
    </head>
    <body>
        <div id="container"></div>
        
        <div class="legend">
            <h4>🧬 Structures</h4>
            {legend_html}
        </div>
        
        <div class="control-panel">
            <h4>⚙️ Display Options</h4>
            
            <div class="control-section">
                <div class="section-title">Structures</div>
                {controls_html}
            </div>
            
            <div class="control-section">
                <div class="section-title">Style</div>
                <select id="styleMode" onchange="updateDisplay()">
                    <option value="sticks">Sticks</option>
                    <option value="cartoon">Cartoon</option>
                    <option value="spheres">Spheres</option>
                    <option value="lines">Lines</option>
                </select>
            </div>
            
            <div class="control-section">
                <div class="section-title">Labels</div>
                <label>
                    <input type="checkbox" id="showLabels" onchange="updateDisplay()">
                    Show Residue Labels
                </label>
            </div>
            
            <div class="control-section">
                <div class="section-title">Download</div>
                <button class="download-btn" onclick="downloadImage()">📸 Save as JPEG</button>
            </div>
        </div>
        
        <script>
            let viewer = $3Dmol.createViewer("container", {{backgroundColor: 'white'}});
            
            {models_js}
            
            {update_js}
            
            // Initialize
            updateDisplay();
            
            // Download functionality
            window.viewer = viewer;  // Make viewer globally accessible
            
            function downloadImage() {{
                try {{
                    // Use 3Dmol.js built-in PNG export
                    const imgData = viewer.pngURI();
                    
                    // Convert PNG to JPEG with white background
                    const img = new Image();
                    img.onload = function() {{
                        const canvas = document.createElement('canvas');
                        canvas.width = img.width;
                        canvas.height = img.height;
                        const ctx = canvas.getContext('2d');
                        
                        // White background for JPEG
                        ctx.fillStyle = 'white';
                        ctx.fillRect(0, 0, canvas.width, canvas.height);
                        ctx.drawImage(img, 0, 0);
                        
                        // Convert to JPEG and download
                        canvas.toBlob(function(blob) {{
                            const url = URL.createObjectURL(blob);
                            const a = document.createElement('a');
                            a.href = url;
                            a.download = 'rna_multi_structure_comparison.jpg';
                            document.body.appendChild(a);
                            a.click();
                            document.body.removeChild(a);
                            URL.revokeObjectURL(url);
                        }}, 'image/jpeg', 0.95);
                    }};
                    img.src = imgData;
                }} catch(e) {{
                    alert('Error capturing image: ' + e.message);
                    console.error(e);
                }}
            }}
        </script>
    </body>
    </html>
    '''
    
    return html