ui_layout.py

"""
ui_layout.py
============
Gradio component layout only.
Imports styles from ui_styles.py and pipeline from backend.py.
Returns build_demo() — called by app.py.
"""

import gradio as gr
from backend import run_pipeline
from ui_styles import HEADER_HTML, FOOTER_HTML, sec

EXAMPLES = [
    ["17", "43071077",  "G", "A", 1, 0],
    ["11", "5226929",   "T", "C", 1, 0],
    ["7",  "117548628", "T", "A", 1, 0],
    ["3",  "37053577",  "A", "C", 0, 1],
    ["19", "44908684",  "G", "T", 1, 0],
]


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="Mutation XAI — Explainability Intelligence System") as demo:

        gr.HTML(HEADER_HTML)

        with gr.Row(equal_height=False):

            # ── LEFT: inputs ──────────────────────────────────────────────────
            with gr.Column(scale=1, min_width=280):

                gr.HTML(sec("IN", "Variant Coordinates"))

                chrom_in = gr.Textbox(label="Chromosome", value="17",
                                      max_lines=1, placeholder="e.g. 17")
                pos_in   = gr.Textbox(label="Position  (hg38 · 1-based)",
                                      value="43071077", max_lines=1)
                with gr.Row():
                    ref_in = gr.Textbox(label="Ref Base", value="G",
                                        max_lines=1, placeholder="A/C/G/T")
                    alt_in = gr.Textbox(label="Alt Base", value="A",
                                        max_lines=1, placeholder="A/C/G/T")
                with gr.Row():
                    exon_in   = gr.Radio([0, 1], label="Exon flag",   value=1)
                    intron_in = gr.Radio([0, 1], label="Intron flag", value=0)

                run_btn = gr.Button("▶  Analyse Variant", variant="primary",
                                    size="lg", elem_classes="run-btn")

                gr.HTML(sec("EX", "Example Variants", "1.4rem"))
                gr.Examples(
                    examples=EXAMPLES,
                    inputs=[chrom_in, pos_in, ref_in, alt_in, exon_in, intron_in],
                    label="", examples_per_page=5,
                )

            # ── RIGHT: outputs ────────────────────────────────────────────────
            with gr.Column(scale=3, min_width=580):

                # ① Explanation — always first
                gr.HTML(sec("01", "Explanation & Signal Summary"))
                summary_out = gr.Markdown(
                    value="*Submit a variant to see the full explanation.*",
                    elem_classes="expl-box",
                )
                final_exp_out = gr.Textbox(
                    label="Final Explanation — grounded in internal model signals",
                    lines=8, max_lines=16, interactive=False,
                )

                # ② XAI metrics
                gr.HTML(sec("02", "Cross-Model Explainability Metrics", "1.4rem"))
                xai_plot = gr.Image(label="XAI Metrics")

                # ③ Per-model signals
                gr.HTML(sec("03", "Internal Model Signals", "1.4rem"))
                with gr.Tabs():

                    with gr.TabItem("🔬  Splice Model"):
                        spl_act_out  = gr.Image(label="conv3 Activation — Splice")
                        spl_risk_out = gr.Image(label="Splice Distance Risk Heatmap")
                        spl_grad_out = gr.Image(label="Gradient Attribution — Splice")

                    with gr.TabItem("🧬  V4 Model"):
                        v4_act_out  = gr.Image(label="conv3 Activation — V4")
                        v4_grad_out = gr.Image(label="Gradient Attribution — V4")

                    with gr.TabItem("📊  Classic Model"):
                        cls_act_out = gr.Image(label="conv3 Activation — Classic")

                    with gr.TabItem("⚗️  Causal Analysis"):
                        cf_out  = gr.Image(label="Counterfactual Analysis")
                        abl_out = gr.Image(label="Feature Ablation Chart")

                    with gr.TabItem("📋  JSON Report"):
                        json_out = gr.Textbox(label="Structured JSON Report",
                                              lines=26, max_lines=60,
                                              interactive=False)
                        dl_out = gr.File(label="⬇  Download JSON")

        gr.HTML(FOOTER_HTML)

        # ── Event wiring ──────────────────────────────────────────────────────
        run_btn.click(
            fn=run_pipeline,
            inputs=[chrom_in, pos_in, ref_in, alt_in, exon_in, intron_in],
            outputs=[
                summary_out, final_exp_out,
                xai_plot,
                spl_act_out, spl_risk_out, spl_grad_out,
                v4_act_out,  v4_grad_out,
                cls_act_out,
                cf_out, abl_out,
                json_out, dl_out,
            ],
            show_progress="full",
        )

    return demo