import warnings
warnings.filterwarnings("ignore")
import os
import re
import sys
import shutil
import random
import subprocess
import torch
import numpy as np
import pandas as pd
import MDAnalysis as mda
from typing import Optional

from pathlib import Path
from tempfile import NamedTemporaryFile

import huggingface_hub
from huggingface_hub.utils import GatedRepoError
from huggingface_hub import get_hf_file_metadata, hf_hub_download, login

import spaces
import gradio as gr

DEVICE = torch.device('cpu')
REPO_URL = "https://github.com/WaymentSteeleLab/Dyna-1.git"
DYNA_MODEL_ID = "gelnesr/Dyna-1"

def setup_environment():
    base_dir = Path(os.getcwd())
    dyna1_dir = base_dir / "Dyna-1"

    for filename in ["dyna1.pt", "dyna1-esm2.pt", "config.json"]:
        if not os.path.exists(f'Dyna-1/model/weights/{filename}'):
            print(f"Downloading {filename} from HuggingFace...")
            try:
                hf_hub_download(
                    repo_id=DYNA_MODEL_ID,
                    filename=filename,
                    repo_type='model',
                    local_dir=f'{dyna1_dir}/model/weights/',
                )
                print(f"Successfully downloaded {filename}")
            except Exception as e:
                print(f"Error downloading {filename}: {str(e)}")

    return dyna1_dir

dyna1_dir = setup_environment()
sys.path.insert(0, str(dyna1_dir))

from model.model import ESM_model
from esm.sdk.api import ESMProtein
from esm.utils.structure.protein_chain import ProteinChain
from transformers import AutoTokenizer
import utils

def check_permissions(token: Optional[str] = None) -> None:
    if token is None:
        raise gr.Error("Please log in to use this Space")
    try:
        url = huggingface_hub.hf_hub_url(repo_id="EvolutionaryScale/esm3-sm-open-v1", repo_type='model', filename="config.json")
        get_hf_file_metadata(url=url)
        return
    except GatedRepoError:
        raise gr.Error("You must have access to ... to run this Space. Please go through the gating process and come back.")

def validate_sequence(sequence):
    if not sequence:
        return None
    alphabets = {'protein': re.compile('^[acdefghiklmnpqrstvwy]*$', re.I)}
    if alphabets['protein'].search(sequence) is None:
        raise gr.Error('Invalid protein sequence. Please use standard amino acid letters.')
    return sequence.upper()

def process_structure(pdb_input, chain_id='A'):
    if not pdb_input:
        return None, None
            
    if isinstance(pdb_input, str) and len(pdb_input) == 4:
        try:
            protein_chain = ProteinChain.from_rcsb(pdb_input.upper(), chain_id=chain_id)
        except Exception as e:
            raise gr.Error(f"Error fetching PDB {pdb_input}: {str(e)}")
    else:
        temp_pdb = NamedTemporaryFile(suffix='.pdb', delete=False)
        try:
            if hasattr(pdb_input, 'name'):
                with open(pdb_input.name, 'rb') as f:
                    pdb_content = f.read()
            else:
                pdb_content = pdb_input.encode() if isinstance(pdb_input, str) else pdb_input
            
            temp_pdb.write(pdb_content)
            temp_pdb.close()
            
            protein_chain = ProteinChain.from_pdb(temp_pdb.name, chain_id=chain_id)
        except Exception as e:
            if os.path.exists(temp_pdb.name):
                os.unlink(temp_pdb.name)
            raise gr.Error(f"Error processing PDB file: {str(e)}")
        
        if os.path.exists(temp_pdb.name):
            os.unlink(temp_pdb.name)
            
    protein = ESMProtein.from_protein_chain(protein_chain)
    return protein, protein_chain

def write_probabilities_to_pdb(protein, probabilities, output_path):
    """Write probabilities to PDB B-factors and save the file."""
    temp_pdb = NamedTemporaryFile(suffix='.pdb', delete=False)
    protein.to_pdb(temp_pdb.name)
    
    curr = mda.Universe(temp_pdb.name)
    curr.add_TopologyAttr('bfactors')
    protein_out = curr.select_atoms("protein")
    
    for residue, prob in zip(protein_out.residues, probabilities):
        for atom in residue.atoms:
            atom.tempfactor = prob
    
    protein_out.write(output_path)
    os.unlink(temp_pdb.name)
    return output_path

