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

# 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()

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_pdb(pdb_id):
    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'
    pdb_path = f'{pdb_id}.pdb'
    response = requests.get(pdb_url)
    if response.status_code == 200:
        with open(pdb_path, 'wb') as f:
            f.write(response.content)
        return pdb_path
    else:
        return None

def process_pdb(pdb_id, segment):
    pdb_path = fetch_pdb(pdb_id)
    if not pdb_path:
        return "Failed to fetch PDB file", None, None
    
    parser = PDBParser(QUIET=1)
    structure = parser.get_structure('protein', pdb_path)
    
    try:
        chain = structure[0][segment]
    except KeyError:
        return "Invalid Chain ID", None, None
    
    # Comprehensive amino acid mapping
    aa_dict = {
        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',
        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',
        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',
        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',
        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'
    }
    
    # Exclude non-amino acid residues
    sequence = [
        residue for residue in chain 
        if residue.get_resname().strip() in aa_dict
    ]
    
    # Prepare input for model prediction
    input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
    with torch.no_grad():
        outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()

    # Calculate scores and normalize them
    scores = expit(outputs[:, 1] - outputs[:, 0])
    normalized_scores = normalize_scores(scores)

    result_str = "\n".join(
        f"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}" 
        for res, score in zip(sequence, normalized_scores)
    )
    
    # Save the predictions to a file
    prediction_file = f"{pdb_id}_predictions.txt"
    with open(prediction_file, "w") as f:
        f.write(result_str)
    
    return result_str, molecule(pdb_path, random_scores, segment), prediction_file

def molecule(input_pdb, scores=None, segment='A'):
    mol = read_mol(input_pdb)  # Read PDB file content
    
    # Prepare high-scoring residues script if scores are provided
    high_score_script = ""
    if scores is not None:
        high_score_script = """
        // Reset all styles first
        viewer.getModel(0).setStyle({}, {});
        
        // Show only the selected chain
        viewer.getModel(0).setStyle(
            {"chain": "%s"}, 
            { cartoon: {colorscheme:"whiteCarbon"} }
        );
        
        // Highlight high-scoring residues only for the selected chain
        let highScoreResidues = [%s];
        viewer.getModel(0).setStyle(
            {"chain": "%s", "resi": highScoreResidues}, 
            {"stick": {"color": "red"}}
        );

        // Highlight high-scoring residues only for the selected chain
        let highScoreResidues2 = [%s];
        viewer.getModel(0).setStyle(
            {"chain": "%s", "resi": highScoreResidues2}, 
            {"stick": {"color": "orange"}}
        );
        """ % (segment, 
               ", ".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),
               segment,
              ", ".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),
               segment)
    
    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);
                viewer.addModel(pdb, "pdb");
                
                // Reset all styles and show only selected chain
                viewer.getModel(0).setStyle(
                    {{"chain": "{segment}"}}, 
                    {{ cartoon: {{ colorscheme:"whiteCarbon" }} }}
                );
                
                {high_score_script}
                
                // Add hover functionality
                viewer.setHoverable(
                    {{}}, 
                    true, 
                    function(atom, viewer, event, container) {{
                        if (!atom.label) {{
                            atom.label = viewer.addLabel(
                                atom.resn + ":" + 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>'

reps =    [
        {
          "model": 0,
          "style": "cartoon",
          "color": "whiteCarbon",
          "residue_range": "",
          "around": 0,
          "byres": False,
        }
    ]

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Protein Binding Site Prediction (Random Scores)")
    with gr.Row():
        pdb_input = gr.Textbox(value="2IWI", label="PDB ID", placeholder="Enter PDB ID here...")
        visualize_btn = gr.Button("Visualize Structure")

    molecule_output2 = Molecule3D(label="Protein Structure", reps=reps)

    with gr.Row():
        pdb_input = gr.Textbox(value="2IWI", label="PDB ID", placeholder="Enter PDB ID here...")
        segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
        prediction_btn = gr.Button("Predict Random Binding Site Scores")

    molecule_output = gr.HTML(label="Protein Structure")
    predictions_output = gr.Textbox(label="Binding Site Predictions")
    download_output = gr.File(label="Download Predictions")
    
    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)
    
    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])
    
    gr.Markdown("## Examples")
    gr.Examples(
        examples=[
            ["2IWI", "A"],
            ["7RPZ", "B"],
            ["3TJN", "C"]
        ],
        inputs=[pdb_input, segment_input],
        outputs=[predictions_output, molecule_output, download_output]
    )

demo.launch(share=True)