"""
RNA Motif Multi-Structure Comparison Tool - Pairwise Mode
Streamlit app for comparing multiple RNA motif structures with separate reference and query sets
Uses dropdown menu for residue configuration and default Backbone + Sugar atom selection
"""

import streamlit as st
import numpy as np
import pandas as pd
from pathlib import Path
import io
import tempfile
import os
from itertools import combinations

# Import our RMSD calculation functions
from rmsd_utils import (
    parse_residue_atoms,
    get_backbone_sugar_and_selectbase_coords_fixed,
    calculate_COM,
    calculate_rotation_rmsd,
    translate_rotate_coords,
    get_backbone_sugar_coords_from_residue,
    get_base_coords_from_residue
)

# Import example data loader
try:
    from example_data_loader import (
        get_example_pdbs,
        load_example_as_uploaded_file,
        get_example_info
    )
    EXAMPLES_AVAILABLE = True
except ImportError:
    EXAMPLES_AVAILABLE = False
    st.warning("Example data loader not available. Please use 'Upload Files' mode.")

# Page configuration
st.set_page_config(
    page_title="RNA Motif Multi-Structure Comparison - Pairwise",
    page_icon="🧬",
    layout="wide",
    initial_sidebar_state="expanded"
)

from image_annotator import annotate_alignment_image


# Custom CSS - IMPROVED VERSION with larger fonts
st.markdown("""
<style>
    /* ========================================
       MAIN CONTENT - LARGER FONTS
       ======================================== */
    
    /* Increase base font size for all main content */
    .main .element-container,
    .main [data-testid="stMarkdownContainer"],
    .main [data-testid="stText"],
    .main p,
    .main span,
    .main div {
        font-size: 1.15rem !important;
    }
    
    /* Headers in main content */
    .main h1 {
        font-size: 2.8rem !important;
        font-weight: 700 !important;
    }
    .main h2 {
        font-size: 2.0rem !important;
        font-weight: 600 !important;
    }
    .main h3 {
        font-size: 1.6rem !important;
        font-weight: 600 !important;
    }
    
    /* Custom header classes */
    .main-header {
        font-size: 2.8rem !important;
        font-weight: bold;
        color: #1f77b4;
        margin-bottom: 1rem;
    }
    .sub-header {
        font-size: 1.4rem !important;
        color: #666;
        margin-bottom: 2rem;
    }
    
    /* Info/warning/success boxes */
    .main [data-testid="stAlert"] p,
    .main [data-testid="stAlert"] {
        font-size: 1.1rem !important;
    }
    
    /* Dataframes and tables */
    .main [data-testid="stDataFrame"],
    .main .dataframe,
    .main table {
        font-size: 1.05rem !important;
    }
    .main .dataframe th,
    .main .dataframe td {
        font-size: 1.05rem !important;
        padding: 8px !important;
    }
    
    /* Metrics */
    .main [data-testid="stMetric"] {
        font-size: 1.15rem !important;
    }
    .main [data-testid="stMetricLabel"] {
        font-size: 1.1rem !important;
    }
    .main [data-testid="stMetricValue"] {
        font-size: 1.8rem !important;
    }
    
    /* Buttons in main content */
    .main button p,
    .main button span {
        font-size: 1.05rem !important;
    }
    
    /* Selectbox, radio, and other inputs in main */
    .main .stSelectbox label,
    .main .stRadio label,
    .main .stNumberInput label,
    .main .stMultiSelect label {
        font-size: 1.1rem !important;
    }
    
    .main .stSelectbox [data-baseweb="select"] div,
    .main .stRadio [role="radiogroup"] label,
    .main .stNumberInput input {
        font-size: 1.05rem !important;
    }
    
    /* Expander headers */
    .main [data-testid="stExpander"] summary {
        font-size: 1.15rem !important;
    }
    
    /* Code blocks */
    .main code,
    .main pre {
        font-size: 1.0rem !important;
    }
    
    /* ========================================
       SIDEBAR - COMPACT & NORMAL FONT
       ======================================== */
    
    /* Ultra-compact sidebar spacing */
    section[data-testid="stSidebar"] {
        padding-top: 0.2rem !important;
    }
    section[data-testid="stSidebar"] > div {
        padding-top: 0.2rem !important;
    }
    
    /* Minimal margins */
    section[data-testid="stSidebar"] [data-testid="stMarkdownContainer"] {
        margin: 0rem !important;
    }
    
    /* Minimal header spacing */
    section[data-testid="stSidebar"] h1, 
    section[data-testid="stSidebar"] h2, 
    section[data-testid="stSidebar"] h3 {
        margin-top: 0.1rem !important;
        margin-bottom: 0.2rem !important;
        padding: 0rem !important;
        line-height: 1.2 !important;
        font-size: 1.0rem !important;
    }
    
    /* Tight widget spacing */
    section[data-testid="stSidebar"] .stSelectbox,
    section[data-testid="stSidebar"] .stNumberInput,
    section[data-testid="stSidebar"] .stRadio,
    section[data-testid="stSidebar"] .stFileUploader {
        margin-top: 0.1rem !important;
        margin-bottom: 0.2rem !important;
    }
    
    section[data-testid="stSidebar"] .stButton {
        margin: 0.2rem 0 !important;
    }
    
    section[data-testid="stSidebar"] .element-container {
        margin: 0.1rem 0 !important;
    }
    
    section[data-testid="stSidebar"] .stAlert {
        padding: 0.3rem 0.5rem !important;
        margin: 0.1rem 0 !important;
    }
    
    section[data-testid="stSidebar"] label {
        margin-bottom: 0.1rem !important;
        font-size: 0.9rem !important;
    }
    
    section[data-testid="stSidebar"] .stCaptionContainer {
        margin: 0.1rem 0 !important;
    }
    
    section[data-testid="stSidebar"] hr {
        margin: 0.2rem 0 !important;
    }
    
    /* Sidebar font sizes - keep normal/small */
    section[data-testid="stSidebar"] * {
        font-size: 0.9rem !important;
    }
    
    section[data-testid="stSidebar"] p,
    section[data-testid="stSidebar"] span,
    section[data-testid="stSidebar"] div {
        font-size: 0.9rem !important;
    }
    
    section[data-testid="stSidebar"] button {
        font-size: 0.9rem !important;
    }
</style>
""", unsafe_allow_html=True)


