"""SugarKi public UI — sugar-chemistry Ki predictor (frontend, MWBC/sugarki).

Pretty Gradio interface that calls the SugarKi backend Space (private,
ZeroGPU). No model weights here — just the UI.
"""
from __future__ import annotations

import os
from gradio_client import Client
import gradio as gr

BACKEND = os.environ.get("SUGARKI_BACKEND", "Umesh1608/sugarki-backend")
HF_TOKEN = os.environ.get("HF_TOKEN")

if not HF_TOKEN:
    print("WARNING: HF_TOKEN not set; backend calls will fail if backend is private.")

# ----------------------------------------------------------------------
# Curated example library
# ----------------------------------------------------------------------

SUGAR_INHIBITORS = {
    "D-mannitol":            "OCC(O)C(O)C(O)C(O)CO",
    "D-fructose":            "OCC(=O)C(O)C(O)C(O)CO",
    "D-sorbitol":            "OCC(O)C(O)C(O)C(O)CO",
    "D-glucose":             "OCC(O)C(O)C(O)C(O)C=O",
    "D-galactose":           "OCC(O)C(O)C(O)C(O)C=O",
    "D-mannose":             "OCC(O)C(O)C(O)C(O)C=O",
    "Xylitol":               "OCC(O)C(O)C(O)CO",
    "Glycerol":              "OCC(O)CO",
    "Glucose-6-phosphate":   "OC(=O)CC(O)C(O)C(O)COP(O)(O)=O",
    "Fructose-1,6-bisphos.": "O=P(O)(O)OCC(O)(C(O)C(O)COP(O)(O)=O)O",
    "Sucrose":               "OCC1OC(OC2(COC1O)OC(CO)C(O)C2O)C(O)C(O)C1O",
    "Trehalose":             "OCC1OC(OC2OC(CO)C(O)C(O)C2O)C(O)C(O)C1O",
    "Custom (paste SMILES)": "",
}

EXAMPLE_MDH_WT = (
    "MGSSHHHHHHSSGLVPRGSHMVKLTLSALPALSPAVAVPAYDPRAQIPGIVHFGVGAFHRSHQAMYLDRL"
    "LNSGRGAGWAICGVGVLPQDARMRDVLAEQDHLYTLVTRSPDGQAQARVIGAIVEFLFAPDDPERVLERL"
    "ADPTTRIVSLTVTEGGYSVSNATGEFDPTPPDIAHDLTPGAVPRTFFGFLTEGLRRRRERGLPPFTVVSC"
    "DNMPGNGEVTRRALTAFARLQDPELGDWIAHNVAFPNSMVDRITPATTEQDRQDIAAAYGIEDAWPVVAE"
    "SFAQWVLEDRFTQGRPALETVGVQVVSDVEPYELMKLRLLNASHQALAYLGLLAGYRFVHEVCQDPLFAR"
    "FLLDYMTQEATPTLRPVPGIDLGAYRRELIARFSNPAIRDPLTRLTVDSSERIPKFLLPVIRDQLARGGE"
    "LARCALVIASWRAYLATVLEEGSASFPDQHAQALAEAVRRDAQQPGAFLDLEAVFGELGRNARFRTAYLS"
    "AWESLRRQGPLGAMRALMGEESSPSNVTSLSGR"
)

EXAMPLES = [
    [EXAMPLE_MDH_WT, "D-mannitol", "substrate", "OCC(=O)C(O)C(O)C(O)CO"],
]

# ----------------------------------------------------------------------
# Backend client
# ----------------------------------------------------------------------

_client = None
def _backend():
    global _client
    if _client is None:
        kwargs = {}
        if HF_TOKEN:
            kwargs["hf_token"] = HF_TOKEN
        _client = Client(BACKEND, **kwargs)
    return _client


