from datetime import datetime
import gradio as gr
import requests
from Bio.PDB import PDBParser, MMCIFParser, PDBIO
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 transformers, datasets
from transformers import AutoTokenizer
from transformers import DataCollatorForTokenClassification

from datasets import Dataset

from scipy.special import expit

from datetime import datetime
import gradio as gr
import requests
from Bio.PDB import PDBParser, MMCIFParser, PDBIO
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

import re
import pandas as pd
import copy

from scipy.special import expit


# Load model and move to device
checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
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()

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

import re
import pandas as pd
import copy

from scipy.special import expit

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

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 = ".") -> Optional[str]:
    """
    Fetch the structure file for a given PDB ID. Prioritizes CIF files.
    If a structure file already exists locally, it uses that.
    """
    file_path = download_structure(pdb_id, output_dir)
    if file_path:
        return file_path
    else:
        return None

def download_structure(pdb_id: str, output_dir: str) -> Optional[str]:
    """
    Attempt to download the structure file in CIF or PDB format.
    Returns the path to the downloaded file, or None if download fails.
    """
    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}"
        try:
            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
        except Exception as e:
            print(f"Download error for {pdb_id}{ext}: {e}")
    return None

def convert_cif_to_pdb(cif_path: str, output_dir: str = ".") -> str:
    """
    Convert a CIF file to PDB format using BioPython and return the PDB file path.
    """
    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)
    return pdb_path

def fetch_pdb(pdb_id):
    pdb_path = fetch_structure(pdb_id)
    if not pdb_path:
        return None
    _, ext = os.path.splitext(pdb_path)
    if ext == '.cif':
        pdb_path = convert_cif_to_pdb(pdb_path)
    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
    """
    # Read the original PDB file
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure('protein', input_pdb)
    
    # Prepare a new structure with only the specified chain and selected residues
    output_pdb = f"{os.path.splitext(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 = 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)
    
    return output_pdb

def calculate_geometric_center(pdb_path: str, high_score_residues: list, chain_id: str):
    """
    Calculate the geometric center of high-scoring residues
    """
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure('protein', pdb_path)
    
    # Collect coordinates of CA atoms from high-scoring residues
    coords = []
    for model in structure:
        for chain in model:
            if chain.id == chain_id:
                for residue in chain:
                    if residue.id[1] in high_score_residues:
                        if 'CA' in residue:  # Use alpha carbon as representative
                            ca_atom = residue['CA']
                            coords.append(ca_atom.coord)
    
    # Calculate geometric center
    if coords:
        center = np.mean(coords, axis=0)
        return center
    return None

def process_pdb(pdb_id_or_file, segment):
    # 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)
    
    if not pdb_path:
        return "Failed to fetch PDB file", None, None
    
    # 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)
    
    try:
        # Parse the structure file
        structure = parser.get_structure('protein', pdb_path)
    except Exception as e:
        return f"Error parsing structure file: {e}", None, None
    
    # Extract the specified chain
    try:
        chain = structure[0][segment]
    except KeyError:
        return "Invalid Chain ID", None, None
    
    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]
    
    scores = np.random.rand(len(sequence))
    normalized_scores = normalize_scores(scores)
    
    # Zip residues with scores to track the residue ID and score
    residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]

    
    # Identify high scoring residues (> 0.5)
    high_score_residues = [resi for resi, score in residue_scores if score > 0.5]
    
    # Preparing the result: only print high scoring residues
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
    result_str += "High-scoring Residues (Score > 0.5):\n"
    result_str += "Columns: Residue Name, Residue Number, One-letter Code, Normalized Score\n\n"
    result_str += "\n".join([
        f"{res.resname} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
        for i, res in enumerate(protein_residues) if res.id[1] in high_score_residues
    ])

    # 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)#, color_map)

    # Improved PyMOL command suggestions
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    pymol_commands = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
    
    pymol_commands += f"""
# PyMOL Visualization Commands
load {os.path.abspath(pdb_path)}, protein
hide everything, all
show cartoon, chain {segment}
color white, chain {segment}
"""
    
    # Color specific residues
    for score_range, color in [
        (high_score_residues, "red")
    ]:
        if score_range:
            resi_list = '+'.join(map(str, score_range))
            pymol_commands += f"""
select high_score_residues, resi {resi_list} and chain {segment}
show sticks, high_score_residues
color {color}, high_score_residues
"""
    # Create prediction and scored PDB files
    prediction_file = f"{pdb_id}_binding_site_residues.txt"
    with open(prediction_file, "w") as f:
        f.write(result_str)
    
    return pymol_commands, mol_vis, [prediction_file,scored_pdb]

def molecule(input_pdb, residue_scores=None, segment='A'):
    # More granular scoring for visualization
    mol = read_mol(input_pdb)  # Read PDB file content

    # 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.5 < score <= 0.6]
        class2_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.7]
        class3_score_residues = [resi for resi, score in residue_scores if 0.7 < score <= 0.8]
        class4_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 0.9]
        class5_score_residues = [resi for resi, score in residue_scores if 0.9 < 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": "blue"}}
        );

        // 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": "lightblue"}}
        );

        // 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": "white"}}
        );

        // 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": "orange"}}
        );

        // 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": "red"}}
        );

        """ % (
            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: 12,
                                    padding: 2
                                }}
                            );
                        }}
                    }},
                    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>
    """
    
    # 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>'

# Gradio UI
with gr.Blocks() 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="4BDU", 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")

    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", placeholder="Enter Chain ID here...")
        prediction_btn = gr.Button("Predict Binding Site")

    molecule_output = gr.HTML(label="Protein Structure")
    explanation_vis = gr.Markdown("""
    Residues with a score > 0.5 are considered binding sites and represented as sticks with the score dependent colorcoding:
    - 0.5-0.6: blue  
    - 0.6–0.7: light blue  
    - 0.7–0.8: white
    - 0.8–0.9: orange
    - 0.9–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")
    
    def process_interface(mode, pdb_id, pdb_file, chain_id):
        if mode == "PDB ID":
            return process_pdb(pdb_id, chain_id)
        elif mode == "Upload File":
            _, ext = os.path.splitext(pdb_file.name)
            file_path = os.path.join('./', f"{_}{ext}")
            if ext == '.cif':
                pdb_path = convert_cif_to_pdb(file_path)
            else:
                pdb_path= file_path
            return process_pdb(pdb_path, chain_id)
        else:
            return "Error: Invalid mode selected", None, None

    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('./', f"{_}{ext}")
            #print(ext)
            if ext == '.cif':
                pdb_path = convert_cif_to_pdb(file_path)
            else:
                pdb_path= file_path
            #print(pdb_path)
            return pdb_path
        else:
            return "Error: Invalid mode selected"

    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], 
        outputs=[predictions_output, molecule_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"],
            ["2F6V", "A"]
        ],
        inputs=[pdb_input, segment_input],
        outputs=[predictions_output, molecule_output, download_output]
    )

demo.launch(share=True)