def handle_name(name=None, pdb_input=None, model_version="ESM3"):
    """Processes the output file name given inputs of name and pdb; otherwise generates a random number"""
    if name:
        pdb_name = name
    elif pdb_input:
        if isinstance(pdb_input, str) and len(pdb_input) == 4:
            pdb_name = pdb_input
        else:
            if hasattr(pdb_input, 'name'):
                pdb_name = Path(pdb_input.name).stem
            else:
                pdb_name = str(random.randint(0, 100000))
    else:
        pdb_name = str(random.randint(0, 100000))
    return f'{pdb_name}-Dyna1{"" if model_version == "ESM3" else "-ESM2"}'

@spaces.GPU(duration=50)
def run_model(model, model_version='ESM2', seq_input=None, struct_input=None, sequence_id=None):
    if model_version == "ESM3":
        logits = model((seq_input, struct_input), sequence_id)
    else:
        logits = model(seq_input, sequence_id)
    return logits.cpu().detach()

def predict_dynamics(sequence=None, pdb_input=None, chain_id='A', use_pdb_seq=False, model_version="ESM3", name=None, oauth_token: Optional[str] = None):
    try:
        # Validate ESM2 requires sequence
        if model_version == "ESM2" and not sequence:
            raise ValueError("ESM-2 model requires a sequence input. Please provide a protein sequence.")
            
        if model_version == "ESM3" and not (sequence or pdb_input):
            raise ValueError("ESM-3 model requires either a sequence, structure (PDB ID/file), or both. Please provide at least one input.")
            
        base_name = handle_name(name, pdb_input, model_version)
        
        seq_input, struct_input = None, None
        sequence = validate_sequence(sequence) if sequence else None
        protein = None
        if model_version == "ESM3":
            model = ESM_model(method='esm3')
            model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1.pt', map_location=torch.device('cpu')), strict=False)
        else:  
            model = ESM_model(method='esm2', nheads=8, nlayers=12, layer=30).to(DEVICE)
            model.load_state_dict(torch.load('Dyna-1/model/weights/dyna1-esm2.pt', map_location=torch.device('cpu')), strict=False)

        model.eval()
        
        if pdb_input and model_version == "ESM3":
            protein, protein_chain = process_structure(pdb_input, chain_id)
            encoder = model.model.encode(protein)
            struct_input = encoder.structure[1:-1].unsqueeze(0)
            pdb_seq = protein.sequence
            seq_input = encoder.sequence[1:-1].unsqueeze(0)
            sequence_id = seq_input != 4099
            
            if not use_pdb_seq:
                seq_input = None
            
            if sequence and len(pdb_seq) != len(sequence):
                raise ValueError('Length of provided sequence does not match length of structure input.')
    
        if sequence:
            tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
            token_seq = tokenizer.encode(sequence, add_special_tokens=False, return_tensors='np')
            seq_input = torch.from_numpy(token_seq).to(DEVICE)
            sequence_id = seq_input != 4099
        
        if not (sequence or (pdb_input and model_version == "ESM3")):
            raise ValueError('Please provide a sequence' + (' or structure input' if model_version == "ESM3" else ''))
        
        logits = run_model(model, model_version, seq_input, struct_input, sequence_id)
        
        probabilities = utils.prob_adjusted(logits).numpy()
        
        seq_to_use = sequence if sequence else pdb_seq if pdb_input else sequence
        results_df = pd.DataFrame({
            'position': np.arange(1, len(probabilities) + 1),
            'residue': np.array(list(seq_to_use)),
            'p_exchange': probabilities,
        })
        
        csv_output = None
        pdb_output = None
        temp_csv = None
        temp_pdb = None
        
        try:
            temp_csv = NamedTemporaryFile(suffix='.csv', delete=False)
            results_df.to_csv(temp_csv.name, index=False)
            csv_output = temp_csv.name
            os.rename(csv_output, f"{base_name}.csv")
            csv_output = f"{base_name}.csv"
            
            if protein is not None and model_version == "ESM3":
                temp_pdb = NamedTemporaryFile(suffix='.pdb', delete=False)
                pdb_output = write_probabilities_to_pdb(protein, probabilities, temp_pdb.name)
                os.rename(pdb_output, f"{base_name}.pdb")
                pdb_output = f"{base_name}.pdb"
                
            return csv_output, pdb_output if pdb_output else None
            
        except Exception as e:
            if temp_csv and os.path.exists(temp_csv.name):
                os.unlink(temp_csv.name)
            if temp_pdb and os.path.exists(temp_pdb.name):
                os.unlink(temp_pdb.name)
            raise gr.Error(f"Error saving output files: {str(e)}")
            
    except Exception as e:
        raise gr.Error(str(e))

