"""
app.py — Meta-LoRA Molecular Generator
Nature-inspired light UI · Gradio 5
"""

import gradio as gr
import pandas as pd
from inference import run_generation, load_models, mol_to_pil

load_models()

# ── Valid scaffold examples (verified with RDKit) ──────────────────────────────
EXAMPLES = {
    " Caffeine-like (xanthines)": "\n".join([
        "Cn1cnc2c1c(=O)n(C)c(=O)n2C",
        "Cn1cnc2[nH]c(=O)n(C)c2c1=O",
        "Cn1cnc2c1c(=O)[nH]c(=O)n2C",
        "O=c1[nH]c(=O)c2[nH]cnc2[nH]1",
        "CCn1cnc2c1c(=O)n(C)c(=O)n2C",
    ]),
    " Aspirin-like (salicylates)": "\n".join([
        "CC(=O)Oc1ccccc1C(=O)O",
        "CC(=O)Oc1ccc(C)cc1C(=O)O",
        "CC(=O)Oc1cccc(C(=O)O)c1",
        "CC(=O)Oc1ccc(F)cc1C(=O)O",
        "CC(=O)Oc1ccc(Cl)cc1C(=O)O",
    ]),
    " Ibuprofen-like (arylpropionic)": "\n".join([
        "CC(C)Cc1ccc(C(C)C(=O)O)cc1",
        "CC(C)Cc1ccc(C(C)C(=O)OC)cc1",
        "CC(C(=O)O)c1ccc(Cl)cc1",
        "CC(C(=O)O)c1cccc(F)c1",
        "CC(C(=O)O)c1ccc(CC)cc1",
    ]),
}

# ── Nature-inspired CSS ────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Lora:ital,wght@0,400;0,600;1,400&family=DM+Sans:wght@300;400;500&display=swap');

:root {
    --sage:       #7C9E7E;
    --sage-light: #B8D4B9;
    --sage-pale:  #EBF3EB;
    --mint:       #C8E6C9;
    --forest:     #3D6B40;
    --earth:      #8B6F47;
    --clay:       #C4956A;
    --cream:      #FAFAF7;
    --warm-white: #F5F2ED;
    --text-dark:  #2C3E2D;
    --text-mid:   #5A7A5C;
    --text-soft:  #8BA88C;
    --shadow-sm:  0 2px 12px rgba(61,107,64,0.10);
    --shadow-md:  0 6px 28px rgba(61,107,64,0.14);
    --radius:     16px;
}

/* ── Global ── */
body, .gradio-container {
    background: var(--cream) !important;
    font-family: 'DM Sans', sans-serif !important;
    color: var(--text-dark) !important;
}

/* Subtle leaf-pattern background */
.gradio-container::before {
    content: '';
    position: fixed;
    inset: 0;
    background-image:
        radial-gradient(circle at 15% 20%, rgba(124,158,126,0.07) 0%, transparent 50%),
        radial-gradient(circle at 85% 75%, rgba(200,230,201,0.12) 0%, transparent 50%),
        radial-gradient(circle at 50% 50%, rgba(235,243,235,0.08) 0%, transparent 70%);
    pointer-events: none;
    z-index: 0;
}

