ui/final.py

"""Final tab / Chapter 6: Your Report — field report generation and export."""

import gradio as gr

from src.plot_config import COLORS


# =====================================================================
# Module-level HTML builders (callable from callbacks)
# =====================================================================

def build_report_html(state, data: dict) -> str:
    """Build a presentation-ready styled HTML field report.

    Parameters
    ----------
    state : AppState
        Current application state with selected_line, backpack, achievements.
    data : dict
        Loaded data dictionary with line_stats, embedding, similarity, etc.

    Returns
    -------
    str
        Self-contained HTML string with inline CSS for portability.
    """
    if state is None or not state.selected_line:
        return (
            '<div style="text-align:center;padding:40px;color:#757575;">'
            '<p style="font-size:16px;">No line selected yet.</p>'
            '<p style="font-size:13px;">Go back to Chapter 0 to choose a line first.</p>'
            "</div>"
        )

    line_id = state.selected_line
    line_stats = data["line_stats"]
    embedding = data["embedding"]
    similarity = data["similarity"]
    gene_freq = data["gene_freq"]
    pav = data.get("pav")

    # -- Gather stats --
    ls_row = line_stats[line_stats["line_id"] == line_id]
    country = ls_row.iloc[0]["country"] if len(ls_row) > 0 else "Unknown"
    genes_present = int(ls_row.iloc[0]["genes_present_count"]) if len(ls_row) > 0 else 0
    unique_genes = int(ls_row.iloc[0]["unique_genes_count"]) if len(ls_row) > 0 else 0

    emb_row = embedding[embedding["line_id"] == line_id]
    cluster_id = int(emb_row.iloc[0]["cluster_id"]) if len(emb_row) > 0 else -1

    # Nearest neighbors
    sim_rows = similarity[similarity["line_id"] == line_id].nlargest(3, "jaccard_score")
    neighbors = [
        (r["neighbor_line_id"], f"{r['jaccard_score']:.3f}")
        for _, r in sim_rows.iterrows()
    ]

    # Core/shell/cloud breakdown
    core_count = shell_count = cloud_count = 0
    if pav is not None and line_id in pav.columns:
        my_genes = set(pav.index[pav[line_id] == 1])
        my_freq = gene_freq[gene_freq["gene_id"].isin(my_genes)]
        core_count = int((my_freq["core_class"] == "core").sum())
        shell_count = int((my_freq["core_class"] == "shell").sum())
        cloud_count = int((my_freq["core_class"] == "cloud").sum())

    # Rare genes
    rare_rows = []
    if pav is not None and line_id in pav.columns:
        my_genes_list = pav.index[pav[line_id] == 1].tolist()
        rare = gene_freq[
            (gene_freq["gene_id"].isin(my_genes_list))
            & (gene_freq["freq_count"] <= 5)
        ].nsmallest(5, "freq_count")
        for _, r in rare.iterrows():
            rare_rows.append((r["gene_id"], int(r["freq_count"]), r["core_class"]))

    # Backpack
    backpack_rows = []
    for g in (state.backpack_genes or []):
        gf = gene_freq[gene_freq["gene_id"] == g]
        if len(gf) > 0:
            backpack_rows.append((g, gf.iloc[0]["core_class"], int(gf.iloc[0]["freq_count"])))
        else:
            backpack_rows.append((g, "unknown", 0))

    # -- Build HTML --
    # Metric cards row
    metrics = [
        (f"{genes_present:,}", "Genes Present", COLORS["core"]),
        (str(unique_genes), "Unique Genes", COLORS["accent"]),
        (str(cluster_id), "Cluster", COLORS["selected"]),
        (str(len(state.backpack_genes)), "Backpack Genes", COLORS["shell"]),
    ]
    metric_cards = ""
    for val, label, color in metrics:
        metric_cards += (
            f'<div style="flex:1;min-width:120px;background:#FFFFFF;border-radius:12px;'
            f"padding:20px;text-align:center;border-top:3px solid {color};"
            f'box-shadow:0 2px 8px rgba(0,0,0,0.06);">'
            f'<div style="font-size:28px;font-weight:700;color:#1A1A1A;">{val}</div>'
            f'<div style="font-size:11px;font-weight:600;text-transform:uppercase;'
            f'letter-spacing:1px;color:#757575;margin-top:4px;">{label}</div>'
            f"</div>"
        )

    # Neighbor rows
    neighbor_html = ""
    for name, score in neighbors:
        neighbor_html += (
            f'<div style="display:flex;justify-content:space-between;padding:8px 0;'
            f'border-bottom:1px solid #F0F0E8;">'
            f'<span style="font-weight:500;">{name}</span>'
            f'<span style="color:#757575;font-family:monospace;">{score}</span>'
            f"</div>"
        )

    # Composition bar
    total_csc = core_count + shell_count + cloud_count
    if total_csc > 0:
        core_pct = core_count / total_csc * 100
        shell_pct = shell_count / total_csc * 100
        cloud_pct = cloud_count / total_csc * 100
    else:
        core_pct = shell_pct = cloud_pct = 0

    composition_bar = (
        f'<div style="display:flex;height:24px;border-radius:6px;overflow:hidden;margin:12px 0;">'
        f'<div style="width:{core_pct:.1f}%;background:{COLORS["core"]};"></div>'
        f'<div style="width:{shell_pct:.1f}%;background:{COLORS["shell"]};"></div>'
        f'<div style="width:{cloud_pct:.1f}%;background:{COLORS["cloud"]};"></div>'
        f"</div>"
        f'<div style="display:flex;justify-content:space-between;font-size:12px;color:#757575;">'
        f'<span style="color:{COLORS["core"]};font-weight:600;">Core: {core_count:,}</span>'
        f'<span style="color:{COLORS["shell"]};font-weight:600;">Shell: {shell_count:,}</span>'
        f'<span style="color:{COLORS["cloud"]};font-weight:600;">Cloud: {cloud_count:,}</span>'
        f"</div>"
    )

    # Rare genes table
    rare_html = ""
    if rare_rows:
        rare_html += (
            '<table style="width:100%;border-collapse:collapse;font-size:13px;">'
            '<tr style="border-bottom:2px solid #E0E0E0;">'
            '<th style="text-align:left;padding:8px 4px;color:#757575;">Gene</th>'
            '<th style="text-align:center;padding:8px 4px;color:#757575;">Lines</th>'
            '<th style="text-align:center;padding:8px 4px;color:#757575;">Class</th>'
            "</tr>"
        )
        for gene, count, cls in rare_rows:
            badge_color = {"core": COLORS["core"], "shell": COLORS["shell"], "cloud": COLORS["cloud"]}.get(cls, "#9E9E9E")
            badge_text_color = "#FFFFFF" if cls != "shell" else "#333333"
            rare_html += (
                f'<tr style="border-bottom:1px solid #F0F0E8;">'
                f'<td style="padding:6px 4px;font-family:monospace;font-weight:500;">{gene}</td>'
                f'<td style="text-align:center;padding:6px 4px;">{count}</td>'
                f'<td style="text-align:center;padding:6px 4px;">'
                f'<span style="display:inline-block;padding:2px 10px;border-radius:12px;'
                f"font-size:11px;font-weight:600;background:{badge_color};color:{badge_text_color};"
                f'">{cls}</span></td></tr>'
            )
        rare_html += "</table>"
    else:
        rare_html = '<p style="color:#757575;font-size:13px;">No rare genes (&le;5 lines) found.</p>'

    # Backpack table
    backpack_html = ""
    if backpack_rows:
        backpack_html += (
            '<table style="width:100%;border-collapse:collapse;font-size:13px;">'
            '<tr style="border-bottom:2px solid #E0E0E0;">'
            '<th style="text-align:left;padding:8px 4px;color:#757575;">Gene</th>'
            '<th style="text-align:center;padding:8px 4px;color:#757575;">Class</th>'
            '<th style="text-align:center;padding:8px 4px;color:#757575;">Lines</th>'
            "</tr>"
        )
        for gene, cls, count in backpack_rows:
            badge_color = {"core": COLORS["core"], "shell": COLORS["shell"], "cloud": COLORS["cloud"]}.get(cls, "#9E9E9E")
            badge_text_color = "#FFFFFF" if cls != "shell" else "#333333"
            backpack_html += (
                f'<tr style="border-bottom:1px solid #F0F0E8;">'
                f'<td style="padding:6px 4px;font-family:monospace;font-weight:500;">{gene}</td>'
                f'<td style="text-align:center;padding:6px 4px;">'
                f'<span style="display:inline-block;padding:2px 10px;border-radius:12px;'
                f"font-size:11px;font-weight:600;background:{badge_color};color:{badge_text_color};"
                f'">{cls}</span></td>'
                f'<td style="text-align:center;padding:6px 4px;">{count}</td>'
                f"</tr>"
            )
        backpack_html += "</table>"
    else:
        backpack_html = '<p style="color:#757575;font-size:13px;">No genes pinned to backpack.</p>'

    # Achievement pills
    achievements_pills = ""
    for a in sorted(state.achievements):
        achievements_pills += (
            f'<span style="display:inline-block;padding:6px 16px;border-radius:20px;'
            f"background:linear-gradient(135deg, #D4A017, #F9A825);color:#333;"
            f'font-weight:600;font-size:12px;margin:4px;box-shadow:0 2px 4px rgba(0,0,0,0.1);">'
            f"{a}</span>"
        )
    if not achievements_pills:
        achievements_pills = '<span style="color:#757575;font-size:13px;">No achievements yet</span>'

    # -- Assemble full report --
    html = f'''
    <div style="font-family:-apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                max-width:720px;margin:0 auto;">

        <!-- Header -->
        <div style="background:{COLORS["bg_dark"]};color:white;padding:32px 36px;
                    border-radius:16px;margin-bottom:24px;">
            <div style="font-size:12px;font-weight:600;text-transform:uppercase;
                        letter-spacing:1.5px;color:{COLORS["text_muted"]};margin-bottom:8px;">
                Field Report</div>
            <h2 style="margin:0 0 6px 0;font-size:26px;font-weight:700;">{line_id}</h2>
            <p style="margin:0;font-size:14px;color:{COLORS["text_muted"]};">
                Origin: {country} &middot; Cluster {cluster_id}</p>
        </div>

        <!-- Metric cards -->
        <div style="display:flex;gap:12px;margin-bottom:24px;flex-wrap:wrap;">
            {metric_cards}
        </div>

        <!-- Nearest Neighbors -->
        <div style="background:#FFFFFF;border-radius:12px;padding:20px 24px;
                    box-shadow:0 2px 8px rgba(0,0,0,0.06);margin-bottom:16px;">
            <h3 style="margin:0 0 12px 0;font-size:15px;font-weight:600;
                       color:{COLORS["text_primary"]};">Nearest Neighbors</h3>
            {neighbor_html}
        </div>

        <!-- Gene Composition -->
        <div style="background:#FFFFFF;border-radius:12px;padding:20px 24px;
                    box-shadow:0 2px 8px rgba(0,0,0,0.06);margin-bottom:16px;">
            <h3 style="margin:0 0 12px 0;font-size:15px;font-weight:600;
                       color:{COLORS["text_primary"]};">Gene Composition</h3>
            {composition_bar}
        </div>

        <!-- Top 5 Rare Genes -->
        <div style="background:#FFFFFF;border-radius:12px;padding:20px 24px;
                    box-shadow:0 2px 8px rgba(0,0,0,0.06);margin-bottom:16px;">
            <h3 style="margin:0 0 12px 0;font-size:15px;font-weight:600;
                       color:{COLORS["text_primary"]};">Top 5 Rare Genes</h3>
            {rare_html}
        </div>

        <!-- Backpack Collection -->
        <div style="background:#FFFFFF;border-radius:12px;padding:20px 24px;
                    box-shadow:0 2px 8px rgba(0,0,0,0.06);margin-bottom:16px;">
            <h3 style="margin:0 0 12px 0;font-size:15px;font-weight:600;
                       color:{COLORS["text_primary"]};">Backpack Collection</h3>
            {backpack_html}
        </div>

        <!-- Achievements -->
        <div style="background:#FFFFFF;border-radius:12px;padding:20px 24px;
                    box-shadow:0 2px 8px rgba(0,0,0,0.06);margin-bottom:16px;">
            <h3 style="margin:0 0 12px 0;font-size:15px;font-weight:600;
                       color:{COLORS["text_primary"]};">Achievements Earned</h3>
            <div>{achievements_pills}</div>
        </div>

        <!-- Footer -->
        <div style="text-align:center;padding:16px 0;font-size:11px;color:#757575;">
            Generated by Pigeon Pea Pangenome Atlas
        </div>
    </div>
    '''
    return html