def predict(sequence, inhibitor_choice, inhibitor_custom, inh_type, substrate_smiles):
    if not sequence or not sequence.strip():
        return "### ⚠️ Please paste an enzyme sequence", None, "", None

    smiles = (
        inhibitor_custom.strip()
        if inhibitor_choice == "Custom (paste SMILES)"
        else SUGAR_INHIBITORS.get(inhibitor_choice, "")
    )
    if not smiles:
        return "### ⚠️ Inhibitor SMILES is empty", None, "", None

    try:
        result = _backend().predict(
            sequence.strip(), smiles, inh_type,
            substrate_smiles.strip() if substrate_smiles else "",
            api_name="/predict_ki",
        )
    except Exception as e:
        return f"### ❌ Backend error\n```\n{e}\n```", None, "", None

    if isinstance(result, dict) and result.get("error"):
        msg = result.get("error_message", str(result))
        return f"### ❌ Prediction error\n```\n{msg}\n```", None, "", result

    # Headline card
    ki_mm = result.get("Ki_mM")
    ki_um = result.get("Ki_uM")
    log_ki = result.get("log10_Ki_mM")
    inh_used = result.get("inh_type_used", "?")
    his_stripped = "✓ stripped" if result.get("his_tag_stripped") else "—"
    sigma = result.get("ensemble_std_log10", "?")

    # Format Ki nicely with appropriate units
    if ki_mm is not None:
        if ki_mm < 0.001:
            ki_display = f"**{ki_mm * 1e6:.2f} nM**"
        elif ki_mm < 1:
            ki_display = f"**{ki_um:.1f} µM**"
        else:
            ki_display = f"**{ki_mm:.3f} mM**"
    else:
        ki_display = "—"

    summary_md = f"""
### 🧪 Predicted Ki: {ki_display}

| | |
|---|---|
| **Ki (mM)** | {ki_mm:.4f} |
| **Ki (µM)** | {ki_um:,.1f} |
| **log₁₀(Ki / mM)** | {log_ki:+.3f} |
| **Inhibition mode** | `{inh_used}` |
| **His-tag prefix** | {his_stripped} |
| **Ensemble σ across 5 modes** | {sigma} |

> **Model:** {result.get('model_version', 'SugarKi')}
"""

    # Per-mode table
    mode_mm = result.get("mode_predictions_mM", {})
    mode_log = result.get("mode_predictions_log10", {})

    mode_rows = []
    for mode in ["external", "product", "substrate", "product_or_substrate", "unknown"]:
        if mode in mode_mm:
            mm = mode_mm[mode]
            um = mm * 1000
            log = mode_log.get(mode, 0)
            highlight = " ⭐" if mode == inh_used else ""
            mode_rows.append([
                f"{mode.replace('_', ' ').title()}{highlight}",
                f"{log:+.3f}",
                f"{mm:.3f}",
                f"{um:,.1f}",
            ])

    notes = result.get("notes", "")
    notes_md = f"\n\n**ℹ️ Notes:** _{notes}_\n" if notes else ""

    return summary_md, mode_rows, notes_md, result


# ----------------------------------------------------------------------
# UI — uses theme-aware colors so text stays readable on any background
# ----------------------------------------------------------------------

