from datetime import datetime
import gradio as gr
import requests
from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select
from Bio.PDB.Polypeptide import is_aa
from Bio.SeqUtils import seq1
from typing import Optional, Tuple
import numpy as np
import os
from gradio_molecule3d import Molecule3D

from model_loader import load_model

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

import re
import pandas as pd
import copy
import gc
import tempfile
import shutil
import atexit
import weakref

import transformers
from transformers import AutoTokenizer, DataCollatorForTokenClassification

from datasets import Dataset

from scipy.special import expit

# Create a temporary directory for this session
TEMP_DIR = tempfile.mkdtemp(prefix="protein_binding_")
print(f"Using temporary directory: {TEMP_DIR}")

# Registry to track created files for cleanup
_file_registry = weakref.WeakSet()

def cleanup_temp_files():
    """Clean up temporary directory on exit"""
    try:
        if os.path.exists(TEMP_DIR):
            shutil.rmtree(TEMP_DIR)
            print(f"Cleaned up temporary directory: {TEMP_DIR}")
    except Exception as e:
        print(f"Error cleaning up temp directory: {e}")

# Register cleanup function
atexit.register(cleanup_temp_files)

# Load model and move to device
checkpoint = 'ThorbenF/prot_t5_xl_uniref50_full_v2'
max_length = 1500
model, tokenizer = load_model(checkpoint, max_length)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()

def cleanup_files(*file_paths):
    """Helper function to clean up files"""
    for path in file_paths:
        if path and os.path.exists(path):
            try:
                os.remove(path)
            except Exception as e:
                print(f"Could not remove {path}: {e}")

def normalize_scores(scores):
    min_score = np.min(scores)
    max_score = np.max(scores)
    normalized = (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
    return normalized

def read_mol(pdb_path):
    """Read PDB file and return its content as a string"""
    with open(pdb_path, 'r') as f:
        return f.read()

def fetch_structure(pdb_id: str, output_dir: str = None) -> str:
    """
    Fetch the structure file for a given PDB ID. Prioritizes CIF files.
    If a structure file already exists locally, it uses that.
    """
    if output_dir is None:
        output_dir = TEMP_DIR
    file_path = download_structure(pdb_id, output_dir)
    return file_path

def download_structure(pdb_id: str, output_dir: str) -> str:
    """
    Attempt to download the structure file in CIF or PDB format.
    Returns the path to the downloaded file.
    """
    for ext in ['.cif', '.pdb']:
        file_path = os.path.join(output_dir, f"{pdb_id}{ext}")
        if os.path.exists(file_path):
            return file_path
        url = f"https://files.rcsb.org/download/{pdb_id}{ext}"
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            with open(file_path, 'wb') as f:
                f.write(response.content)
            return file_path
    return None

def convert_cif_to_pdb(cif_path: str, output_dir: str = None) -> str:
    """
    Convert a CIF file to PDB format using BioPython and return the PDB file path.
    """
    if output_dir is None:
        output_dir = TEMP_DIR
    pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb'))
    parser = MMCIFParser(QUIET=True)
    structure = parser.get_structure('protein', cif_path)
    io = PDBIO()
    io.set_structure(structure)
    io.save(pdb_path)
    
    # Clean up CIF file after conversion
    cleanup_files(cif_path)
    
    return pdb_path

def fetch_pdb(pdb_id):
    pdb_path = fetch_structure(pdb_id, TEMP_DIR)
    _, ext = os.path.splitext(pdb_path)
    if ext == '.cif':
        pdb_path = convert_cif_to_pdb(pdb_path, TEMP_DIR)
    return pdb_path

def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:
    """
    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores
    """
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure('protein', input_pdb)
    
    output_pdb = os.path.join(TEMP_DIR, f"{os.path.splitext(os.path.basename(input_pdb))[0]}_{chain_id}_predictions_scores.pdb")
    
    # Create scores dictionary for easy lookup
    scores_dict = {resi: score for resi, score in residue_scores}

    # Create a custom Select class
    class ResidueSelector(Select):
        def __init__(self, chain_id, selected_residues, scores_dict):
            self.chain_id = chain_id
            self.selected_residues = selected_residues
            self.scores_dict = scores_dict
        
        def accept_chain(self, chain):
            return chain.id == self.chain_id
        
        def accept_residue(self, residue):
            return residue.id[1] in self.selected_residues

        def accept_atom(self, atom):
            if atom.parent.id[1] in self.scores_dict:
                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100
            return True

    # Prepare output PDB with selected chain and residues, modified B-factors
    io = PDBIO()
    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)
    
    io.set_structure(structure[0])
    io.save(output_pdb, selector)
    
    # Clear references
    del structure, io, selector, scores_dict
    
    return output_pdb

def generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, score_type):
    """Generate PyMOL commands based on score type"""
    pymol_commands = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\nScore Type: {score_type}\n\n"
    
    pymol_commands += f"""
    # PyMOL Visualization Commands
    fetch {pdb_id}, protein
    hide everything, all
    show cartoon, chain {segment}
    color white, chain {segment}
    """
    
    # Define colors for each score bracket
    bracket_colors = {
        "0.0-0.2": "white",
        "0.2-0.4": "lightorange",
        "0.4-0.6": "yelloworange",
        "0.6-0.8": "orange",
        "0.8-1.0": "red"
    }
    
    # Add PyMOL commands for each score bracket
    for bracket, residues in residues_by_bracket.items():
        if residues:
            color = bracket_colors[bracket]
            resi_list = '+'.join(map(str, residues))
            pymol_commands += f"""
    select bracket_{bracket.replace('.', '').replace('-', '_')}, resi {resi_list} and chain {segment}
    show sticks, bracket_{bracket.replace('.', '').replace('-', '_')}
    color {color}, bracket_{bracket.replace('.', '').replace('-', '_')}
    """
    return pymol_commands

def generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, scores, current_time, score_type):
    """Generate results text based on score type"""
    result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\nScore Type: {score_type}\n\n"
    result_str += "Residues by Score Brackets:\n\n"
    
    # Add residues for each bracket
    for bracket, residues in residues_by_bracket.items():
        result_str += f"Bracket {bracket}:\n"
        result_str += f"Columns: Residue Name, Residue Number, One-letter Code, {score_type} Score\n"
        result_str += "\n".join([
            f"{res.resname} {res.id[1]} {sequence[i]} {scores[i]:.2f}" 
            for i, res in enumerate(protein_residues) if res.id[1] in residues
        ])
        result_str += "\n\n"
    
    return result_str

