scripts/export_web_data.py

from __future__ import annotations

import json
import re
import sys
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(sys.argv[1]).resolve()
OUT = Path(sys.argv[2]).resolve()
APP_ROOT = ROOT
sys.path.insert(0, str(APP_ROOT))

from app import ANNOTATION_EXAMPLES, ASSEMBLY_SUMMARY, SPECIES_METRICS  # noqa: E402
from explorer.dimensionality import DimensionalityRepository, compute_projection  # noqa: E402

SPECIES = {
    "Species 01": "species01",
    "Species 02": "species02",
    "Species 03": "species03",
}
BASE_SOURCES = [
    "EggNOG_OG", "GO_EggNOG", "GO_ESM2_150M_MF",
    "GO_ESMC_600M_ProteinFunction", "ESM_localization", "GO_union",
    "Pfam_HMMER", "EC_EggNOG", "KEGG_KO_EggNOG",
    "KEGG_KO_KofamScan", "KEGG_KO_union",
]
HIERARCHIES = {
    "EggNOG_OG": ["Direct terms"],
    "GO_EggNOG": ["Direct terms", "GO ancestors"],
    "GO_ESM2_150M_MF": ["Direct terms", "GO ancestors"],
    "GO_ESMC_600M_ProteinFunction": ["Direct terms", "GO ancestors"],
    "ESM_localization": ["Direct terms"],
    "GO_union": ["Direct terms", "GO ancestors"],
    "Pfam_HMMER": ["Direct terms"],
    "EC_EggNOG": ["Direct terms", "EC level 1 class"],
    "KEGG_KO_EggNOG": ["Direct terms", "KEGG pathways", "KEGG level 1", "KEGG level 2"],
    "KEGG_KO_KofamScan": ["Direct terms", "KEGG pathways", "KEGG level 1", "KEGG level 2"],
    "KEGG_KO_union": ["Direct terms", "KEGG pathways", "KEGG level 1", "KEGG level 2"],
}


def clean(value):
    if isinstance(value, str):
        return value.replace("TRINITY_", "").replace("TRINITY", "")
    return value


def slug(value: str) -> str:
    return re.sub(r"[^a-zA-Z0-9]+", "_", value).strip("_").lower()


def write_json(path: Path, value) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, separators=(",", ":")))


def frame_payload(frame: pd.DataFrame, columns: list[str]) -> dict:
    frame = frame[columns].copy()
    for column in frame.columns:
        if column in {"gene_id", "genes", "term", "name", "direction"}:
            frame[column] = frame[column].map(clean)
        if pd.api.types.is_numeric_dtype(frame[column]):
            frame[column] = pd.to_numeric(frame[column], errors="coerce").round(8)
    return json.loads(frame.to_json(orient="split", index=False))


def source_for_hierarchy(base: str, hierarchy: str) -> str:
    if hierarchy == "Direct terms":
        return base
    if hierarchy == "GO ancestors":
        return f"{base}__GO_ANCESTORS"
    if hierarchy == "EC level 1 class":
        return "EC_EggNOG__EC_LEVEL1"
    prefixes = {
        "KEGG pathways": "KEGG_PATHWAY",
        "KEGG level 1": "KEGG_LEVEL1",
        "KEGG level 2": "KEGG_LEVEL2",
    }
    if hierarchy in prefixes:
        method_name = {
            "KEGG_KO_EggNOG": "EggNOG",
            "KEGG_KO_KofamScan": "KofamScan",
            "KEGG_KO_union": "consensus",
        }.get(base)
        if method_name:
            return f"{prefixes[hierarchy]}_{method_name}"
    return base