css = """
.gradio-container {
    font-family: 'Inter', system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
}
.tabs {
    margin-top: 0;
    margin-bottom: 0;
}
.gap {
    gap: 1rem;
}
"""

dyna1_app = gr.Blocks(theme=gr.themes.Soft(), mode="light")

with dyna1_app:

    gr.Markdown("# Dyna-1")
    gr.Markdown("## Predict micro-millisecond protein dynamics from sequence and/or structure")
    gr.Markdown("""[Paper](https://www.biorxiv.org/content/10.1101/2025.03.19.642801v1) |
    [GitHub](https://github.com/WaymentSteeleLab/Dyna-1) | 
    [Model](https://huggingface.co/gelnesr/Dyna-1) | 
    [Datasets](https://huggingface.co/datasets/gelnesr/RelaxDB) |
    [Colab](https://colab.research.google.com/github/WaymentSteeleLab/Dyna-1/blob/main/colab/Dyna_1.ipynb)""")
    
    gr.Markdown("""
    Dyna-1 predicts the probability that each residue experiences micro-millisecond motions.
    You can provide either a protein sequence, a structure (PDB ID or file), or both for the best performance.
    """)
    with gr.Row():
        gr.Markdown("""
            ## Instructions
            - Authorize access to ESM-3 by logging in to HuggingFace (required for ESM-3) 
            - Enter a protein sequence using standard amino acid letters (optional)
            - Provide a PDB ID (e.g., "1ubq") or upload a PDB file (optional)
            - Specify the chain ID if using a structure (default: A)
            - Choose whether to use the sequence from the PDB structure
                    
            You can toggle between using the ESM-3 and ESM-2 versions of the Dyna-1 model. To run with ESM-3, make sure you already 
            have access to the `EvolutionaryScale/esm3-sm-open-v1` weights [here](https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1).

            Note: The model will automatically set up the required environment on first run.

            Use of this HF Space is subject to a [Non-Commercial Use License](https://github.com/WaymentSteeleLab/Dyna-1/blob/main/LICENSE.txt).
            """)
        gr.Image(f"assets/dyna1.png", show_label=False)

    gr.LoginButton()
    model_choice = gr.Dropdown(
        choices=["ESM3", "ESM2"],
        value="ESM3",
        label="Choose model version"
    )

    with gr.Tabs() as tabs:
        with gr.Tab("Input"):
            with gr.Column(visible=True) as esm3_inputs:
                name_input = gr.Text(
                    label="Job Name (optional)",
                    placeholder="Enter name for the job. This will specify the output files. Leave blank to use PDB ID or a random number"
                )
                sequence_input_esm3 = gr.Textbox(
                    label="Protein Sequence",
                    placeholder="Enter protein sequence using standard amino acid letters",
                    lines=1
                )
                pdb_id = gr.Text(
                    label="PDB ID",
                    placeholder="Enter 4-letter PDB ID (e.g. 1UBQ)"
                )
                use_pdb_seq = gr.Checkbox(
                    label="Use sequence from PDB",
                    value=False
                )
                pdb_file = gr.File(
                    label="Or upload PDB file",
                    file_count="single"
                )
                chain_id = gr.Text(
                    label="Chain ID",
                    value="A",
                    placeholder="Enter chain ID"
                )
                submit_btn_esm3 = gr.Button("Predict", variant="primary")

            with gr.Column(visible=False) as esm2_inputs:
                name_input_esm2 = gr.Text(
                    label="Output Name (optional)",
                    placeholder="Enter name for the job. Leave blank to use a random number"
                )
                sequence_input_esm2 = gr.Textbox(
                    label="Protein Sequence",
                    placeholder="Enter protein sequence using standard amino acid letters",
                    lines=1
                )
                submit_btn_esm2 = gr.Button("Predict", variant="primary")

    with gr.Row(visible=True) as examples_esm3:
        label = gr.Textbox(label="Label", visible=False)

        examples = gr.Dataset(
            components=[label, sequence_input_esm3, pdb_id, chain_id, use_pdb_seq],
            samples=[
                ["Structure and its sequence", "-", "1ubq", "A", True],
                ["Structure and unique sequence", "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG", "1ubq", "A", False],
                ["Sequence only", "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG", "-", "-", False],
                ["Structure only", "-", "1ubq", "A", False]
            ],
            label="Examples"
        )

    with gr.Tab("Results"):
        with gr.Row(visible=True) as results_esm3:
            csv_output_esm3 = gr.File(label="Download Results (.csv)")
            pdb_output_esm3 = gr.File(label="Download PDB")
        with gr.Row(visible=False) as results_esm2:
            csv_output_esm2 = gr.File(label="Download Results (.csv)")

    def toggle_model_inputs(choice):
        if choice == "ESM3":
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)
        else:
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)

    model_choice.change(
        fn=toggle_model_inputs,
        inputs=model_choice,
        outputs=[esm3_inputs, esm2_inputs, results_esm3, results_esm2, examples_esm3]
    )

    def predict_esm3(name, sequence, pdb_id, pdb_file, chain_id, use_pdb_seq, oauth_token: gr.OAuthToken | None = None):
        if oauth_token is None:
            raise gr.Error("Please log in to use this Space")
        token_value = oauth_token.token
        check_permissions(token_value)

        csv_output, pdb_output = predict_dynamics(
            sequence=sequence,
            pdb_input=pdb_id if pdb_id else pdb_file,
            chain_id=chain_id,
            use_pdb_seq=use_pdb_seq,
            model_version="ESM3",
            name=name,
            oauth_token=token_value
        )
        return [csv_output, pdb_output]

    def predict_esm2(name, sequence):
        csv_output, _ = predict_dynamics(
            sequence=sequence,
            pdb_input=None,
            chain_id=None,
            use_pdb_seq=False,
            model_version="ESM2",
            name=name
        )
        return [csv_output]

    submit_btn_esm3.click(
        fn=predict_esm3,
        inputs=[name_input, sequence_input_esm3, pdb_id, pdb_file, chain_id, use_pdb_seq],
        outputs=[csv_output_esm3, pdb_output_esm3]
    )

    submit_btn_esm2.click(
        fn=predict_esm2,
        inputs=[name_input_esm2, sequence_input_esm2],
        outputs=[csv_output_esm2]
    )

    gr.Markdown("""
    ---
    This HuggingFace Space was created by Gina El Nesr [@ginaelnesr](https://twitter.com/ginaelnesr).
    """)

    gr.Markdown("""If you are using our code, datasets, or model, please use the following citation:
    ```bibtex
    @article {Dyna-1,
        author = {Wayment-Steele, Hannah K. and El Nesr, Gina and Hettiarachchi, Ramith and Kariyawasam, Hasindu and Ovchinnikov, Sergey and Kern, Dorothee},
        title = {Learning millisecond protein dynamics from what is missing in NMR spectra},
        year = {2025},
        doi = {10.1101/2025.03.19.642801},
        journal = {bioRxiv}
    }
    ```
    """)

if __name__ == "__main__":
    dyna1_app.launch(
        share=True
    )