biolmnet/data.py

from __future__ import annotations

import hashlib
import io
import re
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from pathlib import Path
from typing import BinaryIO

import numpy as np
import pandas as pd
from scipy.stats import hypergeom


UPSTREAM_REPOSITORY = "https://github.com/bozdaglab/BioLM-NET"
UPSTREAM_RAW = "https://raw.githubusercontent.com/bozdaglab/BioLM-NET/main"
GENEPT_REPOSITORY = "honicky/genept-composable-embeddings"
DATASET_FILES = {
    "gene": "Gene_Expression.csv",
    "dna": "DNA_Methylation.csv",
    "labels": "label.csv",
    "gene_pathways": "ge_target_to_KEGG_significant.csv",
    "dna_pathways": "dna_target_to_KEGG_significant.csv",
}


@dataclass
class BranchPriors:
    input_genes: list[str]
    hidden_genes: list[str]
    biological_mask: np.ndarray
    pathways: list[str] = field(default_factory=list)
    pathway_mask: np.ndarray | None = None
    embeddings: np.ndarray | None = None
    pdi_edges: int = 0
    ppi_edges: int = 0
    missing_embedding_genes: int = 0


@dataclass
class PreparedWorkspace:
    gene_expression: np.ndarray
    dna_methylation: np.ndarray
    labels: np.ndarray
    label_names: list[str]
    gene_branch: BranchPriors
    dna_branch: BranchPriors
    source_name: str
    warnings: list[str] = field(default_factory=list)


def _normalise_columns(frame: pd.DataFrame) -> pd.DataFrame:
    result = frame.copy()
    result.columns = [str(value).strip() for value in result.columns]
    return result


def read_csv(source: str | Path | BinaryIO) -> pd.DataFrame:
    if hasattr(source, "read"):
        return _normalise_columns(pd.read_csv(source))
    text = str(source)
    if text.startswith(("http://", "https://")):
        request = urllib.request.Request(
            text, headers={"User-Agent": "BioLM-NET-HuggingFace-Space/1.0"}
        )
        with urllib.request.urlopen(request, timeout=60) as response:
            payload = response.read()
        return _normalise_columns(pd.read_csv(io.BytesIO(payload)))
    return _normalise_columns(pd.read_csv(text))


def github_folder_to_raw_base(folder_url: str) -> str:
    value = folder_url.strip().rstrip("/")
    if value.startswith("https://raw.githubusercontent.com/"):
        return value
    match = re.match(
        r"https://github\.com/([^/]+)/([^/]+)/(?:tree|blob)/([^/]+)(?:/(.*))?$",
        value,
    )
    if match:
        owner, repository, branch, folder = match.groups()
        suffix = f"/{folder}" if folder else ""
        return (
            f"https://raw.githubusercontent.com/{owner}/{repository}/"
            f"{branch}{suffix}"
        )
    match = re.match(r"https://github\.com/([^/]+)/([^/]+)$", value)
    if match:
        owner, repository = match.groups()
        return f"https://raw.githubusercontent.com/{owner}/{repository}/main"
    raise ValueError(
        "Use a GitHub repository/folder URL such as "
        "https://github.com/owner/repo/tree/main/Dataset/BRCA."
    )


def upstream_example_sources(dataset: str) -> dict[str, str]:
    if dataset not in {"BRCA", "COAD", "GBM", "scTrioseq2"}:
        raise ValueError(f"Unknown BioLM-NET example dataset: {dataset}")
    base = f"{UPSTREAM_RAW}/Dataset/{dataset}"
    return {key: f"{base}/{filename}" for key, filename in DATASET_FILES.items()}


def github_dataset_sources(folder_url: str) -> dict[str, str]:
    base = github_folder_to_raw_base(folder_url)
    return {key: f"{base}/{filename}" for key, filename in DATASET_FILES.items()}


def upstream_interaction_sources() -> tuple[str, str]:
    return (
        f"{UPSTREAM_RAW}/Dataset/PDI/PDI.csv",
        f"{UPSTREAM_RAW}/Dataset/PPI/PPI.csv",
    )