CSS = """
/* Force a consistent DARK theme regardless of system preference.
   Dark backgrounds with light text — both set together so contrast is guaranteed. */

html, body, .gradio-container, .gradio-container .main {
    background: #0b1220 !important;
    color: #f3f4f6 !important;
}

/* All Gradio block-level containers — force dark backgrounds */
.gradio-container .block,
.gradio-container .form,
.gradio-container .panel,
.gradio-container .gr-box,
.gradio-container .gr-block,
.gradio-container .gr-form,
.gradio-container .gr-panel,
.gradio-container fieldset,
.gradio-container details,
.gradio-container summary,
.gradio-container .label-wrap,
.gradio-container .wrap,
.gradio-container .accordion,
.gradio-container .gr-accordion,
.gradio-container .prose,
.gradio-container .markdown,
.gradio-container .markdown-body {
    background-color: #111827 !important;
    color: #f3f4f6 !important;
    border-color: #374151 !important;
}

/* Inputs — slightly lighter slate so they stand out from the page */
.gradio-container input,
.gradio-container textarea,
.gradio-container select,
.gradio-container .gr-input,
.gradio-container .gr-textarea {
    background-color: #1f2937 !important;
    color: #f9fafb !important;
    border-color: #374151 !important;
}
.gradio-container input::placeholder,
.gradio-container textarea::placeholder {
    color: #9ca3af !important;
}

/* All text */
.gradio-container, .gradio-container * {
    color: #f3f4f6 !important;
}
.gradio-container h1,
.gradio-container h2,
.gradio-container h3,
.gradio-container h4,
.gradio-container h5,
.gradio-container h6,
.gradio-container strong,
.gradio-container b {
    color: #ffffff !important;
}

/* Tables */
.gradio-container table {
    background-color: #111827 !important;
    border-color: #374151 !important;
}
.gradio-container table th,
.gradio-container table td {
    color: #e5e7eb !important;
    background-color: #111827 !important;
    border-color: #374151 !important;
}
.gradio-container table th {
    background-color: #1f2937 !important;
    color: #ffffff !important;
}

/* Links and inline code — bright accent colors that pop on dark */
.gradio-container a { color: #93c5fd !important; }
.gradio-container code {
    background-color: #1f2937 !important;
    color: #fbcfe8 !important;
}

/* Primary buttons — keep vivid blue with white text */
.gradio-container button.primary,
.gradio-container button.lg.primary,
.gradio-container .primary,
.gradio-container .primary * {
    background: #2563eb !important;
    color: #ffffff !important;
    border-color: #1d4ed8 !important;
}

/* Title block layout */
.title-block { text-align: center; padding: 1em 0; margin-bottom: 0.5em; }

/* Ki result card */
.ki-card {
    background: #111827 !important;
    border: 1px solid #374151 !important;
    border-radius: 8px !important;
    padding: 1em 1.5em !important;
}
"""