def process_pdb(pdb_id_or_file, segment, score_type='normalized'):
    # Determine if input is a PDB ID or file path
    if pdb_id_or_file.endswith('.pdb'):
        pdb_path = pdb_id_or_file
        pdb_id = os.path.splitext(os.path.basename(pdb_path))[0]
    else:
        pdb_id = pdb_id_or_file
        pdb_path = fetch_pdb(pdb_id)
    
    # Determine the file format and choose the appropriate parser
    _, ext = os.path.splitext(pdb_path)
    parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)
    
    # Parse the structure file
    structure = parser.get_structure('protein', pdb_path)
    
    # Extract the specified chain
    chain = structure[0][segment]
    
    protein_residues = [res for res in chain if is_aa(res)]
    sequence = "".join(seq1(res.resname) for res in protein_residues)
    sequence_id = [res.id[1] for res in protein_residues]

    input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
    
    with torch.no_grad():
        outputs = model(input_ids).logits
        outputs_cpu = outputs.detach().cpu().numpy().squeeze()
        
        # Explicitly delete GPU tensors
        del outputs, input_ids
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    # Calculate scores and normalize them
    raw_scores = expit(outputs_cpu[:, 1] - outputs_cpu[:, 0])
    normalized_scores = normalize_scores(raw_scores)
    
    # Clear outputs_cpu
    del outputs_cpu
    
    # Choose which scores to use based on score_type
    display_scores = normalized_scores if score_type == 'normalized' else raw_scores
    
    # Zip residues with scores to track the residue ID and score
    residue_scores = [(resi, score) for resi, score in zip(sequence_id, display_scores)]
    
    # Also save both score types for later use
    raw_residue_scores = [(resi, score) for resi, score in zip(sequence_id, raw_scores)]
    norm_residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]
    
    # Define the score brackets
    score_brackets = {
        "0.0-0.2": (0.0, 0.2),
        "0.2-0.4": (0.2, 0.4),
        "0.4-0.6": (0.4, 0.6),
        "0.6-0.8": (0.6, 0.8),
        "0.8-1.0": (0.8, 1.0)
    }
    
    # Initialize a dictionary to store residues by bracket
    residues_by_bracket = {bracket: [] for bracket in score_brackets}
    
    # Categorize residues into brackets
    for resi, score in residue_scores:
        for bracket, (lower, upper) in score_brackets.items():
            if lower <= score < upper:
                residues_by_bracket[bracket].append(resi)
                break
    
    # Generate timestamp
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Generate result text and PyMOL commands based on score type
    display_score_type = "Normalized" if score_type == 'normalized' else "Raw"
    result_str = generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, 
                                      display_scores, current_time, display_score_type)
    pymol_commands = generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, display_score_type)
    
    # Create chain-specific PDB with scores in B-factor
    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)

    # Molecule visualization with updated script with color mapping
    mol_vis = molecule(pdb_path, residue_scores, segment)
    
    # Create prediction file
    prediction_file = os.path.join(TEMP_DIR, f"{pdb_id}_{display_score_type.lower()}_binding_site_residues.txt")
    with open(prediction_file, "w") as f:
        f.write(result_str)
    
    scored_pdb_name = os.path.join(TEMP_DIR, f"{pdb_id}_{segment}_{display_score_type.lower()}_predictions_scores.pdb")
    os.rename(scored_pdb, scored_pdb_name)
    
    # Clear large objects from memory
    del structure, chain, protein_residues, raw_scores, normalized_scores, display_scores
    gc.collect()
    
    return pymol_commands, mol_vis, [prediction_file, scored_pdb_name], raw_residue_scores, norm_residue_scores, pdb_id, segment