def validate_and_align_omics(
    gene_frame: pd.DataFrame,
    dna_frame: pd.DataFrame,
    labels_frame: pd.DataFrame,
    *,
    allow_preset_trim: bool = False,
) -> tuple[pd.DataFrame, pd.DataFrame, np.ndarray, list[str], list[str]]:
    warnings: list[str] = []
    if gene_frame.columns.duplicated().any() or dna_frame.columns.duplicated().any():
        raise ValueError("Omics files must have unique gene-name columns.")
    if labels_frame.shape[1] != 1:
        raise ValueError("The label file must contain exactly one column.")
    if gene_frame.empty or dna_frame.empty or labels_frame.empty:
        raise ValueError("Gene expression, DNA methylation, and labels cannot be empty.")

    counts = [len(gene_frame), len(dna_frame), len(labels_frame)]
    if len(set(counts)) != 1:
        if not allow_preset_trim:
            raise ValueError(
                "The two omics files and label file must contain the same number "
                f"of rows; received {counts}."
            )
        common = min(counts)
        warnings.append(
            f"The upstream example has row counts {counts}; all inputs were "
            f"aligned to the first {common} rows, matching repository order."
        )
        gene_frame = gene_frame.iloc[:common].reset_index(drop=True)
        dna_frame = dna_frame.iloc[:common].reset_index(drop=True)
        labels_frame = labels_frame.iloc[:common].reset_index(drop=True)

    for name, frame in (
        ("gene expression", gene_frame),
        ("DNA methylation", dna_frame),
    ):
        converted = frame.apply(pd.to_numeric, errors="coerce")
        invalid = int(converted.isna().sum().sum())
        if invalid:
            raise ValueError(
                f"{name.title()} contains {invalid:,} missing or non-numeric values."
            )
        if not np.isfinite(converted.to_numpy(dtype=np.float64)).all():
            raise ValueError(f"{name.title()} contains infinite values.")
        if name == "gene expression":
            gene_frame = converted
        else:
            dna_frame = converted

    raw_labels = labels_frame.iloc[:, 0]
    if raw_labels.isna().any():
        raise ValueError("Labels cannot be empty.")
    labels = raw_labels.astype(str).str.strip()
    if labels.eq("").any():
        raise ValueError("Labels cannot be empty.")
    unique_labels = sorted(labels.unique().tolist())
    if len(unique_labels) < 2:
        raise ValueError("Training requires at least two label classes.")
    label_to_index = {label: index for index, label in enumerate(unique_labels)}
    encoded = labels.map(label_to_index).to_numpy(dtype=np.int64)
    return gene_frame, dna_frame, encoded, unique_labels, warnings


def _clean_interactions(
    pdi_frame: pd.DataFrame, ppi_frame: pd.DataFrame
) -> tuple[pd.DataFrame, pd.DataFrame]:
    required_pdi = {"TF", "Target"}
    required_ppi = {"protein1", "protein2", "combined_score"}
    if not required_pdi.issubset(pdi_frame.columns):
        raise ValueError("PDI.csv must contain TF and Target columns.")
    if not required_ppi.issubset(ppi_frame.columns):
        raise ValueError(
            "PPI.csv must contain protein1, protein2, and combined_score columns."
        )

    pdi = pdi_frame.loc[:, ["TF", "Target"]].dropna().copy()
    pdi["TF"] = pdi["TF"].astype(str).str.strip()
    pdi["Target"] = pdi["Target"].astype(str).str.strip()
    pdi = pdi[(pdi["TF"] != "") & (pdi["Target"] != "")].drop_duplicates()

    ppi = ppi_frame.loc[
        :, ["protein1", "protein2", "combined_score"]
    ].dropna().copy()
    ppi["protein1"] = ppi["protein1"].astype(str).str.strip()
    ppi["protein2"] = ppi["protein2"].astype(str).str.strip()
    ppi["combined_score"] = pd.to_numeric(
        ppi["combined_score"], errors="coerce"
    )
    ppi = ppi.dropna()
    if ppi["combined_score"].max() > 1:
        ppi["combined_score"] = ppi["combined_score"] / 1000.0
    ppi = ppi[ppi["combined_score"] > 0.7]
    if ppi.empty:
        raise ValueError("No PPI interactions remain above score 0.7.")
    threshold = float(ppi["combined_score"].quantile(0.9))
    ppi = ppi[ppi["combined_score"] >= threshold].drop_duplicates(
        ["protein1", "protein2"]
    )
    return pdi, ppi