/* ── Header ── */
.app-header {
    background: linear-gradient(135deg, var(--forest) 0%, var(--sage) 60%, #9EC4A0 100%);
    border-radius: var(--radius);
    padding: 36px 40px 32px;
    margin-bottom: 8px;
    position: relative;
    overflow: hidden;
    box-shadow: var(--shadow-md);
}
.app-header::after {
    content: '⬡ ⬡ ⬡';
    position: absolute;
    right: 32px; top: 20px;
    font-size: 48px;
    opacity: 0.08;
    letter-spacing: 8px;
    color: white;
}
.app-header h1 {
    font-family: 'Lora', serif !important;
    font-size: 2rem !important;
    font-weight: 600 !important;
    color: white !important;
    margin: 0 0 8px !important;
    line-height: 1.2 !important;
}
.app-header p {
    color: rgba(255,255,255,0.85) !important;
    font-size: 0.95rem !important;
    margin: 0 !important;
    font-weight: 300 !important;
    max-width: 560px !important;
}

/* ── Panels ── */
.panel-card {
    background: white !important;
    border-radius: var(--radius) !important;
    padding: 28px !important;
    box-shadow: var(--shadow-sm) !important;
    border: 1px solid rgba(124,158,126,0.15) !important;
}

/* ── Labels ── */
label span, .label-wrap span {
    font-family: 'DM Sans', sans-serif !important;
    font-weight: 500 !important;
    font-size: 0.82rem !important;
    color: var(--text-mid) !important;
    text-transform: uppercase !important;
    letter-spacing: 0.06em !important;
}

/* ── Dropdown ── */
.gr-dropdown, select {
    border: 1.5px solid var(--sage-light) !important;
    border-radius: 10px !important;
    background: var(--sage-pale) !important;
    color: var(--text-dark) !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 0.9rem !important;
    transition: border-color 0.2s !important;
}
.gr-dropdown:focus, select:focus {
    border-color: var(--sage) !important;
    outline: none !important;
    box-shadow: 0 0 0 3px rgba(124,158,126,0.18) !important;
}

/* ── Textbox ── */
textarea, .gr-textbox textarea {
    border: 1.5px solid var(--sage-light) !important;
    border-radius: 12px !important;
    background: var(--sage-pale) !important;
    color: var(--text-dark) !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 0.88rem !important;
    line-height: 1.7 !important;
    padding: 14px !important;
    transition: border-color 0.2s, box-shadow 0.2s !important;
    resize: vertical !important;
}
textarea:focus {
    border-color: var(--sage) !important;
    box-shadow: 0 0 0 3px rgba(124,158,126,0.18) !important;
    outline: none !important;
    background: white !important;
}

/* ── Slider ── */
input[type=range] {
    accent-color: var(--sage) !important;
}
.slider-container { background: transparent !important; }

/* ── Generate button ── */
#gen-btn {
    background: linear-gradient(135deg, var(--forest), var(--sage)) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 1rem !important;
    font-weight: 500 !important;
    padding: 14px 28px !important;
    letter-spacing: 0.03em !important;
    cursor: pointer !important;
    box-shadow: 0 4px 16px rgba(61,107,64,0.28) !important;
    transition: transform 0.15s, box-shadow 0.15s !important;
    width: 100% !important;
}
#gen-btn:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 8px 24px rgba(61,107,64,0.36) !important;
}
#gen-btn:active { transform: translateY(0) !important; }

/* ── Error/status message ── */
#error-msg p {
    color: #C0392B !important;
    background: #FFF0EE !important;
    border: 1px solid #F5C6C2 !important;
    border-radius: 8px !important;
    padding: 10px 14px !important;
    font-size: 0.85rem !important;
}

/* ── Metrics card ── */
#metrics-box {
    background: linear-gradient(135deg, var(--sage-pale), #F0FAF0) !important;
    border: 1.5px solid var(--sage-light) !important;
    border-radius: var(--radius) !important;
    padding: 22px !important;
}
#metrics-box table {
    width: 100% !important;
    border-collapse: collapse !important;
}
#metrics-box th {
    color: var(--text-soft) !important;
    font-size: 0.78rem !important;
    text-transform: uppercase !important;
    letter-spacing: 0.06em !important;
    padding: 6px 10px !important;
    text-align: left !important;
}
#metrics-box td {
    padding: 7px 10px !important;
    font-size: 0.92rem !important;
    color: var(--text-dark) !important;
    border-bottom: 1px solid rgba(124,158,126,0.12) !important;
}
#metrics-box td:last-child { color: var(--forest) !important; font-weight: 600 !important; }

/* ── Gallery ── */
.gallery-item, .thumbnail-item {
    border-radius: 12px !important;
    overflow: hidden !important;
    border: 1.5px solid var(--sage-light) !important;
    box-shadow: var(--shadow-sm) !important;
    background: white !important;
    transition: transform 0.15s, box-shadow 0.15s !important;
}
.gallery-item:hover { transform: translateY(-2px) !important; box-shadow: var(--shadow-md) !important; }

/* ── Table ── */
.gr-dataframe table { border-collapse: separate !important; border-spacing: 0 4px !important; }
.gr-dataframe th {
    background: var(--sage-pale) !important;
    color: var(--text-mid) !important;
    font-size: 0.78rem !important;
    text-transform: uppercase !important;
    letter-spacing: 0.05em !important;
    padding: 10px 14px !important;
    border: none !important;
}
.gr-dataframe td {
    background: white !important;
    padding: 9px 14px !important;
    border-bottom: 1px solid var(--sage-pale) !important;
    font-size: 0.87rem !important;
    color: var(--text-dark) !important;
}

/* ── Section headings ── */
.section-title {
    font-family: 'Lora', serif !important;
    font-size: 1rem !important;
    font-weight: 600 !important;
    color: var(--forest) !important;
    margin-bottom: 14px !important;
    display: flex !important;
    align-items: center !important;
    gap: 8px !important;
}

/* ── Footer ── */
.app-footer {
    text-align: center !important;
    padding: 20px !important;
    color: var(--text-soft) !important;
    font-size: 0.78rem !important;
    border-top: 1px solid var(--sage-light) !important;
    margin-top: 8px !important;
}