with gr.Blocks(title="SugarKi — Ki prediction", theme=gr.themes.Default(), css=CSS) as demo:

    # ----- header -----
    gr.Markdown(
        """
        <div align='center'>

        # 🧪 SugarKi

        ##### Family-specialized Ki prediction for sugar-chemistry enzymes
        ##### polyol DHs · sugar kinases · glycosidases · phosphatases · aldolases · isomerases · phosphomutases

        </div>
        """
    )

    # ----- about + inputs explainer -----
    with gr.Accordion("📐 How does SugarKi work?", open=True):
        gr.Markdown(
            """
            SugarKi takes an **enzyme sequence** and an **inhibitor SMILES** and predicts the
            inhibition constant **Ki** specifically for sugar-chemistry enzymes — where current
            SOTA models like CatPred fail catastrophically (R² = −0.95 on monosaccharide
            inhibitors).

            **Inputs**

            - `Enzyme sequence` — paste any sugar-family enzyme (His-tag auto-stripped). Best for EC 1.1.1.x, 2.7.1.x, 3.1.3.x, 3.2.1.x, 4.1.2.x, 5.3.1.x, 5.4.2.x.
            - `Inhibitor SMILES` — SMILES of the small molecule whose Ki you want.
            - `Inhibition type` — choose `product` for product inhibition (e.g., mannitol on mannitol DH), `substrate` for substrate-mode at high [S], or `auto` for an ensemble across all 5 modes.
            - `Substrate SMILES` (optional) — native substrate, helps disambiguate product-inhibition cases.

            ### Validation

            - WT MDH-006 (mannitol DH) + D-mannitol → predicted **11.78 mM** in substrate mode vs literature **12 mM** (1.8% error).
            - Sugar-Ki test set R² = **0.702** vs CatPred 0.243 / SELFprot 0.623.
            - Typical MAE ≈ 0.6 log units; relative ranking is more reliable than absolute Ki values.
            """
        )

    # ----- input + output -----
    with gr.Row():
        with gr.Column(scale=3):
            seq_in = gr.Textbox(
                label="Enzyme amino acid sequence",
                placeholder="MGSSHHHHHH... or just the catalytic domain (30–1000 aa)",
                lines=10,
                value=EXAMPLE_MDH_WT,
                show_copy_button=True,
            )
            with gr.Row():
                inh_choice = gr.Dropdown(
                    choices=list(SUGAR_INHIBITORS.keys()),
                    value="D-mannitol",
                    label="Inhibitor",
                )
                inh_type = gr.Dropdown(
                    ["auto", "external", "product", "substrate",
                     "product_or_substrate", "unknown"],
                    value="auto",
                    label="Inhibition type",
                    info="`product` for product inhibition; `substrate` for substrate inhibition at high [S]",
                )
            inh_custom = gr.Textbox(
                label="Custom inhibitor SMILES (only used if inhibitor = 'Custom (paste SMILES)')",
                value="",
            )
            sub_smiles = gr.Textbox(
                label="Native substrate SMILES (optional, helps product-inhibition prediction)",
                value="",
            )
            btn = gr.Button("🧬 Predict Ki", variant="primary", size="lg")

        with gr.Column(scale=2):
            summary_out = gr.Markdown(elem_classes=["ki-card"])
            gr.Markdown(
                """
                **What each inhibition mode means**

                - **External** — small-molecule inhibitor unrelated to substrate or product (classical competitive/non-competitive ligand).
                - **Product** — the inhibitor IS the reaction product (product inhibition; e.g., mannitol on mannitol dehydrogenase).
                - **Substrate** — the inhibitor IS the native substrate; binding becomes inhibitory at high [S] (substrate inhibition).
                - **Product or Substrate** — chemically ambiguous: the molecule could play either role for this enzyme.
                - **Unknown** — mode not annotated; treat as a generic ensemble estimate.

                The ⭐ row is the mode you selected (or the auto-resolved default).
                """
            )
            mode_table = gr.Dataframe(
                headers=["Inhibition mode", "log₁₀(Ki/mM)", "Ki (mM)", "Ki (µM)"],
                label="Per-mode predictions (across 5 inhibition assumptions)",
                wrap=True,
                interactive=False,
            )
            notes_out = gr.Markdown()
            with gr.Accordion("Raw JSON response", open=False):
                raw_out = gr.JSON()

    btn.click(
        predict,
        inputs=[seq_in, inh_choice, inh_custom, inh_type, sub_smiles],
        outputs=[summary_out, mode_table, notes_out, raw_out],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[seq_in, inh_choice, inh_type, sub_smiles],
        label="Example: MDH-006 WT + D-mannitol (substrate-mode)",
    )

    # ----- footer -----
    gr.Markdown(
        """
        ---

        ### How SugarKi compares to existing Ki predictors

        | Model | Sugar-chemistry Ki | General Ki |
        |---|---|---|
        | CatPred zero-shot | R² = 0.243 (catastrophic on monosaccharides/polyols) | R² = 0.578 |
        | SELFprot zero-shot | R² = 0.623 | R² = 0.314 |
        | **SugarKi specialist** | **R² = 0.702** | (router falls back to CatPred) |

        ### Limitations

        - Test MAE ~0.6 log units; **rankings** more reliable than **absolute** values.
        - Best on EC families 1.1.1.x / 2.7.1.x / 3.2.1.x / 3.1.3.x / 4.1.2.x / 5.3.1.x / 5.4.2.x.
        - Allosteric inhibition not separately modeled.
        - First call cold-starts ~30 s while ESMFold loads; subsequent calls re-use cached structures.

        ### Citation

        Paper in preparation. Hosted by [MWBC](https://huggingface.co/MWBC), backed by
        a private SugarKi inference Space ([Umesh1608](https://huggingface.co/Umesh1608)).
        """
    )


if __name__ == "__main__":
    demo.launch()