ui/quest3.py

"""Quest 3: Genome Explorer — Circos-style polar ring chart and contig drill-down.

Features a circular genome overview showing variability hotspots across
contigs, plus a linear gene-track view when drilling into a single contig.
"""

import json
import gradio as gr
import plotly.graph_objects as go

from src.utils import PRECOMPUTED_DIR
from src.plot_config import COLORS, apply_template

# ---------------------------------------------------------------------------
# Load polar contig layout once at import time
# ---------------------------------------------------------------------------
_POLAR_PATH = PRECOMPUTED_DIR / "polar_contig_layout.json"
_POLAR_DATA: list | None = None
if _POLAR_PATH.exists():
    with open(_POLAR_PATH) as _f:
        _POLAR_DATA = json.load(_f)


# ---------------------------------------------------------------------------
# Standalone chart builder (callable from callbacks / app.py)
# ---------------------------------------------------------------------------

def build_circos_figure(polar_layout=None):
    """Build a Circos-style circular genome plot.

    Parameters
    ----------
    polar_layout : list or None
        List of contig dicts from ``polar_contig_layout.json``.  Falls back
        to the module-level ``_POLAR_DATA`` when *None*.
    """
    if polar_layout is None:
        polar_layout = _POLAR_DATA

    if not polar_layout:
        fig = go.Figure()
        fig.add_annotation(text="Polar layout data not available", showarrow=False)
        return fig

    fig = go.Figure()

    for contig in polar_layout:
        contig_id = contig["contig_id"]
        # Truncate contig name for readability
        short_name = (
            contig_id.split("|")[-2]
            if "|" in contig_id
            else contig_id[-15:]
        )

        n_bins = max(len(contig["bins"]), 1)
        bin_width = (contig["theta_end"] - contig["theta_start"]) / n_bins

        # Add bins as polar bars
        for b in contig["bins"]:
            score = b.get("variability_score", 0)
            # Color by variability: green -> amber -> red
            if score == 0:
                color = COLORS["core"]
            elif score < 2:
                color = COLORS["shell"]
            else:
                color = COLORS["cloud"]

            fig.add_trace(go.Barpolar(
                r=[b["total_genes"]],
                theta=[b["theta"]],
                width=[bin_width],
                marker_color=color,
                marker_line_color="rgba(255,255,255,0.3)",
                marker_line_width=0.5,
                opacity=0.85,
                hovertemplate=(
                    f"{short_name}<br>"
                    f"Genes: %{{r}}<br>"
                    f"Score: {score:.1f}"
                    "<extra></extra>"
                ),
                showlegend=False,
            ))

        # Add a contig label at the midpoint
        mid_theta = (contig["theta_start"] + contig["theta_end"]) / 2
        max_r = (
            max(b["total_genes"] for b in contig["bins"]) + 3
            if contig["bins"]
            else 5
        )
        fig.add_trace(go.Scatterpolar(
            r=[max_r],
            theta=[mid_theta],
            mode="text",
            text=[short_name],
            textfont=dict(size=8, color=COLORS["text_secondary"]),
            showlegend=False,
            hoverinfo="skip",
        ))

    fig.update_layout(
        polar=dict(
            bgcolor="rgba(0,0,0,0)",
            radialaxis=dict(showticklabels=False, showline=False, showgrid=False),
            angularaxis=dict(
                showticklabels=False,
                showline=False,
                showgrid=False,
                direction="clockwise",
            ),
        ),
        height=500,
        margin=dict(l=20, r=20, t=20, b=20),
        paper_bgcolor="rgba(0,0,0,0)",
        showlegend=False,
    )
    return fig


# ---------------------------------------------------------------------------
# Explanation header HTML
# ---------------------------------------------------------------------------

_EXPLANATION_HTML = (
    '<div style="padding:12px 20px;color:#424242;font-size:14px;line-height:1.6;">'
    "Where in the genome does genetic diversity concentrate? "
    "The outer ring shows variability &mdash; "
    '<span style="color:#C62828;font-weight:600;">red</span> regions are hotspots '
    "with rare or variable genes, "
    '<span style="color:#F9A825;font-weight:600;">amber</span> marks moderate variation, '
    "and "
    '<span style="color:#2E7D32;font-weight:600;">green</span> regions are conserved.'
    "</div>"
)


# ---------------------------------------------------------------------------
# Tab builder
# ---------------------------------------------------------------------------

def build_quest3(contig_choices: list[str]):
    """Build Quest 3 tab components.

    Returns a dict whose keys match the original wiring contract:
        tab, contig_dropdown, heatmap_plot, track_plot, region_table,
        region_gene_text
    """
    with gr.Tab("Genome Explorer", id="quest3") as tab:
        # --- A) Explanation header ---
        gr.HTML(value=_EXPLANATION_HTML)

        # --- B) Circos-style polar ring chart (replaces old heatmap hero) ---
        heatmap_plot = gr.Plot(
            label="Circular Genome Overview",
            value=build_circos_figure(),
        )

        # --- C) Contig drill-down ---
        gr.Markdown("### Contig Detail")
        contig_dropdown = gr.Dropdown(
            choices=contig_choices,
            label="Select contig (top contigs by gene count)",
            interactive=True,
        )

        track_plot = gr.Plot(label="Gene track (colored by class)", visible=False)

        region_table = gr.Dataframe(
            headers=["gene_id", "start", "end", "strand", "core_class", "freq_pct"],
            label="Genes in selected region",
            interactive=False,
        )

        region_gene_text = gr.Textbox(
            label="Selected gene from region",
            interactive=False,
            visible=False,
        )

    return {
        "tab": tab,
        "contig_dropdown": contig_dropdown,
        "heatmap_plot": heatmap_plot,
        "track_plot": track_plot,
        "region_table": region_table,
        "region_gene_text": region_gene_text,
    }