def molecule(input_pdb, residue_scores=None, segment='A'):
    # Read PDB file content
    mol = read_mol(input_pdb)
    
    # Prepare high-scoring residues script if scores are provided
    high_score_script = ""
    if residue_scores is not None:
        # Filter residues based on their scores
        class1_score_residues = [resi for resi, score in residue_scores if 0.0 < score <= 0.2]
        class2_score_residues = [resi for resi, score in residue_scores if 0.2 < score <= 0.4]
        class3_score_residues = [resi for resi, score in residue_scores if 0.4 < score <= 0.6]
        class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]
        class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]
        
        high_score_script = """
        // Load the original model and apply white cartoon style
        let chainModel = viewer.addModel(pdb, "pdb");
        chainModel.setStyle({}, {});
        chainModel.setStyle(
            {"chain": "%s"}, 
            {"cartoon": {"color": "white"}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class1Model = viewer.addModel(pdb, "pdb");
        class1Model.setStyle({}, {});
        class1Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFFFFFF", "opacity": 0.5}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class2Model = viewer.addModel(pdb, "pdb");
        class2Model.setStyle({}, {});
        class2Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFFD580", "opacity": 0.7}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class3Model = viewer.addModel(pdb, "pdb");
        class3Model.setStyle({}, {});
        class3Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFFA500", "opacity": 1}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class4Model = viewer.addModel(pdb, "pdb");
        class4Model.setStyle({}, {});
        class4Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFF4500", "opacity": 1}}
        );

        // Create a new model for high-scoring residues and apply red sticks style
        let class5Model = viewer.addModel(pdb, "pdb");
        class5Model.setStyle({}, {});
        class5Model.setStyle(
            {"chain": "%s", "resi": [%s]}, 
            {"stick": {"color": "0xFF0000", "alpha": 1}}
        );

        """ % (
            segment,
            segment,
            ", ".join(str(resi) for resi in class1_score_residues),
            segment,
            ", ".join(str(resi) for resi in class2_score_residues),
            segment,
            ", ".join(str(resi) for resi in class3_score_residues),
            segment,
            ", ".join(str(resi) for resi in class4_score_residues),
            segment,
            ", ".join(str(resi) for resi in class5_score_residues)
        )
    
    # Generate the full HTML content
    html_content = f"""
    <!DOCTYPE html>
    <html>
    <head>    
        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
        <style>
        .mol-container {{
            width: 100%;
            height: 700px;
            position: relative;
        }}
        </style>
        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js"></script>
        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
    </head>
    <body>
        <div id="container" class="mol-container"></div>
        <script>
            let pdb = `{mol}`; // Use template literal to properly escape PDB content
            $(document).ready(function () {{
                let element = $("#container");
                let config = {{ backgroundColor: "white" }};
                let viewer = $3Dmol.createViewer(element, config);
                
                {high_score_script}
                
                // Add hover functionality
                viewer.setHoverable(
                    {{}}, 
                    true, 
                    function(atom, viewer, event, container) {{
                        if (!atom.label) {{
                            atom.label = viewer.addLabel(
                                atom.resn + ":" +atom.resi + ":" + atom.atom, 
                                {{
                                    position: atom, 
                                    backgroundColor: 'mintcream', 
                                    fontColor: 'black',
                                    fontSize: 18,
                                    padding: 4
                                }}
                            );
                        }}
                    }},
                    function(atom, viewer) {{
                        if (atom.label) {{
                            viewer.removeLabel(atom.label);
                            delete atom.label;
                        }}
                    }}
                );
                
                viewer.zoomTo();
                viewer.render();
                viewer.zoom(0.8, 2000);
            }});
        </script>
    </body>
    </html>
    """
    
    # Clear mol from memory after use
    del mol
    
    # Return the HTML content within an iframe safely encoded for special characters
    return f'<iframe width="100%" height="700" srcdoc="{html_content.replace(chr(34), "&quot;").replace(chr(39), "&#39;")}"></iframe>'