def save_uploaded_file(uploaded_file, directory):
    """Save an uploaded file to a temporary directory"""
    file_path = os.path.join(directory, uploaded_file.name)
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    return file_path


def get_structure_info(pdb_path):
    """
    Get information about a structure's residues.
    
    Args:
        pdb_path: Path to PDB file
    
    Returns:
        List of dicts with residue info: [{index, resnum, resname, full_name}, ...]
    """
    residues = parse_residue_atoms(pdb_path)
    
    structure_info = []
    for idx, res in enumerate(residues):
        structure_info.append({
            'index': idx,
            'resnum': res['resnum'],
            'resname': res['resname'],
            'full_name': f"{idx+1}. {res['resname']} (residue #{res['resnum']})"
        })
    
    return structure_info


def load_structure_data(uploaded_files, temp_dir):
    """Load structure data from uploaded files"""
    structure_data = []
    
    for uploaded_file in uploaded_files:
        file_path = save_uploaded_file(uploaded_file, temp_dir)
        residues = parse_residue_atoms(file_path)
        
        structure_data.append({
            'name': uploaded_file.name,
            'path': file_path,
            'residues': residues,
            'num_residues': len(residues)
        })
    
    return structure_data



def extract_window_coords(residues, window_indices):
    """
    Extract coordinates for a specific window of residues.
    
    Args:
        residues: List of all residues
        window_indices: List of indices to extract
    
    Returns:
        numpy array of coordinates
    """
    from rmsd_utils import get_backbone_sugar_coords_from_residue, get_base_coords_from_residue
    
    all_coords = []
    for idx in window_indices:
        if idx < len(residues):
            residue = residues[idx]
            # Get backbone and sugar coordinates
            backbone_coords = get_backbone_sugar_coords_from_residue(residue)
            all_coords.extend(backbone_coords)
            # Get base coordinates
            base_coords = get_base_coords_from_residue(residue)
            all_coords.extend(base_coords)
    
    return np.asarray(all_coords)


def generate_windows_from_selection(selected_indices, window_size, window_type):
    """Generate windows from selected residue indices"""
    if len(selected_indices) < window_size:
        return []
    
    windows = []

    
    if len(selected_indices) == window_size:
        windows.append(selected_indices)
        return windows
        
    if window_type == "contiguous":
        # Only sliding windows
        for i in range(len(selected_indices) - window_size + 1):
            windows.append(selected_indices[i:i+window_size])
    
    elif window_type == "non-contiguous":
        from itertools import combinations
        all_combos = list(combinations(selected_indices, window_size))
        
        # Get the contiguous windows (to exclude them)
        contiguous_windows = []
        for i in range(len(selected_indices) - window_size + 1):
            contiguous_windows.append(tuple(selected_indices[i:i+window_size]))
        
        # Filter: keep only combinations that are NOT in contiguous_windows
        for combo in all_combos:
            if combo not in contiguous_windows:
                windows.append(list(combo))
    else:
        from itertools import combinations
        all_combos = list(combinations(selected_indices, window_size))
       
        # Filter: keep only combinations that are NOT in contiguous_windows
        for combo in all_combos:
            windows.append(list(combo))
    return windows

