import gradio as gr
import torch
from pathlib import Path
from helpers import load_model, parse_fasta_string
from io import StringIO
import csv
import tempfile
import transformers
import traceback
# mute esm warning for weights
transformers.logging.set_verbosity_error()

# Constants
MODEL_NAME = "esm2_t33_650M_UR50D"
CURRENT_DIR = Path(__file__).parent
PATH_MODEL = CURRENT_DIR / "model"
DEVICE = "cpu"
# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # arount 2 mins for fireprot
VALID_AMINO_ACIDS = set("ACDEFGHIKLMNPQRSTVWY")

model, tokenizer = load_model(MODEL_NAME, PATH_MODEL, DEVICE)

def _error(message):
    """Return outputs that surface a clear error in the UI (results, download, status)."""
    return None, None, gr.update(value=f"⚠️ **Input error:** {message}", visible=True)


def render_results(results, threshold):
    """Build the ranked table and clear any error. CSV is generated on-demand via download callback.

    Pure post-processing: `results` is already sorted by Tm at prediction time, and the
    threshold only controls the Yes/No label. Runs in well under a second with no model
    inference. Returns the same (output, error_box) pair as the error/success paths so
    every entry point owns the full result area.
    """
    clear_error = gr.update(value="", visible=False)
    if not results:
        # No predictions yet (e.g. threshold changed before any successful run).
        return None, clear_error

    cutoff = float(threshold)
    table = [
        [i + 1, r["id"], r["tm"], "Yes" if r["tm"] > cutoff else "No"]
        for i, r in enumerate(results)
    ]

    return table, clear_error


def generate_csv(results, threshold):
    """Generate CSV file from predictions. Called on-demand when user clicks download."""
    if not results:
        return None

    cutoff = float(threshold)
    table = [
        [i + 1, r["id"], r["tm"], "Yes" if r["tm"] > cutoff else "No"]
        for i, r in enumerate(results)
    ]

    csv_buffer = StringIO()
    writer = csv.writer(csv_buffer)
    writer.writerow(["Rank", "ID", "Predicted Tm [°C]", f"Thermostable (threshold > {threshold}°C)"])
    writer.writerows(table)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".csv", mode="w", encoding="utf-8") as tmp:
        tmp.write(csv_buffer.getvalue())
        tmp_path = tmp.name

    return tmp_path


def predict_tm(seq_text, seq_file, threshold):
    # predict_tm drives one extra output (results_state) beyond render_results, so every
    # exit appends the cached-predictions slot — [] on any error, the predictions on success.
    def error(message):
        table, error_box = _error(message)[:2]
        return table, error_box, ([],)

    # --- Input validation (always returns clear messages) ---
    if seq_file is not None:
        try:
            with open(seq_file.name, "r", encoding="utf-8") as f:
                fasta_str = f.read()
        except Exception as e:
            return error(f"Could not read uploaded file: {str(e)}")
    elif seq_text.strip():
        fasta_str = seq_text
    else:
        return error("Please provide a sequence via text or file.")

    try:
        records = parse_fasta_string(fasta_str)
    except Exception as e:
        return error(f"FASTA parsing failed: {str(e)}")

    if not records:
        return error("No valid sequences found in the input (no FASTA records detected). "
                     "Make sure each entry starts with a '>' header line.")

    # --- Validate each sequence before inference ---
    for record in records:
        seq = record["sequence"].upper()
        if len(seq) < 20:
            return error(f"Sequence '{record['id']}' is too short ({len(seq)} amino acids). Minimum length is 20.")
        if len(seq) > 2000:
            return error(f"Sequence '{record['id']}' is too long ({len(seq)} amino acids). Maximum length is 2000.")
        invalid_chars = set(seq) - VALID_AMINO_ACIDS
        if invalid_chars:
            return error(f"Invalid characters in sequence '{record['id']}': {''.join(sorted(invalid_chars))}. "
                         f"Only standard amino acids (ACDEFGHIKLMNPQRSTVWY) are allowed.")

    # --- Model inference (explicit try/except to surface meaningful error) ---
    results = []
    for i, record in enumerate(records, 1):
        seq = record["sequence"].upper()
        try:
            inputs = tokenizer(seq, return_tensors="pt", max_length=512, truncation=True, padding=True)
            inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = model(**inputs)
                prediction = outputs.logits.squeeze().item()
            results.append({"id": record["id"], "tm": round(prediction, 2)})
        except torch.cuda.OutOfMemoryError:
            return error(f"GPU out-of-memory error while processing sequence '{record['id']}' "
                         f"(length {len(seq)}). Try on CPU or reduce input size.")
        except Exception as e:
            return error(f"Model inference failed for sequence '{record['id']}': {str(e)}. "
                         f"This may indicate an incompatibility between the model and the input.")

    # Sort once here; the rank order is threshold-independent, so the cached results
    # can be re-filtered on threshold changes without re-sorting.
    results.sort(key=lambda x: x["tm"], reverse=True)

    # --- Build output ---
    try:
        outputs = render_results(results, threshold)
    except Exception as e:
        return error(f"Failed to generate output: {str(e)}")

    # Cache the sorted predictions so changing the threshold can re-filter without inference.
    return outputs + (results,)


