bioai/ui/charts.py

"""
Biopesticide-AI visualization helpers.

Generates matplotlib charts for the Gradio UI:
- Efficacy distribution chart (bar chart of top candidates)
- Off-target heatmap (species x candidates)
- Half-life comparison chart
- GC content gauge

All charts use the project palette and are saved as PNG for Gradio Image display.
"""

from __future__ import annotations

import io
from pathlib import Path
from typing import Dict, List

import matplotlib

matplotlib.use("Agg")  # non-interactive backend
import matplotlib.font_manager as fm
import matplotlib.pyplot as plt
import numpy as np

# Font setup — rely on fontconfig discovery (don't addfont variable fonts)
plt.rcParams["font.sans-serif"] = ["Noto Sans SC", "DejaVu Sans", "Liberation Sans"]
plt.rcParams["axes.unicode_minus"] = False

# Project palette (matches the PDF and UI)
PALETTE = {
    "accent": "#2f86b2",
    "accent_2": "#ba5a6a",
    "header_fill": "#4a616c",
    "text_primary": "#242627",
    "text_muted": "#71777a",
    "border": "#a1b9c6",
    "card_bg": "#ecedee",
    "page_bg": "#f4f5f6",
    "success": "#449f63",
    "warning": "#b69045",
    "error": "#964039",
    "info": "#4c7094",
}


def _style_axes(ax, title: str = "", xlabel: str = "", ylabel: str = ""):
    """Apply consistent styling to an axes object."""
    ax.set_title(title, fontsize=12, fontweight="bold", color=PALETTE["text_primary"], pad=12)
    ax.set_xlabel(xlabel, fontsize=10, color=PALETTE["text_muted"])
    ax.set_ylabel(ylabel, fontsize=10, color=PALETTE["text_muted"])
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color(PALETTE["border"])
    ax.spines["bottom"].set_color(PALETTE["border"])
    ax.tick_params(colors=PALETTE["text_muted"], labelsize=9)
    ax.grid(axis="y", linestyle="--", alpha=0.3, color=PALETTE["border"])


def efficacy_bar_chart(candidates: List[Dict]) -> str:
    """Horizontal bar chart of top candidates by efficacy score.

    Returns path to saved PNG.
    """
    if not candidates:
        return _empty_chart("No candidates to display")

    n = len(candidates)
    fig, ax = plt.subplots(figsize=(7, max(2.5, 0.5 * n + 1)), constrained_layout=True)
    fig.patch.set_facecolor("white")

    labels = [f"#{i+1} {c.get('sirna_seq', '')[:8]}..." for i, c in enumerate(candidates)]
    efficacies = [c.get("efficacy", 0) for c in candidates]
    scores = [c.get("final_score", 0) for c in candidates]

    y = np.arange(n)
    # Color gradient: top candidates get accent blue, lower ones get muted
    colors = [PALETTE["accent"] if s > 0.3 else PALETTE["text_muted"] for s in scores]

    bars = ax.barh(y, efficacies, color=colors, height=0.6, edgecolor="white", linewidth=0.5)
    ax.set_yticks(y)
    ax.set_yticklabels(labels, fontsize=9, color=PALETTE["text_primary"])
    ax.invert_yaxis()  # top candidate at top
    ax.set_xlim(0, 1.0)

    # Add value labels on bars
    for bar, eff in zip(bars, efficacies):
        ax.text(bar.get_width() + 0.01, bar.get_y() + bar.get_height() / 2,
                f"{eff:.3f}", va="center", fontsize=9, color=PALETTE["text_primary"])

    _style_axes(ax, title="Efficacy Scores (Top Candidates)", xlabel="Predicted Efficacy (0-1)")

    out = _save_temp(fig, "efficacy_chart.png")
    return out