def main() -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "de").mkdir(exist_ok=True)
    (OUT / "enrichment").mkdir(exist_ok=True)

    app_data = {
        "title": "Transcriptomes Explorer",
        "assemblySummary": ASSEMBLY_SUMMARY.to_dict(orient="records"),
        "annotationExamples": {
            species: [
                dict(zip([
                    "Approach", "Tool / result", "Database or type", "Proteins", "Genes",
                    "Protein -> gene", "Example annotation", "Description / evidence",
                ], [clean(x) for x in row]))
                for row in rows
            ]
            for species, rows in ANNOTATION_EXAMPLES.items()
        },
        "species": [],
        "deIndex": {},
        "enrichmentIndex": {},
        "sourceOrder": BASE_SOURCES,
        "hierarchies": HIERARCHIES,
    }

    for species_label, species_slug in SPECIES.items():
        data_dir = ROOT / "data" / species_slug
        manifest = json.loads((data_dir / "manifest.json").read_text())
        contrasts = pd.read_parquet(data_dir / "contrast_summary.parquet")
        de = pd.read_parquet(data_dir / "de_significant.parquet")
        enrichment = pd.read_parquet(data_dir / "enrichment.parquet")
        species_meta = SPECIES_METRICS[species_label]
        app_data["species"].append({
            "label": species_label,
            "slug": species_slug,
            "genes": species_meta["genes"],
            "proteins": species_meta["proteins"],
            "contrasts": len(contrasts),
            "samples": manifest.get("sample_count", 0),
        })
        app_data["deIndex"][species_slug] = []
        app_data["enrichmentIndex"][species_slug] = []
        contrast_meta = contrasts.set_index("contrast_id").to_dict(orient="index")

        for contrast_id, meta in contrast_meta.items():
            contrast_id = str(contrast_id)
            for method in ("DESeq2", "edgeR"):
                method_frame = de[(de["contrast_id"].astype(str) == contrast_id) & (de["method"] == method)].copy()
                if method == "DESeq2":
                    columns = ["gene_id", "baseMean", "log2FoldChange", "lfcSE", "stat", "pvalue", "padj", "direction"]
                else:
                    columns = ["gene_id", "logFC", "logCPM", "F", "PValue", "FDR", "direction"]
                columns = [c for c in columns if c in method_frame.columns]
                filename = f"{species_slug}_{slug(contrast_id)}_{method.lower()}.json"
                write_json(OUT / "de" / filename, frame_payload(method_frame, columns))
                app_data["deIndex"][species_slug].append({
                    "contrast_id": contrast_id,
                    "method": method,
                    "file": f"de/{filename}",
                    "group_a": str(meta.get("group_a", "group 1")),
                    "group_b": str(meta.get("group_b", "group 2")),
                    "n_group_a": int(meta.get("n_group_a", 0)),
                    "n_group_b": int(meta.get("n_group_b", 0)),
                    "n_genes_input": int(meta.get("n_genes_input", 0)),
                    "n_genes_tested": int(meta.get("n_genes_tested", 0)),
                })

                eframe = enrichment[
                    (enrichment["contrast_id"].astype(str) == contrast_id)
                    & (enrichment["method"] == method)
                ].copy()
                ecolumns = [
                    "source_id", "direction", "term", "name", "overlap",
                    "foreground_size", "term_size", "background_size",
                    "fold_enrichment", "pvalue", "padj", "genes",
                ]
                rows = []
                for source_id, source_frame in eframe.groupby("source_id", sort=False):
                    source_frame = source_frame.sort_values(["padj", "pvalue"], na_position="last")
                    rows.append(source_frame.head(40))
                    app_data["enrichmentIndex"][species_slug].append({
                        "contrast_id": contrast_id,
                        "method": method,
                        "source_id": str(source_id),
                        "file": f"enrichment/{species_slug}_{slug(contrast_id)}_{method.lower()}.json",
                        "universe_terms": int(source_frame["term"].nunique()),
                    })
                eweb = pd.concat(rows, ignore_index=True) if rows else eframe.head(0)
                write_json(OUT / "enrichment" / f"{species_slug}_{slug(contrast_id)}_{method.lower()}.json", frame_payload(eweb, ecolumns))

        # Static 3D coordinates are precomputed once for a stable browser-only demo.
        dim_root = ROOT / "data" / "dimensionality"
        dim_repo = DimensionalityRepository(dim_root)
        projections = {}
        for method in ("PCA", "UMAP"):
            for label, limit in (("all", None), ("5000", 5000), ("1000", 1000)):
                expression, metadata = dim_repo.load(species_label)
                result = compute_projection(expression, metadata, method, limit)
                projections[f"{method}_{label}"] = {
                    "axes": list(result.axis_columns),
                    "selected_gene_count": result.selected_gene_count,
                    "expressed_gene_count": result.expressed_gene_count,
                    "explained_variance": list(result.explained_variance) if result.explained_variance else None,
                    "rows": json.loads(result.coordinates.to_json(orient="records")),
                }
        write_json(OUT / "dimensionality" / f"{species_slug}.json", projections)

    write_json(OUT / "app-data.json", app_data)


if __name__ == "__main__":
    main()