def main():
    st.markdown('<h1 class="main-header">🧬 RNA Motif Multi-Structure Comparison</h1>', unsafe_allow_html=True)
    st.markdown('<p class="sub-header">Pairwise comparison: Reference structures vs Query structures</p>', unsafe_allow_html=True)
    
    # Create temporary directory
    if 'temp_dir' not in st.session_state:
        st.session_state['temp_dir'] = tempfile.mkdtemp()
    temp_dir = st.session_state['temp_dir']
    
    # Initialize session state
    if 'data_mode' not in st.session_state:
        st.session_state['data_mode'] = 'upload'
    if 'ref_selections' not in st.session_state:
        st.session_state['ref_selections'] = {}
    if 'query_selections' not in st.session_state:
        st.session_state['query_selections'] = {}
    
    # Sidebar: Step 1 - Data Source Selection
    st.sidebar.title("⚙️ Configuration")
    st.sidebar.subheader("1️⃣ Data Source")
    
    # Check if examples are available
    if EXAMPLES_AVAILABLE:
        data_mode = st.sidebar.radio(
            "Choose data source",
            ["Upload Files", "Use Example Data"],
            key="data_mode_radio",
            help="Upload your own PDB files or use provided examples"
        )
    else:
        st.sidebar.info("ℹ️ Example data not available. Using upload mode.")
        data_mode = "Upload Files"
    
    # Update data mode
    if data_mode == "Upload Files":
        st.session_state['data_mode'] = 'upload'
        # Reset example initialization when switching to upload mode
        if 'example_mode_initialized' in st.session_state:
            del st.session_state['example_mode_initialized']
    else:
        st.session_state['data_mode'] = 'example'
    
    # Step 2: File Upload/Selection - SEPARATE FOR REFERENCE AND QUERY
    st.sidebar.subheader("2️⃣ Structure Files")
    
    reference_files = []
    query_files = []
    
    if st.session_state['data_mode'] == 'upload':
        st.sidebar.markdown("**Upload Reference Structures**")
        ref_uploaded = st.sidebar.file_uploader(
            "Reference PDB files",
            type=['pdb'],
            accept_multiple_files=True,
            key="ref_uploader",
            help="Upload one or more reference structures (e.g., Pentaloop)"
        )
        
        st.sidebar.markdown("**Upload Query Structures**")
        query_uploaded = st.sidebar.file_uploader(
            "Query PDB files",
            type=['pdb'],
            accept_multiple_files=True,
            key="query_uploader",
            help="Upload one or more query structures (e.g., Tetraloop)"
        )
        
        reference_files = ref_uploaded if ref_uploaded else []
        query_files = query_uploaded if query_uploaded else []
        
    else:  # Example data mode
        if not EXAMPLES_AVAILABLE:
            st.sidebar.error("❌ Example data loader module not found")
            reference_files = []
            query_files = []
        else:
            try:
                examples = get_example_pdbs()
                
                if not examples or len(examples) == 0:
                    st.sidebar.error("❌ No example data available. Please add PDB files to 'data/' folder")
                    st.sidebar.info("💡 Create a 'data/' folder in the same directory as the app and add .pdb files")
                    reference_files = []
                    query_files = []
                else:
                    example_names = sorted(list(examples.keys()))
                    
                    # Auto-select examples when first switching to example mode
                    if 'example_mode_initialized' not in st.session_state:
                        st.session_state['example_mode_initialized'] = True
                        # Auto-select first half as reference, second half as query
                        mid_point = max(1, len(example_names) // 2)
                        st.session_state['auto_ref_examples'] = example_names[:mid_point]
                        st.session_state['auto_query_examples'] = example_names[mid_point:mid_point*2]
                    
                    st.sidebar.markdown("**Select Reference Examples**")
                    ref_example_names = st.sidebar.multiselect(
                        "Reference structures",
                        options=example_names,
                        default=st.session_state.get('auto_ref_examples', []),
                        key="ref_examples",
                        help="Select example reference structures"
                    )
                    
                    if ref_example_names:
                        st.sidebar.success(f"✅ {len(ref_example_names)} reference file(s) selected")
                    
                    st.sidebar.markdown("**Select Query Examples**")
                    query_example_names = st.sidebar.multiselect(
                        "Query structures",
                        options=example_names,
                        default=st.session_state.get('auto_query_examples', []),
                        key="query_examples",
                        help="Select example query structures"
                    )
                    
                    if query_example_names:
                        st.sidebar.success(f"✅ {len(query_example_names)} query file(s) selected")
                    
                    # Convert names to paths and load files
                    try:
                        reference_files = [load_example_as_uploaded_file(examples[name]) for name in ref_example_names]
                        query_files = [load_example_as_uploaded_file(examples[name]) for name in query_example_names]
                            
                    except Exception as load_error:
                        st.sidebar.error(f"Error loading files: {str(load_error)}")
                        import traceback
                        st.sidebar.code(traceback.format_exc())
                        reference_files = []
                        query_files = []
            except Exception as e:
                st.sidebar.error(f"❌ Error loading examples: {str(e)}")
                import traceback
                st.sidebar.code(traceback.format_exc())
                reference_files = []
                query_files = []
    
    # Show upload status
    if reference_files and query_files:
        st.sidebar.success(f"✅ {len(reference_files)} reference + {len(query_files)} query structures")
    elif reference_files:
        st.sidebar.info(f"ℹ️ {len(reference_files)} reference structures loaded")
    elif query_files:
        st.sidebar.info(f"ℹ️ {len(query_files)} query structures loaded")
    else:
        st.sidebar.warning("⚠️ Upload or select structures")
    
    # Residue trimming controls - add early so they're available when needed
    st.sidebar.markdown("---")
    st.sidebar.markdown("**🔧 5'/3' Base Trimming (Reference) **")
    col1, col2 = st.sidebar.columns(2)
    with col1:
        n_term_trim_ref = st.number_input(
            "5' trim_ref",
            min_value=0,
            max_value=10,
            value=2,
            step=1,
            help="Number of bases to remove from 5' end",
            key="n_term_trim_ref"
        )
    with col2:
        c_term_trim_ref = st.number_input(
            "3' trim_ref",
            min_value=0,
            max_value=10,
            value=2,
            step=1,
            help="Number of bases to remove from 3' end",
            key="c_term_trim_ref"
        )
    
    
    # Residue trimming controls - add early so they're available when needed
    st.sidebar.markdown("---")
    st.sidebar.markdown("**🔧 5'/3' Base Trimming (Query) **")
    col1, col2 = st.sidebar.columns(2)
    with col1:
        n_term_trim_query = st.number_input(
            "5' trim_query",
            min_value=0,
            max_value=10,
            value=2,
            step=1,
            help="Number of bases to remove from 5' end",
            key="n_term_trim_query"
        )
    with col2:
        c_term_trim_query = st.number_input(
            "3' trim_query",
            min_value=0,
            max_value=10,
            value=2,
            step=1,
            help="Number of bases to remove from 3' end",
            key="c_term_trim_query"
        )
    
    # Load structure data
    ref_structure_data = []
    query_structure_data = []
    
    if reference_files:
        ref_structure_data = load_structure_data(reference_files, temp_dir)
    
    if query_files:
        query_structure_data = load_structure_data(query_files, temp_dir)
    
    # Track current files to reset selections if files change
    current_ref_files = set([s['name'] for s in ref_structure_data])
    current_query_files = set([s['name'] for s in query_structure_data])
    
    if 'current_ref_files' not in st.session_state:
        st.session_state['current_ref_files'] = current_ref_files
    if 'current_query_files' not in st.session_state:
        st.session_state['current_query_files'] = current_query_files
    
    # Reset selections if files changed
    if st.session_state['current_ref_files'] != current_ref_files:
        st.session_state['current_ref_files'] = current_ref_files
        st.session_state['ref_selections'] = {}
        if 'ref_auto_initialized' in st.session_state:
            del st.session_state['ref_auto_initialized']
    
    if st.session_state['current_query_files'] != current_query_files:
        st.session_state['current_query_files'] = current_query_files
        st.session_state['query_selections'] = {}
        if 'query_auto_initialized' in st.session_state:
            del st.session_state['query_auto_initialized']
    
    # Auto-initialize selections (exclude first and last residue by default)
    if 'ref_auto_initialized' not in st.session_state and ref_structure_data:
        for struct in ref_structure_data:
            num_res = struct['num_residues']
            if num_res > n_term_trim_ref + c_term_trim_ref:
                auto_selection = list(range(n_term_trim_ref, num_res - c_term_trim_ref))
                st.session_state['ref_selections'][struct['name']] = auto_selection
            else:
                st.session_state['ref_selections'][struct['name']] = list(range(num_res))
        st.session_state['ref_auto_initialized'] = True
    
    if 'query_auto_initialized' not in st.session_state and query_structure_data:
        for struct in query_structure_data:
            num_res = struct['num_residues']
            if num_res > n_term_trim_query + c_term_trim_query:
                auto_selection = list(range(n_term_trim_query, num_res - c_term_trim_query))
                st.session_state['query_selections'][struct['name']] = auto_selection
            else:
                st.session_state['query_selections'][struct['name']] = list(range(num_res))
        st.session_state['query_auto_initialized'] = True
    
    # Step 3: Configure Atom Selections in Main Area
    st.markdown("---")
    st.subheader("🔬 Configure Atom Selections")
    st.info(f"""ℹ️ **Atom Selection:** Backbone + Sugar\n
    - For purines (A, G): N9, C8, C4\n
    - For pyrimidines (C, U): N1, C2, C6\n
    - For backbone and sugar atoms: "P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"\n
    """)
    
    
    # Create two columns for Reference and Query
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### 📋 Reference Structures")
        if ref_structure_data:
            selected_ref_name = st.selectbox(
                "Select structure to configure (excluding two bases in 5' and 3' by default)",
                options=[s['name'] for s in ref_structure_data],
                key="ref_dropdown",
                help="Choose a reference structure to configure its residue selection"
            )
            
            selected_ref = next((s for s in ref_structure_data if s['name'] == selected_ref_name), None)
            
            if selected_ref:
                st.markdown(f"**{selected_ref['name']}** ({selected_ref['num_residues']} residues)")
                
                # Display residue table
                structure_info = get_structure_info(selected_ref['path'])
                info_df = pd.DataFrame(structure_info)[['index', 'resnum', 'resname']]
                info_df.columns = ['Index (0-based)', 'Residue Number', 'Base Type']
                info_df['Index (1-based)'] = info_df['Index (0-based)'] + 1
                info_df = info_df[['Index (1-based)', 'Index (0-based)', 'Residue Number', 'Base Type']]
                
                with st.expander("📋 View Residue Table", expanded=False):
                    st.dataframe(info_df, use_container_width=True, height=min(300, len(structure_info) * 35 + 38))
                
                # Selection method
                selection_method = st.radio(
                    f"Selection method for {selected_ref['name']}",
                    ["Select by range", "Select specific residues", "Use all residues"],
                    key=f"method_ref_{selected_ref['name']}",
                    index=1,
                    horizontal=True
                )
                
                selected_indices = []
                
                if selection_method == "Select by range":
                    current_selection = st.session_state['ref_selections'].get(selected_ref['name'], [])
                    default_start = current_selection[0] + n_term_trim_ref if current_selection else n_term_trim_ref
                    default_end = current_selection[-1] + 1 if current_selection else max(n_term_trim_ref, len(structure_info) - c_term_trim_ref)
                    
                    c1, c2 = st.columns(2)
                    with c1:
                        start_idx = st.number_input(
                            "Start index (1-based)",
                            min_value=1,
                            max_value=len(structure_info),
                            value=default_start,
                            key=f"start_ref_{selected_ref['name']}"
                        )
                    with c2:
                        end_idx = st.number_input(
                            "End index (1-based, inclusive)",
                            min_value=1,
                            max_value=len(structure_info),
                            value=default_end,
                            key=f"end_ref_{selected_ref['name']}"
                        )
                    
                    if start_idx <= end_idx:
                        selected_indices = list(range(start_idx - 1, end_idx))
                        st.success(f"✓ Selected residues: {[i+1 for i in selected_indices]}")
                        # Auto-save the selection
                        st.session_state['ref_selections'][selected_ref['name']] = selected_indices
                    else:
                        st.error("Start index must be ≤ end index")
                        
                elif selection_method == "Select specific residues":
                    # Always use current trim values for default selection (updates when trim values change)
                    default_names = [structure_info[i]['full_name'] for i in range(n_term_trim_ref, len(structure_info)-c_term_trim_ref)]
                    
                    selected_names = st.multiselect(
                        "Select residues",
                        options=[info['full_name'] for info in structure_info],
                        default=default_names,
                        key=f"specific_ref_{selected_ref['name']}_n{n_term_trim_ref}_c{c_term_trim_ref}"
                    )
                    
                    
                    name_to_idx = {info['full_name']: info['index'] for info in structure_info}
                    selected_indices = [name_to_idx[name] for name in selected_names]
                    selected_indices.sort()
                    
                    if selected_indices:
                        st.success(f"✓ Selected {len(selected_indices)} residues: {[i+1 for i in selected_indices]}")
                        # Auto-save the selection
                        st.session_state['ref_selections'][selected_ref['name']] = selected_indices
                        
                else:  # Use all residues
                    selected_indices = list(range(len(structure_info)))
                    st.info(f"✓ Using all {len(selected_indices)} residues")
                    # Auto-save the selection
                    st.session_state['ref_selections'][selected_ref['name']] = selected_indices
                
                # Show current saved selection (now always up-to-date)
                if selected_ref['name'] in st.session_state['ref_selections']:
                    saved_indices = st.session_state['ref_selections'][selected_ref['name']]
                    st.info(f"**Current saved selection:** {len(saved_indices)} residues: {[i+1 for i in saved_indices]}")
        else:
            st.info("Upload reference structures to configure")
    
    with col2:
        st.markdown("### 📋 Query Structures")
        if query_structure_data:
            selected_query_name = st.selectbox(
                "Select structure to configure (excluding two bases in 5' and 3' by default)",
                options=[s['name'] for s in query_structure_data],
                key="query_dropdown",
                help="Choose a query structure to configure its residue selection"
            )
            
            selected_query = next((s for s in query_structure_data if s['name'] == selected_query_name), None)
            
            if selected_query:
                st.markdown(f"**{selected_query['name']}** ({selected_query['num_residues']} residues)")
                
                # Display residue table
                structure_info = get_structure_info(selected_query['path'])
                info_df = pd.DataFrame(structure_info)[['index', 'resnum', 'resname']]
                info_df.columns = ['Index (0-based)', 'Residue Number', 'Base Type']
                info_df['Index (1-based)'] = info_df['Index (0-based)'] + 1
                info_df = info_df[['Index (1-based)', 'Index (0-based)', 'Residue Number', 'Base Type']]
                
                with st.expander("📋 View Residue Table", expanded=False):
                    st.dataframe(info_df, use_container_width=True, height=min(300, len(structure_info) * 35 + 38))
                
                # Selection method
                selection_method = st.radio(
                    f"Selection method for {selected_query['name']}",
                    ["Select by range", "Select specific residues", "Use all residues"],
                    key=f"method_query_{selected_query['name']}",
                    index=1,
                    horizontal=True
                )
                
                selected_indices = []
                
                if selection_method == "Select by range":
                    current_selection = st.session_state['query_selections'].get(selected_query['name'], [])
                    default_start = current_selection[0] + n_term_trim_query if current_selection else 3
                    default_end = current_selection[-1] + 1 if current_selection else max(2, len(structure_info) - c_term_trim_query)
                    
                    c1, c2 = st.columns(2)
                    with c1:
                        start_idx = st.number_input(
                            "Start index (1-based)",
                            min_value=1,
                            max_value=len(structure_info),
                            value=default_start,
                            key=f"start_query_{selected_query['name']}"
                        )
                    with c2:
                        end_idx = st.number_input(
                            "End index (1-based, inclusive)",
                            min_value=1,
                            max_value=len(structure_info),
                            value=default_end,
                            key=f"end_query_{selected_query['name']}"
                        )
                    
                    if start_idx <= end_idx:
                        selected_indices = list(range(start_idx - 1, end_idx))
                        st.success(f"✓ Selected residues: {[i+1 for i in selected_indices]}")
                        # Auto-save the selection
                        st.session_state['query_selections'][selected_query['name']] = selected_indices
                    else:
                        st.error("Start index must be ≤ end index")
                        
                elif selection_method == "Select specific residues":
                    # Always use current trim values for default selection (updates when trim values change)
                    default_names = [structure_info[i]['full_name'] for i in range(n_term_trim_query, len(structure_info)-c_term_trim_query)]
                    
                    selected_names = st.multiselect(
                        "Select residues",
                        options=[info['full_name'] for info in structure_info],
                        default=default_names,
                        key=f"specific_query_{selected_query['name']}_n{n_term_trim_query}_c{c_term_trim_query}"
                    )
                    
                    name_to_idx = {info['full_name']: info['index'] for info in structure_info}
                    selected_indices = [name_to_idx[name] for name in selected_names]
                    selected_indices.sort()
                    
                    if selected_indices:
                        st.success(f"✓ Selected {len(selected_indices)} residues: {[i+1 for i in selected_indices]}")
                        # Auto-save the selection
                        st.session_state['query_selections'][selected_query['name']] = selected_indices
                        
                else:  # Use all residues
                    selected_indices = list(range(len(structure_info)))
                    st.info(f"✓ Using all {len(selected_indices)} residues")
                    # Auto-save the selection
                    st.session_state['query_selections'][selected_query['name']] = selected_indices
                
                # Show current saved selection (now always up-to-date)
                if selected_query['name'] in st.session_state['query_selections']:
                    saved_indices = st.session_state['query_selections'][selected_query['name']]
                    st.info(f"**Current saved selection:** {len(saved_indices)} residues: {[i+1 for i in saved_indices]}")
        else:
            st.info("Upload query structures to configure")
    
    # Step 4: Window Configuration
    st.sidebar.subheader("3️⃣ Window Configuration")
    
    # Check if all structures have selections
    all_ref_have_selections = all(s['name'] in st.session_state['ref_selections'] for s in ref_structure_data)
    all_query_have_selections = all(s['name'] in st.session_state['query_selections'] for s in query_structure_data)
    
    if all_ref_have_selections and all_query_have_selections and ref_structure_data and query_structure_data:
        # Find minimum selection size
        all_selections = list(st.session_state['ref_selections'].values()) + list(st.session_state['query_selections'].values())
        min_selection_size = min(len(sel) for sel in all_selections)
        
        window_size = st.sidebar.number_input(
            "Window Size",
            min_value=2,
            max_value=min_selection_size,
            value=min(4, min_selection_size),
            step=1,
            help="Number of residues per comparison window"
        )
        
        window_type = st.sidebar.radio(
            "Window Type",
            ["contiguous", "non-contiguous", "both"],
            index=0,
            help="Contiguous: sliding windows. Non-contiguous: all combinations"
        )
    else:
        st.sidebar.warning("⚠️ Configure selections first")
        window_size = 4
        window_type = "contiguous"
    
    
    # Step 5: Run Analysis
    st.sidebar.subheader("4️⃣ Run Analysis")
    
    can_run = (all_ref_have_selections and all_query_have_selections and 
               ref_structure_data and query_structure_data)
    
    if st.sidebar.button("🚀 Run Pairwise Analysis", type="primary", disabled=not can_run):
        if not can_run:
            st.error("Please upload and configure both reference and query structures")
            return
        
        # Run comparisons
        with st.spinner("Analyzing structures..."):
            results = []
            
            # For each reference structure
            for ref_struct in ref_structure_data:
                ref_indices = st.session_state['ref_selections'][ref_struct['name']]
                ref_windows = generate_windows_from_selection(ref_indices, window_size, window_type)
                
                if not ref_windows:
                    continue
                
                # For each reference window
                for ref_window in ref_windows:
                    # Extract reference coords
                    ref_coords = extract_window_coords(ref_struct['residues'], ref_window)
                    ref_com = calculate_COM(ref_coords)
                    ref_sequence = ''.join([ref_struct['residues'][i]['resname'] for i in ref_window])
                    
                    # Compare against all query structures
                    for query_struct in query_structure_data:
                        query_indices = st.session_state['query_selections'][query_struct['name']]
                        query_windows = generate_windows_from_selection(query_indices, window_size, window_type)
                        
                        for query_window in query_windows:
                            # Extract query coords
                            query_coords = extract_window_coords(query_struct['residues'], query_window)
                            query_com = calculate_COM(query_coords)
                            query_sequence = ''.join([query_struct['residues'][i]['resname'] for i in query_window])
                            
                            # Calculate RMSD
                            U, RMSD = calculate_rotation_rmsd(ref_coords, query_coords, ref_com, query_com)
                            
                            if U is None or RMSD is None:
                                RMSD = 999.0
                                U = np.eye(3)
                            
                            results.append({
                                'Reference': ref_struct['name'],
                                'Ref_Window': ref_window,
                                'Ref_Sequence': ref_sequence,
                                'Query': query_struct['name'],
                                'Query_Window': query_window,
                                'Query_Sequence': query_sequence,
                                'RMSD': RMSD,
                                'Rotation_Matrix': U,
                                'Ref_COM': ref_com,
                                'Query_COM': query_com,
                                'Ref_Path': ref_struct['path'],
                                'Query_Path': query_struct['path']
                            })
            
            results_df = pd.DataFrame(results)
            st.session_state['results'] = results_df
            st.session_state['ref_structure_data'] = ref_structure_data
            st.session_state['query_structure_data'] = query_structure_data
            
            st.success(f"✅ Analysis complete! {len(results_df)} comparisons performed.")
    
    # Display results
    if 'results' in st.session_state:
        results_df = st.session_state['results']
        
        st.markdown("---")
        st.subheader("📊 Results Summary")
        
        # RMSD threshold filter
        col1, col2 = st.columns([1, 3])
        with col1:
            rmsd_threshold = st.slider(
                "RMSD Threshold (Å)",
                min_value=0.0,
                max_value=10.0,
                value=3.0,
                step=0.1
            )
        
        filtered_df = results_df[results_df['RMSD'] <= rmsd_threshold]
        
        with col2:
            st.metric("Comparisons Below Threshold", f"{len(filtered_df)} / {len(results_df)}")
        
        # Best match per Reference-Query pair
        st.markdown("### 🏆 Best Match per Reference-Query Pair")
        
        if len(filtered_df) > 0:
            # Group by Reference and Query to find best match for each pair
            best_matches = filtered_df.loc[filtered_df.groupby(['Reference', 'Query'])['RMSD'].idxmin()]
            
            best_display = best_matches[['Reference', 'Query', 'Ref_Sequence', 'Query_Sequence', 'RMSD']].copy()
            best_display['RMSD'] = best_display['RMSD'].round(3)
            best_display.columns = ['Reference', 'Query', 'Ref Sequence', 'Query Sequence', 'RMSD (Å)']
            st.dataframe(best_display, use_container_width=True)
        else:
            st.warning("No matches found below threshold")
        
        # Full results
        with st.expander("📋 All Comparison Results"):
            if len(filtered_df) > 0:
                display_df = filtered_df[['Reference', 'Ref_Window', 'Ref_Sequence', 'Query', 'Query_Window', 'Query_Sequence', 'RMSD']].copy()
                
                # Format the window indices to be 1-based
                display_df['Ref_Residues'] = display_df['Ref_Window'].apply(lambda x: ','.join([str(i+1) for i in x]))
                display_df['Query_Residues'] = display_df['Query_Window'].apply(lambda x: ','.join([str(i+1) for i in x]))
                
                # Reorder columns
                display_df = display_df[['Reference', 'Ref_Residues', 'Ref_Sequence', 'Query', 'Query_Residues', 'Query_Sequence', 'RMSD']]
                display_df['RMSD'] = display_df['RMSD'].round(3)
                display_df = display_df.sort_values('RMSD').reset_index(drop=True)
                
                # Rename columns for better display
                display_df.columns = ['Reference', 'Ref_Indices', 'Ref_Sequence', 'Query', 'Query_Indices', 'Query_Sequence', 'RMSD (Å)']
                
                st.dataframe(display_df, use_container_width=True, height=400)
            else:
                st.info("No results to display")
        
        # Visualization
        st.markdown("---")
        st.subheader("🔬 3D Structure Visualization")
        
        if len(filtered_df) > 0:
            st.markdown("**Select a comparison to visualize:**")
            
            # Create dropdown options
            viz_options = []
            for idx, row in filtered_df.iterrows():
                ref_res_str = ','.join([str(i+1) for i in row['Ref_Window']])
                query_res_str = ','.join([str(i+1) for i in row['Query_Window']])
                option_text = f"{row['Reference']}[{ref_res_str}] ({row['Ref_Sequence']}) vs {row['Query']}[{query_res_str}] ({row['Query_Sequence']}) | RMSD: {row['RMSD']:.3f} Å"
                viz_options.append((idx, option_text))
            
            # Sort by RMSD
            viz_options.sort(key=lambda x: filtered_df.loc[x[0], 'RMSD'])
            
            selected_viz_idx = st.selectbox(
                "Choose comparison to visualize",
                options=[opt[0] for opt in viz_options],
                format_func=lambda idx: next(opt[1] for opt in viz_options if opt[0] == idx),
                help="All comparisons below RMSD threshold, sorted by RMSD"
            )
            
            # Get the selected comparison
            selected_row = filtered_df.loc[selected_viz_idx]
            
            # Import visualization function
            from visualization import create_structure_visualization
            
            # Display RMSD info
            #st.info(f"**RMSD: {selected_row['RMSD']:.3f} Å** ({len(selected_row['Query_Indices'])} residues) | Reference: {selected_row['Reference']}{selected_row['Ref_Residues']} ({selected_row['Ref_Sequence']}) | Query: {selected_row['Query']}{selected_row['Query_Residues']} ({selected_row['Query_Sequence']})")
            
            # Create visualization - wider display
            col1, col2, col3 = st.columns([0.5, 4, 0.5])
            
            with col2:
                try:
                    viz_html = create_structure_visualization(
                        selected_row['Ref_Path'],
                        selected_row['Query_Path'],
                        selected_row['Ref_Window'],
                        selected_row['Query_Window'],
                        selected_row['Rotation_Matrix'],
                        selected_row['Ref_COM'],
                        selected_row['Query_COM'],
                        selected_row['RMSD'],
                        ref_name=selected_row['Reference'],
                        query_name=selected_row['Query'],
                        ref_sequence=selected_row['Ref_Sequence'],
                        query_sequence=selected_row['Query_Sequence']
                    )
                    st.components.v1.html(viz_html, width=1400, height=750, scrolling=False)
                except Exception as e:
                    st.error(f"Error creating visualization: {str(e)}")
                    import traceback
                    st.code(traceback.format_exc())
            
            # Automatic Annotation Info
            st.markdown("---")
            st.success("✅ **Automatic Annotation:** When you click 'Download PNG' in the 3D viewer above, the image automatically includes RMSD, structure names, and sequences!")
            st.info("💡 **Customize font size:** Use the 'Annotation Font Size' dropdown in the viewer controls (top-right) to choose from Small, Medium, Large (default), or Extra Large fonts!")





            # Show transformation details
            with st.expander("🔧 Transformation Details"):
                col1, col2 = st.columns(2)
                
                with col1:
                    st.markdown("**Rotation Matrix (U):**")
                    st.dataframe(
                        pd.DataFrame(selected_row['Rotation_Matrix']).round(4),
                        use_container_width=True
                    )
                
                with col2:
                    st.markdown("**Translation Vectors:**")
                    st.write(f"Reference COM: [{selected_row['Ref_COM'][0]:.3f}, {selected_row['Ref_COM'][1]:.3f}, {selected_row['Ref_COM'][2]:.3f}]")
                    st.write(f"Query COM: [{selected_row['Query_COM'][0]:.3f}, {selected_row['Query_COM'][1]:.3f}, {selected_row['Query_COM'][2]:.3f}]")
            
            
            # Download aligned structures
            with st.expander("💾 Download Structure Files"):
                st.markdown("**Download extracted and aligned structures for external visualization**")
                
                from visualization import extract_window_pdb, transform_pdb_string
                
                # Extract reference window
                ref_pdb = extract_window_pdb(
                    selected_row['Ref_Path'],
                    selected_row['Ref_Window']
                )
                
                # Extract and transform query window
                query_pdb = extract_window_pdb(
                    selected_row['Query_Path'],
                    selected_row['Query_Window']
                )
                
                query_aligned_pdb = transform_pdb_string(
                    query_pdb,
                    selected_row['Rotation_Matrix'],
                    selected_row['Query_COM'],
                    selected_row['Ref_COM']
                )
                
                col1, col2, col3 = st.columns(3)
                
                with col1:
                    # Reference structure
                    ref_filename = f"ref_{selected_row['Reference'].replace('.pdb', '')}_{'_'.join(map(str, [i+1 for i in selected_row['Ref_Window']]))}.pdb"
                    st.download_button(
                        label="📥 Reference PDB",
                        data=ref_pdb,
                        file_name=ref_filename,
                        mime="chemical/x-pdb",
                        help="Original reference structure (selected residues only)"
                    )
                
                with col2:
                    # Query structure (original position)
                    query_filename = f"query_{selected_row['Query'].replace('.pdb', '')}_{'_'.join(map(str, [i+1 for i in selected_row['Query_Window']]))}.pdb"
                    st.download_button(
                        label="📥 Query PDB (Original)",
                        data=query_pdb,
                        file_name=query_filename,
                        mime="chemical/x-pdb",
                        help="Original query structure (selected residues only)"
                    )
                
                with col3:
                    # Query structure (aligned)
                    query_aligned_filename = f"query_aligned_{selected_row['Query'].replace('.pdb', '')}_{'_'.join(map(str, [i+1 for i in selected_row['Query_Window']]))}.pdb"
                    st.download_button(
                        label="📥 Query PDB (Aligned)",
                        data=query_aligned_pdb,
                        file_name=query_aligned_filename,
                        mime="chemical/x-pdb",
                        help="Query structure aligned to reference"
                    )
                
                # Combined aligned structure
                st.markdown("---")
                st.markdown("**Combined Aligned Structure (Reference + Query)**")
                
                # Create combined PDB with both structures
                combined_pdb_lines = []
                
                # Add header information as REMARK records
                combined_pdb_lines.append(f"REMARK Reference: {selected_row['Reference']}")
                combined_pdb_lines.append(f"REMARK Reference Residues: {','.join(map(str, [i+1 for i in selected_row['Ref_Window']]))}")
                combined_pdb_lines.append(f"REMARK Reference Sequence: {selected_row['Ref_Sequence']}")
                combined_pdb_lines.append(f"REMARK Query: {selected_row['Query']}")
                combined_pdb_lines.append(f"REMARK Query Residues: {','.join(map(str, [i+1 for i in selected_row['Query_Window']]))}")
                combined_pdb_lines.append(f"REMARK Query Sequence: {selected_row['Query_Sequence']}")
                combined_pdb_lines.append(f"REMARK RMSD: {selected_row['RMSD']:.3f} Angstroms")
                combined_pdb_lines.append("MODEL        1")
                
                # Add reference atoms with chain A
                for line in ref_pdb.split('\n'):
                    if line.startswith(('ATOM', 'HETATM')):
                        # Set chain to A for reference
                        modified_line = line[:21] + 'A' + line[22:]
                        combined_pdb_lines.append(modified_line)
                
                combined_pdb_lines.append("ENDMDL")
                combined_pdb_lines.append("MODEL        2")
                
                # Add aligned query atoms with chain B
                for line in query_aligned_pdb.split('\n'):
                    if line.startswith(('ATOM', 'HETATM')):
                        # Set chain to B for query
                        modified_line = line[:21] + 'B' + line[22:]
                        combined_pdb_lines.append(modified_line)
                
                combined_pdb_lines.append("ENDMDL")
                combined_pdb_lines.append("END")
                
                combined_pdb = '\n'.join(combined_pdb_lines)
                
                combined_filename = f"aligned_{selected_row['Reference'].replace('.pdb', '')}_{selected_row['Query'].replace('.pdb', '')}_rmsd_{selected_row['RMSD']:.3f}.pdb"
                
                st.download_button(
                    label="📥 Download Combined Aligned Structure",
                    data=combined_pdb,
                    file_name=combined_filename,
                    mime="chemical/x-pdb",
                    help="Reference (chain A) and aligned query (chain B) in one file",
                    use_container_width=True
                )
                
                st.info("💡 **Tip:** The combined PDB contains reference (chain A) and aligned query (chain B) - ready for PyMOL/Chimera")
        
        else:
            st.warning("No comparisons below RMSD threshold to visualize")
        
        # Export Results
        st.markdown("---")
        st.subheader("💾 Export Results")
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.markdown("**Download Results Table**")
            if len(filtered_df) > 0:
                export_df = filtered_df[['Reference', 'Ref_Window', 'Ref_Sequence', 'Query', 'Query_Window', 'Query_Sequence', 'RMSD']].copy()
                export_df['Ref_Residues'] = export_df['Ref_Window'].apply(lambda x: ','.join([str(i+1) for i in x]))
                export_df['Query_Residues'] = export_df['Query_Window'].apply(lambda x: ','.join([str(i+1) for i in x]))
                export_df = export_df[['Reference', 'Ref_Residues', 'Ref_Sequence', 'Query', 'Query_Residues', 'Query_Sequence', 'RMSD']]
                export_df = export_df.sort_values('RMSD').reset_index(drop=True)
                
                csv = export_df.to_csv(index=False)
                st.download_button(
                    label="📥 Download Results (CSV)",
                    data=csv,
                    file_name="rna_pairwise_comparison_results.csv",
                    mime="text/csv"
                )
            else:
                st.info("No results to export")
        
        with col2:
            st.markdown("**Download Aligned Structures**")
            if len(filtered_df) > 0 and st.button("📦 Generate PDB Archive"):
                with st.spinner("Creating archive..."):
                    import zipfile
                    from visualization_multi import extract_window_pdb, transform_pdb_string
                    
                    zip_buffer = io.BytesIO()
                    
                    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
                        for idx, row in filtered_df.iterrows():
                            comp_name = f"comp_{idx:03d}_rmsd_{row['RMSD']:.3f}"
                            
                            # Reference
                            ref_pdb = extract_window_pdb(row['Ref_Path'], row['Ref_Window'])
                            zip_file.writestr(f"{comp_name}/reference.pdb", ref_pdb)
                            
                            # Query original
                            query_pdb = extract_window_pdb(row['Query_Path'], row['Query_Window'])
                            zip_file.writestr(f"{comp_name}/query_original.pdb", query_pdb)
                            
                            # Query aligned
                            query_aligned = transform_pdb_string(
                                query_pdb,
                                row['Rotation_Matrix'],
                                row['Query_COM'],
                                row['Ref_COM']
                            )
                            zip_file.writestr(f"{comp_name}/query_aligned.pdb", query_aligned)
                            
                            # README
                            readme = f"""Comparison #{idx}
RMSD: {row['RMSD']:.3f} Å
Atom Selection: Backbone + Sugar (default)

Reference: {row['Reference']}
  Residues: {','.join([str(i+1) for i in row['Ref_Window']])}
  Sequence: {row['Ref_Sequence']}

Query: {row['Query']}
  Residues: {','.join([str(i+1) for i in row['Query_Window']])}
  Sequence: {row['Query_Sequence']}
"""
                            zip_file.writestr(f"{comp_name}/README.txt", readme)
                    
                    zip_buffer.seek(0)
                    
                    st.download_button(
                        label="📥 Download PDB Archive (ZIP)",
                        data=zip_buffer.getvalue(),
                        file_name="aligned_structures.zip",
                        mime="application/zip",
                        help=f"Contains {len(filtered_df)} comparison sets with reference, original query, and aligned query PDBs"
                    )
                    
                    st.success(f"✅ Archive ready! Contains {len(filtered_df)} comparisons with 3 PDB files each.")




if __name__ == "__main__":
    main()