def build_achievements_html(state) -> str:
    """Build styled HTML for achievement badges.

    Parameters
    ----------
    state : AppState
        Current application state.

    Returns
    -------
    str
        HTML string with styled badge pills.
    """
    if state is None or not state.achievements:
        return (
            '<div style="padding:12px;text-align:center;color:#757575;">'
            "Complete quests to earn badges!</div>"
        )

    pills = []
    for a in sorted(state.achievements):
        pills.append(
            f'<span class="achievement-badge">{a}</span>'
        )
    return '<div style="display:flex;flex-wrap:wrap;gap:6px;">' + "".join(pills) + "</div>"


# =====================================================================
# Gradio UI builder
# =====================================================================

def build_final_tab():
    """Build Final Report tab components. Returns dict of components.

    Returned keys (prefixed with ``final_`` by layout.py):
        tab, generate_btn, report_md, download_json, download_csv, achievements_html
    """
    with gr.Tab("Your Report", id="final") as tab:
        # -- Header --
        gr.HTML(
            '<div style="padding:4px 0 12px 0;">'
            '<h2 style="margin:0 0 4px 0;font-size:22px;font-weight:700;'
            f'color:{COLORS["text_primary"]};">Chapter 6: Your Report</h2>'
            '<p style="margin:0;font-size:14px;color:#757575;line-height:1.5;">'
            "Generate a presentation-ready summary of your pangenome exploration, "
            "including your selected line, findings, and backpack collection."
            "</p></div>"
        )

        generate_btn = gr.Button(
            "Generate Report",
            variant="primary",
            size="lg",
        )

        report_md = gr.Markdown(
            value="*Click 'Generate Report' to create your field report.*"
        )

        with gr.Row():
            download_json = gr.File(label="Download JSON", visible=False)
            download_csv = gr.File(label="Download CSV", visible=False)

        gr.HTML(
            '<div style="margin-top:16px;padding:4px 0;">'
            '<h3 style="margin:0 0 4px 0;font-size:17px;font-weight:600;'
            f'color:{COLORS["text_primary"]};">Achievements Earned</h3>'
            "</div>"
        )
        achievements_html = gr.HTML(
            value=(
                '<div style="padding:12px;text-align:center;color:#757575;">'
                "Complete quests to earn badges!</div>"
            ),
        )

    return {
        "tab": tab,
        "generate_btn": generate_btn,
        "report_md": report_md,
        "download_json": download_json,
        "download_csv": download_csv,
        "achievements_html": achievements_html,
    }