/* ── Pill badge ── */
.pill {
    display: inline-block;
    background: var(--sage-pale);
    border: 1px solid var(--sage-light);
    color: var(--forest);
    border-radius: 20px;
    padding: 2px 12px;
    font-size: 0.75rem;
    font-weight: 500;
    margin: 2px;
}
"""

def load_example(choice):
    return EXAMPLES.get(choice, "")

def validate_smiles_input(smiles_text):
    from rdkit import Chem
    lines = [s.strip() for s in smiles_text.strip().split('\n') if s.strip()]
    if len(lines) < 3:
        return None, "❌ Enter at least 3 SMILES (one per line)."
    if len(lines) > 10:
        return None, "❌ Maximum 10 SMILES in the support set."
    valid = []
    for smi in lines:
        if Chem.MolFromSmiles(smi) is None:
            return None, f"❌ Invalid SMILES: `{smi}`"
        valid.append(smi)
    return valid, None

def generate_molecules(smiles_text, n_generate, progress=gr.Progress()):
    valid_smiles, err = validate_smiles_input(smiles_text)
    if err:
        return None, err, None, None

    TEMPERATURE = 0.8   # fixed — optimal balance of diversity and validity

    progress(0.15, desc="Encoding support set…")
    try:
        progress(0.40, desc="Generating molecules…")
        results = run_generation(valid_smiles, n=int(n_generate), temperature=TEMPERATURE)
    except Exception as e:
        return None, f"❌ Generation failed: {str(e)}", None, None

    progress(0.85, desc="Computing metrics…")
    m = results

    summary_md = f"""
<div id="metrics-box">

| Metric | Value |
|---|---|
|  Generated | {m['n_generated']} |
|  Valid | {m['n_valid']} &nbsp;({m['validity']:.1f}%) |
|  Unique | {m['n_unique']} &nbsp;({m['uniqueness']:.1f}%) |
|  Novel | {m['n_novel']} &nbsp;({m['novelty']:.1f}%) |
|  Avg Tanimoto | {m['avg_tanimoto']:.4f} |

</div>
"""

    rows = []
    for item in results['images']:
        props = item.get('props') or {}
        rows.append({
            "SMILES":  item['smiles'],
            "QED":     props.get("qed", "—"),
            "LogP":    props.get("logp", "—"),
            "MW":      props.get("mw", "—"),
            "HBD":     props.get("hbd", "—"),
            "HBA":     props.get("hba", "—"),
        })

    df      = pd.DataFrame(rows) if rows else pd.DataFrame(columns=["SMILES","QED","LogP","MW","HBD","HBA"])
    gallery = [item['image'] for item in results['images']]

    progress(1.0, desc="Done!")
    return summary_md, "", gallery, df


# ── Layout ────────────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, title="Meta-LoRA Molecular Generator") as demo:

    # Header
    gr.HTML("""
<div class="app-header">
    <div style="display:flex; align-items:center; gap:20px; margin-bottom:12px">
        
        <h1 style="margin:0!important">Meta-LoRA Molecular Generator</h1>
    </div>
    """)

    with gr.Row(equal_height=False):

        # ── Left panel ──
        with gr.Column(scale=1, min_width=320):
            gr.HTML('<div class="section-title"> Support Set</div>')

            example_dropdown = gr.Dropdown(
                choices=list(EXAMPLES.keys()),
                label="Load an example scaffold family",
                value=None,
            )

            smiles_input = gr.Textbox(
                label="SMILES strings · one per line · 3 – 10 molecules",
                placeholder="Paste SMILES here…\nCn1cnc2c1c(=O)n(C)c(=O)n2C\n…",
                lines=8,
            )

            n_generate = gr.Slider(
                minimum=1, maximum=10, value=5, step=1,
                label="Molecules to generate",
            )

            generate_btn = gr.Button(" Generate Molecules", elem_id="gen-btn", variant="primary")
            error_box    = gr.Markdown(elem_id="error-msg")

        # ── Right panel ──
        with gr.Column(scale=2):
            gr.HTML('<div class="section-title"> Results</div>')
            metrics_md  = gr.Markdown(elem_id="metrics-box")
            mol_gallery = gr.Gallery(
                label="Generated molecules (novel · valid)",
                columns=3,
                height="auto",
                object_fit="contain",
            )

    gr.HTML('<div class="section-title" style="margin-top:20px"> Molecule Properties</div>')
    props_table = gr.Dataframe(
        headers=["SMILES","QED","LogP","MW","HBD","HBA"],
        interactive=False,
        wrap=True,
    )

    gr.HTML("""
    <div class="app-footer">
        Scaffold-Episodic Meta-Learning · Context-Conditioned LoRA (rank 16) · ZINC250k ·
        Validity ~96.8% · Uniqueness ~99% · Novelty ~98% · Tanimoto ~0.42
    </div>
    """)

    # Events
    example_dropdown.change(fn=load_example, inputs=example_dropdown, outputs=smiles_input)
    generate_btn.click(
        fn=generate_molecules,
        inputs=[smiles_input, n_generate],
        outputs=[metrics_md, error_box, mol_gallery, props_table],
    )

demo.launch()