# Override the theme's grey label / helper text directly (more reliable than CSS),
# then use CSS only for the file-upload box sizing.
theme = gr.themes.Origin().set(
    body_text_color="#1f2937",
    body_text_color_subdued="#374151",
    block_label_text_color="#1f2937",
    block_label_text_weight="700",
    block_label_text_size="1rem",
    block_title_text_color="#1f2937",
    block_title_text_weight="700",
    block_title_text_size="1rem",
    block_info_text_color="#374151",
)

# Label / info text colors live in the theme tokens above. CSS is only for what no
# theme token can express: the file-upload box size and its (inherited) inner font.
CUSTOM_CSS = """
#fasta_upload .wrap {
    min-height: 110px !important;
    font-size: 0.8rem !important;
}
"""

demo = gr.Blocks(theme=theme, css=CUSTOM_CSS)
with demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Image("assets/TmProt_logo.png", width=100, height=100, show_label=False, show_download_button=False, container=False, show_share_button=False, show_fullscreen_button=False, interactive=False)
        with gr.Column(scale=7):
            gr.Markdown("""
                        # TmProt
                        ## Protein Thermostability Predictor
                        """)
            gr.Markdown(value="""
                        ### TmProt is a machine-learning-based protein thermostability predictor that leverages a fine-tuned ESM-2 protein language model to estimate melting temperatures (Tm) of protein sequences. It enables users to upload protein sequences in FASTA format (either pasted as text or uploaded as a file), and outputs predicted Tm values ranked by a user-defined thermostability threshold.
                        
                        📄 **Paper:** [https://doi.org/10.64898/2026.05.07.723192](https://doi.org/10.64898/2026.05.07.723192)<br>
                        🐙 **GitHub:** [https://github.com/loschmidt/TmProt](https://github.com/loschmidt/TmProt)
                        """
                        )
            
        with gr.Column(scale=1):
            gr.Image("assets/logo.png", width=100, height=100, show_label=False, show_download_button=False, container=False, show_share_button=False, show_fullscreen_button=False, interactive=False)
    with gr.Row():
        with gr.Column(scale=4):
            seq_text = gr.Textbox(
                label="FASTA sequences",
                lines=6,
                placeholder=">seq\nMKTIIALSYIFCLVFA",
                value="",
            )
            seq_file = gr.File(label="Or upload FASTA file", file_types=[".fasta", ".fa", ".txt"], type="filepath", height=120, elem_id="fasta_upload")
            btn = gr.Button("Predict")

            def toggle_input_visibility(txt, f):
                if txt.strip():
                    return gr.update(visible=False)
                elif f is not None:
                    return gr.update(visible=False)
                else:
                    return gr.update(visible=True)

            # Text input and file upload share the space: using one hides the other,
            # and clearing it brings the alternative back.
            seq_text.input(
                lambda txt: toggle_input_visibility(txt, None),
                inputs=seq_text,
                outputs=seq_file,
            )
            seq_file.change(
                lambda f: toggle_input_visibility("", f),
                inputs=seq_file,
                outputs=seq_text,
            )
            cutoff_bins = [str(x) for x in range(20, 101, 10)]
            cutoff_bar = gr.Radio(
                choices=cutoff_bins,
                label="Thermostability threshold (°C)",
                info="Only affects the Yes/No label — Tm predictions are independent of this setting. Changing it re-filters existing results instantly, with no need to re-run the model.",
                value="60"
            )
        with gr.Column(scale=4):
            error_box = gr.Markdown(visible=False)
            output = gr.Dataframe(headers=["Rank", "ID", "Predicted Tm [°C]", "Thermostable"], label="Results")
            download_btn = gr.DownloadButton(label="Download CSV")
    # Holds the raw per-sequence predictions from the last run, so changing the
    # threshold can re-filter instantly without re-running the model.
    results_state = gr.State([])
    btn.click(
        predict_tm,
        inputs=[seq_text, seq_file, cutoff_bar],
        outputs=[output, error_box, results_state]
    )
    # Changing the threshold only re-labels existing results — no inference. Shares the
    # full result-area contract (incl. error_box) so a stale error banner is cleared too.
    cutoff_bar.change(
        render_results,
        inputs=[results_state, cutoff_bar],
        outputs=[output, error_box]
    )
    # Generate CSV on-demand when user clicks download, using current results and threshold.
    download_btn.click(
        generate_csv,
        inputs=[results_state, cutoff_bar],
        outputs=download_btn
    )
    with gr.Row():
        examples = gr.Examples(
            examples = [
            [""">I1W5V5
MSIENLSSNKSFGGWHKQYSHVSNTLNCAMRFAIYLPPQASTGAKVPVLYWLSGLTCSDENFMQKAGAQRLAAELGIAIVAPDTSPRGEGVADDEGYDLGQGAGFYVNATQAPWNRHYQMYDYVVNELPELIESMFPVSDKRAIAGHSMGGHGALTIALRNPERYQSVSAFSPINNPVNCPWGQKAFTAYLGKDTDTWREYDASLLMRAAKQYVPALVDQGEADNFLAEQLKPEVLEAAASSNNYPLELRSHEGYDHSYYFIASFIEDHLRFHSNYLNA""",
                        None,           # seq_file is None (we use text)
                        "60"            # threshold
                    ],
                    [""">R4YJ85
MINLEKALAGRRILIVDDLVEARSSLKKMATILGGDNIDVATDGIEAMSLIHEHEYDIVLSDYNLGRTKDGQQILEEARFTQRLRATSLFIVITGENAIDMVMGALEYDPDGYITKPYTLNMLKERLIRIITIKEELRKVNKAIDLQKYDLAIKYCLEVLDSNPRLRLPASRILGQLLMRQKRFQQALKIYSQLLNERSVSWAKLGQAICIFKLGDPNSALALLNRALVDHPLYVQCYDWIAKILLTLDKPLEAQAALEKAIVISPKAVLRQMELGRIAYENGDMVTAEPAFKYSVRLGRFSCHKSAKNYLQFVRSAQALLINPKERQTQNKANEAFRALTELKQDFSDDKDSLFEASIVESKTHLKMENLDEAKRSANDAEDMLAKLECPKIDYKLQMTETFIETDQSVKAQKMIDELKSAELSDKQIIMLNRLDNDLNGEALKRHSTSLNDQGVSHYEKGELEEAIIAFDQATHYEQAGISVLLNSIQAKISLMERDSPDKKILKNVRSLLIRIGEIAKDDERFARYSRLRKTYDRLCRAAAK""",
                        None,
                        "50"
                    ],
            ],
            inputs=[seq_text, seq_file, cutoff_bar],
            label="Click an example to try TmProt instantly",
            examples_per_page=2,
        )
        # An example fills the text box (and leaves the file empty), so restore the
        # default split: text box visible, file upload hidden.
        examples.dataset.click(
            lambda: gr.update(visible=False),
            outputs=seq_file,
        )
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(value="""
                ## Features

                    - Predict protein melting temperature (Tm) from amino acid sequences

                    - Accepts input via FASTA text or FASTA file upload

                    - Supports sequences from 20 to 2000 amino acids in length

                    - Outputs a ranked table with predicted Tm and thermostability status based on user-chosen threshold

                    - CSV download option for easy export and downstream analysis
            """
            )
        with gr.Column(scale=1):
            gr.Markdown(value="""
                ## Model Overview

                    - Base Model: facebook/esm2_t33_650M_UR50D (650M parameters)

                    - Fine-tuning method: LoRA (Low-Rank Adaptation) using PEFT framework

                    - Task: Regression prediction of protein melting temperature (Tm)

                    - Training Data: ProMelt dataset (merged Meltome Atlas + ProTherm) with ~45,000 protein sequences and experimental Tm values

                    - Output: Single linear regression output neuron predicting Tm in °C
            """
            )

if __name__ == "__main__":
    demo.launch(share=True, allowed_paths=['./assets'])