def build_biological_mask(
    input_genes: list[str],
    pdi_frame: pd.DataFrame,
    ppi_frame: pd.DataFrame,
) -> BranchPriors:
    pdi, ppi = _clean_interactions(pdi_frame, ppi_frame)
    input_genes = [str(gene).strip() for gene in input_genes]
    input_index = {gene: index for index, gene in enumerate(input_genes)}
    input_set = set(input_genes)

    # The paper retains PDI targets that are DE/HVG; both TF and target must
    # therefore be represented in the input branch.
    pdi_selected = pdi[
        pdi["TF"].isin(input_set) & pdi["Target"].isin(input_set)
    ].copy()

    # STRING PPI is undirected. Add the partner of every input protein,
    # regardless of which endpoint it occupies in the source file.
    forward = ppi[ppi["protein1"].isin(input_set)].rename(
        columns={"protein1": "source", "protein2": "target"}
    )
    reverse = ppi[ppi["protein2"].isin(input_set)].rename(
        columns={"protein2": "source", "protein1": "target"}
    )
    ppi_selected = pd.concat(
        [
            forward[["source", "target", "combined_score"]],
            reverse[["source", "target", "combined_score"]],
        ],
        ignore_index=True,
    ).drop_duplicates(["source", "target"])

    hidden_genes = sorted(
        set(pdi_selected["Target"].tolist())
        | set(ppi_selected["target"].tolist())
    )
    if not hidden_genes:
        raise ValueError(
            "None of the input genes have retained PDI/PPI connections. "
            "Check that columns use HGNC gene symbols."
        )
    hidden_index = {gene: index for index, gene in enumerate(hidden_genes)}
    mask = np.zeros((len(input_genes), len(hidden_genes)), dtype=np.float32)

    for row in pdi_selected.itertuples(index=False):
        mask[input_index[row.TF], hidden_index[row.Target]] = 1.0
    for row in ppi_selected.itertuples(index=False):
        source_index = input_index[row.source]
        target_index = hidden_index[row.target]
        mask[source_index, target_index] = max(
            mask[source_index, target_index], float(row.combined_score)
        )

    active = mask.sum(axis=0) > 0
    return BranchPriors(
        input_genes=input_genes,
        hidden_genes=[
            gene for gene, keep in zip(hidden_genes, active, strict=True) if keep
        ],
        biological_mask=mask[:, active],
        pdi_edges=int(len(pdi_selected)),
        ppi_edges=int(len(ppi_selected)),
    )


def _benjamini_hochberg(p_values: np.ndarray) -> np.ndarray:
    count = len(p_values)
    order = np.argsort(p_values)
    ranked = p_values[order]
    adjusted = ranked * count / np.arange(1, count + 1)
    adjusted = np.minimum.accumulate(adjusted[::-1])[::-1]
    output = np.empty_like(adjusted)
    output[order] = np.clip(adjusted, 0.0, 1.0)
    return output


def build_pathway_mask(
    hidden_genes: list[str],
    pathway_frame: pd.DataFrame,
    *,
    precomputed_significant: bool,
    adjusted_p_threshold: float = 0.05,
) -> tuple[list[str], np.ndarray, pd.DataFrame]:
    required = {"SYMBOL", "PathwayID"}
    if not required.issubset(pathway_frame.columns):
        raise ValueError("Pathway data must contain SYMBOL and PathwayID columns.")
    mapping = pathway_frame.loc[:, ["SYMBOL", "PathwayID"]].dropna().copy()
    mapping["SYMBOL"] = mapping["SYMBOL"].astype(str).str.strip()
    mapping["PathwayID"] = mapping["PathwayID"].astype(str).str.strip()
    mapping = mapping[
        (mapping["SYMBOL"] != "") & (mapping["PathwayID"] != "")
    ].drop_duplicates()

    hidden_set = set(hidden_genes)
    overlap = mapping[mapping["SYMBOL"].isin(hidden_set)]
    if overlap.empty:
        raise ValueError(
            "No PDI/PPI hidden genes overlap the supplied pathway annotations."
        )

    rows: list[dict[str, float | int | str]] = []
    if precomputed_significant:
        for pathway, group in overlap.groupby("PathwayID"):
            rows.append(
                {
                    "PathwayID": pathway,
                    "overlap_genes": int(group["SYMBOL"].nunique()),
                    "adjusted_p_value": np.nan,
                }
            )
    else:
        universe = set(mapping["SYMBOL"])
        selected = hidden_set & universe
        population = len(universe)
        draws = len(selected)
        for pathway, group in mapping.groupby("PathwayID"):
            members = set(group["SYMBOL"])
            successes = len(members)
            observed = len(selected & members)
            if observed == 0:
                continue
            p_value = float(
                hypergeom.sf(observed - 1, population, successes, draws)
            )
            rows.append(
                {
                    "PathwayID": pathway,
                    "overlap_genes": observed,
                    "p_value": p_value,
                }
            )
        if rows:
            p_values = np.array([float(row["p_value"]) for row in rows])
            adjusted = _benjamini_hochberg(p_values)
            for row, value in zip(rows, adjusted, strict=True):
                row["adjusted_p_value"] = float(value)
            rows = [
                row
                for row in rows
                if float(row["adjusted_p_value"]) < adjusted_p_threshold
            ]

    enrichment = pd.DataFrame(rows)
    if enrichment.empty:
        raise ValueError(
            "No significantly enriched pathways remain at BH-adjusted p < 0.05. "
            "Upload a broader gene-to-pathway annotation catalog or revise the "
            "input feature selection."
        )
    enrichment = enrichment.sort_values(
        ["overlap_genes", "PathwayID"], ascending=[False, True]
    ).reset_index(drop=True)
    pathways = enrichment["PathwayID"].astype(str).tolist()
    gene_index = {gene: index for index, gene in enumerate(hidden_genes)}
    pathway_index = {
        pathway: index for index, pathway in enumerate(pathways)
    }
    mask = np.zeros((len(hidden_genes), len(pathways)), dtype=bool)
    kept_mapping = overlap[overlap["PathwayID"].isin(pathway_index)]
    for row in kept_mapping.itertuples(index=False):
        mask[gene_index[row.SYMBOL], pathway_index[row.PathwayID]] = True

    active_pathways = mask.sum(axis=0) > 0
    pathways = [
        pathway
        for pathway, keep in zip(pathways, active_pathways, strict=True)
        if keep
    ]
    return pathways, mask[:, active_pathways], enrichment


