Update app.py

app.py

@@ -999,6 +999,67 @@ def main():
                     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_Indices']
+                )
+                
+                # Extract and transform query window
+                query_pdb = extract_window_pdb(
+                    selected_row['Query_Path'],
+                    selected_row['Query_Indices']
+                )
+                
+                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_Indices']]))}.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_Indices']]))}.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_Indices']]))}.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"
+                    )
+                
+                st.info("💡 **Tip:** Load reference and aligned query together in PyMOL/Chimera to examine the superposition")
         
         else:
             st.warning("No comparisons below RMSD threshold to visualize")
@@ -1076,11 +1137,16 @@ Query: {row['Query']}
                     zip_buffer.seek(0)
                     
                     st.download_button(
-                        label="📥 Download ZIP",
+                        label="📥 Download PDB Archive (ZIP)",
                         data=zip_buffer.getvalue(),
                         file_name="aligned_structures.zip",
-                        mime="application/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__":