with gr.Blocks(css="""
    /* Customize Gradio button colors */
    #visualize-btn, #predict-btn {
        background-color: #FF7300; /* Deep orange */
        color: white;
        border-radius: 5px;
        padding: 10px;
        font-weight: bold;
    }
    #visualize-btn:hover, #predict-btn:hover {
        background-color: #CC5C00; /* Darkened orange on hover */
    }
""") as demo:
    gr.Markdown("# Protein Binding Site Prediction")
    
    # Mode selection
    mode = gr.Radio(
        choices=["PDB ID", "Upload File"],
        value="PDB ID",
        label="Input Mode",
        info="Choose whether to input a PDB ID or upload a PDB/CIF file."
    )

    # Input components based on mode
    pdb_input = gr.Textbox(value="2F6V", label="PDB ID", placeholder="Enter PDB ID here...")
    pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
    visualize_btn = gr.Button("Visualize Structure", elem_id="visualize-btn")

    molecule_output2 = Molecule3D(label="Protein Structure", reps=[
        {
            "model": 0,
            "style": "cartoon",
            "color": "whiteCarbon",
            "residue_range": "",
            "around": 0,
            "byres": False,
        }
    ])

    with gr.Row():
        segment_input = gr.Textbox(value="A", label="Chain ID (protein)", placeholder="Enter Chain ID here...",
        info="Choose in which chain to predict binding sites.")
        prediction_btn = gr.Button("Predict Binding Site", elem_id="predict-btn")

    # Add score type selector
    score_type = gr.Radio(
        choices=["Normalized Scores", "Raw Scores"],
        value="Normalized Scores",
        label="Score Visualization Type",
        info="Choose which score type to visualize"
    )

    molecule_output = gr.HTML(label="Protein Structure")
    explanation_vis = gr.Markdown("""
    Score dependent colorcoding:
    - 0.0-0.2: white  
    - 0.2–0.4: light orange  
    - 0.4–0.6: yellow orange
    - 0.6–0.8: orange
    - 0.8–1.0: red
    """)
    predictions_output = gr.Textbox(label="Visualize Prediction with PyMol")
    gr.Markdown("### Download:\n- List of predicted binding site residues\n- PDB with score in beta factor column")
    download_output = gr.File(label="Download Files", file_count="multiple")
    
    # Store these as state variables so we can switch between them
    raw_scores_state = gr.State(None)
    norm_scores_state = gr.State(None)
    last_pdb_path = gr.State(None)
    last_segment = gr.State(None)
    last_pdb_id = gr.State(None)
    
    def process_interface(mode, pdb_id, pdb_file, chain_id, score_type_val):
        try:
            selected_score_type = 'normalized' if score_type_val == "Normalized Scores" else 'raw'
            
            # First get the actual PDB file path
            if mode == "PDB ID":
                pdb_path = fetch_pdb(pdb_id)
                
                pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_id_result, segment = process_pdb(pdb_path, chain_id, selected_score_type)
                return pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_path, chain_id, pdb_id_result
            elif mode == "Upload File":
                _, ext = os.path.splitext(pdb_file.name)
                file_path = os.path.join(TEMP_DIR, f"{os.path.basename(pdb_file.name)}")
                
                # Copy uploaded file to temp directory
                shutil.copy(pdb_file.name, file_path)
                
                if ext == '.cif':
                    pdb_path = convert_cif_to_pdb(file_path, TEMP_DIR)
                else:
                    pdb_path = file_path
                
                pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_id_result, segment = process_pdb(pdb_path, chain_id, selected_score_type)
                return pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_path, chain_id, pdb_id_result
        finally:
            # Force garbage collection after processing
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    
    def update_visualization_and_files(score_type_val, raw_scores, norm_scores, pdb_path, segment, pdb_id):
        if raw_scores is None or norm_scores is None or pdb_path is None or segment is None or pdb_id is None:
            return None, None, None
        
        try:
            # Choose scores based on radio button selection
            selected_score_type = 'normalized' if score_type_val == "Normalized Scores" else 'raw'
            selected_scores = norm_scores if selected_score_type == 'normalized' else raw_scores
            
            # Generate visualization with selected scores
            mol_vis = molecule(pdb_path, selected_scores, segment)
            
            # Generate PyMOL commands and downloadable files
            # Get structure for residue info
            _, ext = os.path.splitext(pdb_path)
            parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)
            structure = parser.get_structure('protein', pdb_path)
            chain = structure[0][segment]
            protein_residues = [res for res in chain if is_aa(res)]
            sequence = "".join(seq1(res.resname) for res in protein_residues)
            
            # Define score brackets
            score_brackets = {
                "0.0-0.2": (0.0, 0.2),
                "0.2-0.4": (0.2, 0.4),
                "0.4-0.6": (0.4, 0.6),
                "0.6-0.8": (0.6, 0.8),
                "0.8-1.0": (0.8, 1.0)
            }
            
            # Initialize a dictionary to store residues by bracket
            residues_by_bracket = {bracket: [] for bracket in score_brackets}
            
            # Categorize residues into brackets
            for resi, score in selected_scores:
                for bracket, (lower, upper) in score_brackets.items():
                    if lower <= score < upper:
                        residues_by_bracket[bracket].append(resi)
                        break
            
            # Generate timestamp
            current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            
            # Generate result text and PyMOL commands based on score type
            display_score_type = "Normalized" if selected_score_type == 'normalized' else "Raw"
            scores_array = [score for _, score in selected_scores]
            result_str = generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, 
                                               scores_array, current_time, display_score_type)
            pymol_commands = generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, display_score_type)
            
            # Create chain-specific PDB with scores in B-factor
            scored_pdb = create_chain_specific_pdb(pdb_path, segment, selected_scores, protein_residues)
            
            # Create prediction file
            prediction_file = os.path.join(TEMP_DIR, f"{pdb_id}_{display_score_type.lower()}_binding_site_residues.txt")
            with open(prediction_file, "w") as f:
                f.write(result_str)
            
            scored_pdb_name = os.path.join(TEMP_DIR, f"{pdb_id}_{segment}_{display_score_type.lower()}_predictions_scores.pdb")
            os.rename(scored_pdb, scored_pdb_name)
            
            # Clear memory
            del structure, chain, protein_residues, scores_array
            
            return mol_vis, pymol_commands, [prediction_file, scored_pdb_name]
        finally:
            gc.collect()

    def fetch_interface(mode, pdb_id, pdb_file):
        if mode == "PDB ID":
            return fetch_pdb(pdb_id)
        elif mode == "Upload File":
            _, ext = os.path.splitext(pdb_file.name)
            file_path = os.path.join(TEMP_DIR, f"{os.path.basename(pdb_file.name)}")
            shutil.copy(pdb_file.name, file_path)
            if ext == '.cif':
                pdb_path = convert_cif_to_pdb(file_path, TEMP_DIR)
            else:
                pdb_path = file_path
            return pdb_path

    def toggle_mode(selected_mode):
        if selected_mode == "PDB ID":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)
        
    mode.change(
        toggle_mode,
        inputs=[mode],
        outputs=[pdb_input, pdb_file]
    )

    prediction_btn.click(
        process_interface, 
        inputs=[mode, pdb_input, pdb_file, segment_input, score_type], 
        outputs=[predictions_output, molecule_output, download_output, 
                raw_scores_state, norm_scores_state, last_pdb_path, last_segment, last_pdb_id]
    )
    
    # Update visualization, PyMOL commands, and files when score type changes
    score_type.change(
        update_visualization_and_files,
        inputs=[score_type, raw_scores_state, norm_scores_state, last_pdb_path, last_segment, last_pdb_id],
        outputs=[molecule_output, predictions_output, download_output]
    )

    visualize_btn.click(
        fetch_interface, 
        inputs=[mode, pdb_input, pdb_file], 
        outputs=molecule_output2
    )

    gr.Markdown("## Examples")
    gr.Examples(
        examples=[
            ["7RPZ", "A"],
            ["2IWI", "B"],
            ["7LCJ", "R"],
            ["4OBE", "A"]
        ],
        inputs=[pdb_input, segment_input],
        outputs=[predictions_output, molecule_output, download_output]
    )

    def predict_utils(sequence):
        try:
            input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
            with torch.no_grad():
                outputs = model(input_ids).logits
                outputs_cpu = outputs.detach().cpu().numpy().squeeze()
                
                # Explicitly delete GPU tensors
                del outputs, input_ids
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
        
            raw_scores = expit(outputs_cpu[:, 1] - outputs_cpu[:, 0])
            normalized_scores = normalize_scores(raw_scores)
            
            result = {
                "raw_scores": raw_scores.tolist(),
                "normalized_scores": normalized_scores.tolist()
            }
            
            # Clear memory
            del outputs_cpu, raw_scores, normalized_scores
            gc.collect()
            
            return result
        except Exception as e:
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            raise e

    dummy_input = gr.Textbox(visible=False)
    dummy_output = gr.Textbox(visible=False)
    
    dummy_btn = gr.Button("Predict Sequence", visible=False)
    dummy_btn.click(
        predict_utils,
        inputs=[dummy_input],
        outputs=[dummy_output]
    )

demo.launch(share=True)