def offtarget_heatmap(candidates: List[Dict], species_names: List[str]) -> str:
    """Heatmap of off-target risk: candidates (rows) x species (cols).

    Returns path to saved PNG.
    """
    if not candidates or not species_names:
        return _empty_chart("No off-target data to display")

    n_cand = len(candidates)
    n_sp = len(species_names)
    matrix = np.zeros((n_cand, n_sp))
    for i, c in enumerate(candidates):
        per_sp = c.get("offtarget_per_species", {})
        for j, sp in enumerate(species_names):
            matrix[i, j] = per_sp.get(sp, 0.0)

    fig, ax = plt.subplots(figsize=(max(7, 0.5 * n_sp + 4), max(3, 0.5 * n_cand + 2)), constrained_layout=True)
    fig.patch.set_facecolor("white")

    # Custom colormap: white -> warning yellow -> error red
    from matplotlib.colors import LinearSegmentedColormap
    cmap = LinearSegmentedColormap.from_list("risk", ["#ffffff", "#fef6e4", "#b69045", "#964039"])

    im = ax.imshow(matrix, aspect="auto", cmap=cmap, vmin=0, vmax=max(0.1, matrix.max()))
    ax.set_xticks(np.arange(n_sp))
    ax.set_xticklabels([sp.replace("_", " ") for sp in species_names],
                       rotation=45, ha="right", fontsize=8, color=PALETTE["text_primary"])
    ax.set_yticks(np.arange(n_cand))
    ax.set_yticklabels([f"#{i+1}" for i in range(n_cand)], fontsize=9, color=PALETTE["text_primary"])

    # Add value annotations
    for i in range(n_cand):
        for j in range(n_sp):
            val = matrix[i, j]
            if val > 0:
                ax.text(j, i, f"{val:.2f}", ha="center", va="center",
                        fontsize=7, color=PALETTE["text_primary"])

    ax.set_title("Off-Target Risk Heatmap", fontsize=12, fontweight="bold",
                 color=PALETTE["text_primary"], pad=12)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    # Colorbar
    cbar = fig.colorbar(im, ax=ax, shrink=0.6, pad=0.02)
    cbar.set_label("Risk (0-1)", fontsize=9, color=PALETTE["text_muted"])
    cbar.ax.tick_params(colors=PALETTE["text_muted"], labelsize=8)

    out = _save_temp(fig, "offtarget_heatmap.png")
    return out


def halflife_chart(candidates: List[Dict]) -> str:
    """Bar chart of predicted half-lives with risk-tier color coding.

    Returns path to saved PNG.
    """
    if not candidates:
        return _empty_chart("No half-life data to display")

    n = len(candidates)
    fig, ax = plt.subplots(figsize=(7, max(2.5, 0.5 * n + 1)), constrained_layout=True)
    fig.patch.set_facecolor("white")

    labels = [f"#{i+1} {c.get('sirna_seq', '')[:8]}..." for i, c in enumerate(candidates)]
    half_lives = [c.get("half_life_hours", 0) for c in candidates]

    # Color by risk tier: <24h = warning, 24-72h = success, >72h = info
    colors = []
    for hl in half_lives:
        if hl < 24:
            colors.append(PALETTE["error"])
        elif hl < 72:
            colors.append(PALETTE["success"])
        else:
            colors.append(PALETTE["info"])

    y = np.arange(n)
    bars = ax.barh(y, half_lives, color=colors, height=0.6, edgecolor="white", linewidth=0.5)
    ax.set_yticks(y)
    ax.set_yticklabels(labels, fontsize=9, color=PALETTE["text_primary"])
    ax.invert_yaxis()

    # Add value labels
    for bar, hl in zip(bars, half_lives):
        ax.text(bar.get_width() + 2, bar.get_y() + bar.get_height() / 2,
                f"{hl:.1f}h ({hl/24:.1f}d)", va="center", fontsize=9, color=PALETTE["text_primary"])

    _style_axes(ax, title="Predicted Environmental Half-Life",
                xlabel="Half-life (hours)")

    # Add reference lines for risk tiers
    ax.axvline(x=24, color=PALETTE["error"], linestyle=":", alpha=0.5, linewidth=1)
    ax.axvline(x=72, color=PALETTE["success"], linestyle=":", alpha=0.5, linewidth=1)
    ax.text(24, -0.7, "1 day", fontsize=7, color=PALETTE["text_muted"], ha="center")
    ax.text(72, -0.7, "3 days", fontsize=7, color=PALETTE["text_muted"], ha="center")

    out = _save_temp(fig, "halflife_chart.png")
    return out


def _empty_chart(message: str) -> str:
    """Generate a placeholder chart with a message."""
    fig, ax = plt.subplots(figsize=(6, 3), constrained_layout=True)
    fig.patch.set_facecolor("white")
    ax.text(0.5, 0.5, message, ha="center", va="center",
            fontsize=12, color=PALETTE["text_muted"], style="italic")
    ax.axis("off")
    return _save_temp(fig, "empty.png")


def _save_temp(fig, filename: str) -> str:
    """Save figure to a temp directory and close it. Returns the path."""
    import tempfile
    import os
    tmp_dir = Path(tempfile.gettempdir()) / "bioai_charts"
    tmp_dir.mkdir(parents=True, exist_ok=True)
    out = str(tmp_dir / filename)
    fig.savefig(out, dpi=150, facecolor="white", bbox_inches=None)
    plt.close(fig)
    return out