def deterministic_gene_embeddings(
    genes: list[str], dimensions: int = 64
) -> pd.DataFrame:
    """Deterministic test/fallback embeddings, never silently used for GenePT."""

    vectors = []
    for gene in genes:
        digest = hashlib.sha256(gene.encode("utf-8")).digest()
        seed = int.from_bytes(digest[:8], "little")
        generator = np.random.default_rng(seed)
        vector = generator.normal(0, 1, dimensions).astype(np.float32)
        vector /= max(float(np.linalg.norm(vector)), 1e-8)
        vectors.append(vector)
    return pd.DataFrame(vectors, index=genes)


def load_genept_embeddings(filename: str) -> pd.DataFrame:
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise RuntimeError(
            "huggingface_hub is required to retrieve GenePT embeddings."
        ) from exc
    path = hf_hub_download(
        repo_id=GENEPT_REPOSITORY,
        filename=filename,
        repo_type="model",
    )
    frame = pd.read_parquet(path)
    frame.index = frame.index.astype(str).str.strip()
    return frame


def attach_embeddings_and_pathways(
    branch: BranchPriors,
    embedding_frame: pd.DataFrame,
    pathway_frame: pd.DataFrame,
    *,
    precomputed_significant: bool,
) -> pd.DataFrame:
    embedding_index = set(embedding_frame.index.astype(str))
    keep = np.array(
        [gene in embedding_index for gene in branch.hidden_genes], dtype=bool
    )
    branch.missing_embedding_genes = int((~keep).sum())
    if not keep.any():
        raise ValueError(
            "No retained PDI/PPI genes have embeddings in the selected GenePT file."
        )
    branch.hidden_genes = [
        gene
        for gene, retained in zip(branch.hidden_genes, keep, strict=True)
        if retained
    ]
    branch.biological_mask = branch.biological_mask[:, keep]
    branch.embeddings = (
        embedding_frame.loc[branch.hidden_genes].to_numpy(dtype=np.float32)
    )
    (
        branch.pathways,
        branch.pathway_mask,
        enrichment,
    ) = build_pathway_mask(
        branch.hidden_genes,
        pathway_frame,
        precomputed_significant=precomputed_significant,
    )

    genes_with_pathways = branch.pathway_mask.sum(axis=1) > 0
    if not genes_with_pathways.any():
        raise ValueError("No embedded hidden genes belong to a retained pathway.")
    branch.hidden_genes = [
        gene
        for gene, retained in zip(
            branch.hidden_genes, genes_with_pathways, strict=True
        )
        if retained
    ]
    branch.biological_mask = branch.biological_mask[:, genes_with_pathways]
    branch.embeddings = branch.embeddings[genes_with_pathways]
    branch.pathway_mask = branch.pathway_mask[genes_with_pathways]
    return enrichment