viewers.py

"""Structure viewers and display formatting for the Gradio UI."""

import re

import gradio as gr
import pandas as pd

from file_utils import _df_to_csv


def _format_display_number(value) -> str:
    if pd.isna(value):
        return ""
    return f"{float(value):.2f}".rstrip("0").rstrip(".")


def _file_output(value: str | None) -> dict:
    return gr.update(value=value, visible=bool(value))


def _csv_download_output(df: pd.DataFrame | None) -> dict:
    return _file_output(_df_to_csv(df))


def _format_results_display(df: pd.DataFrame):
    numeric_columns = [col for col in df.columns[-4:] if pd.api.types.is_numeric_dtype(df[col])]
    if not numeric_columns:
        return df
    return df.style.format({col: _format_display_number for col in numeric_columns})


def _get_best_sc_sample(df: pd.DataFrame) -> str | None:
    if df.empty or "Sample" not in df.columns:
        return None
    if "tmalign_score" in df.columns and df["tmalign_score"].notna().any():
        return str(df.loc[df["tmalign_score"].idxmax(), "Sample"])
    return str(df.iloc[0]["Sample"])


def _render_af2_viewer(sample_name: str | None, af2_pdb_data: dict[str, str], color_mode: str = "plddt") -> str:
    if not sample_name or not af2_pdb_data or sample_name not in af2_pdb_data:
        return ""
    import molview as mv

    v = mv.view(width=840, height=500)
    v.addModel(af2_pdb_data[sample_name], name=f"AF2: {sample_name}")
    v.setColorMode(color_mode)
    v.setBackgroundColor("#000000")
    return v._repr_html_()


def _render_reference_viewer(sample_name: str | None, input_pdb_data: dict[str, str], color_mode: str = "chain") -> str:
    if not sample_name or not input_pdb_data:
        return ""
    input_key = re.sub(r"_sample\d+$", "", sample_name)
    if input_key not in input_pdb_data:
        return ""
    import molview as mv

    v = mv.view(width=840, height=500)
    v.addModel(input_pdb_data[input_key], name=f"Reference: {input_key}")
    v.setColorMode(color_mode)
    v.setBackgroundColor("#000000")
    return v._repr_html_()


def _update_viewers(
    best_sample: str,
    af2_pdb_data: dict[str, str],
    input_pdb_data: dict[str, str],
    show_overlay: bool,
    color_mode: str = "plddt",
    ref_color_mode: str = "chain",
):
    af2_html = _render_af2_viewer(best_sample, af2_pdb_data, color_mode)
    if show_overlay:
        ref_html = _render_reference_viewer(best_sample, input_pdb_data, ref_color_mode)
        return af2_html, gr.update(value=ref_html, visible=True), gr.update(visible=True)
    return af2_html, gr.update(value="", visible=False), gr.update(visible=False)