utils.py

import html
import logging
from typing import List, Tuple

import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit.Chem import Draw

logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


def _draw_molecule_svg(smiles: str, size: Tuple[int, int]) -> str:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return f"<pre class='gt4sd-smiles'>{html.escape(smiles)}</pre>"
    svg = Draw.MolsToGridImage([mol], molsPerRow=1, subImgSize=size, useSVG=True)
    return str(svg).replace("<?xml version='1.0' encoding='iso-8859-1'?>", "")


def _draw_static_grid(
    result_df: pd.DataFrame, n_cols: int, size: Tuple[int, int]
) -> str:
    columns = [column for column in result_df.columns if column != "SMILES"]
    cards = []
    for _, row in result_df.iterrows():
        smiles = str(row["SMILES"])
        details = "".join(
            f"<dt>{html.escape(str(column))}</dt>"
            f"<dd>{html.escape(str(row[column]))}</dd>"
            for column in columns
        )
        cards.append(
            "<section class='gt4sd-molecule-card'>"
            f"<div class='gt4sd-molecule-image'>{_draw_molecule_svg(smiles, size)}</div>"
            f"<dl class='gt4sd-molecule-details'>{details}</dl>"
            "</section>"
        )

    return f"""
<style>
.gt4sd-molecule-grid {{
  display: grid;
  grid-template-columns: repeat({max(1, n_cols)}, minmax(180px, 1fr));
  gap: 12px;
  width: 100%;
  max-height: 900px;
  overflow: auto;
}}
.gt4sd-molecule-card {{
  border: 1px solid #e5e7eb;
  border-radius: 6px;
  background: #fff;
  padding: 12px;
  min-width: 0;
}}
.gt4sd-molecule-image svg {{
  width: 100%;
  max-width: {size[0]}px;
  height: auto;
  display: block;
  margin: 0 auto 8px;
}}
.gt4sd-molecule-details {{
  display: grid;
  grid-template-columns: max-content minmax(0, 1fr);
  gap: 4px 10px;
  margin: 0;
  font: 13px/1.35 system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}}
.gt4sd-molecule-details dt {{ color: #6b7280; font-weight: 600; }}
.gt4sd-molecule-details dd {{ margin: 0; min-width: 0; overflow-wrap: anywhere; }}
.gt4sd-smiles {{ white-space: pre-wrap; overflow-wrap: anywhere; }}
@media (max-width: 720px) {{
  .gt4sd-molecule-grid {{ grid-template-columns: 1fr; }}
}}
</style>
<div class='gt4sd-molecule-grid'>
{''.join(cards)}
</div>
"""


def draw_grid_predict(
    sequences: List[str], properties: np.array, property_names: List[str], domain: str
) -> str:
    """
    Uses RDKit SVGs to draw a HTML grid for the prediction

    Args:
        sequences: Sequences for which properties are predicted.
        properties: Predicted properties. Array of shape (n_samples, n_properties).
        names: List of property names
        domain: Domain of the prediction (molecules or proteins).

    Returns:
        HTML to display
    """

    if domain not in ["Molecules", "Proteins"]:
        raise ValueError(f"Unsupported domain {domain}")

    if domain == "Proteins":
        converter = lambda x: Chem.MolToSmiles(Chem.MolFromFASTA(x))
    else:
        converter = lambda x: x

    smiles = []
    for sequence in sequences:
        try:
            seq = converter(sequence)
            smiles.append(seq)
        except Exception:
            logger.warning(f"Could not draw sequence {seq}")

    result = pd.DataFrame({"SMILES": smiles})
    for i, name in enumerate(property_names):
        result[name] = properties[:, i]
    n_cols = min(3, len(result))
    size = (140, 200) if len(result) > 3 else (600, 700)
    return _draw_static_grid(result, n_cols=n_cols, size=size)