"""Build and cache the BRCA workspace with both GenePT files, using the app's own prepare step.
Usage: build_workspaces.py <cache_dir>"""
import pickle
import sys
from pathlib import Path

import numpy as np

import app

CACHE = Path(sys.argv[1])
CACHE.mkdir(parents=True, exist_ok=True)
OPTIONS = {
    "bulk": "Bulk · large-3 context",
    "singlecell": "Single-cell · cell type, tissue, drug & pathway",
}
for key, option in OPTIONS.items():
    result = app.prepare_workspace(
        "BioLM-NET examples", "BRCA", "", None, None, None, None, None, None, None,
        True, option, progress=lambda *a, **k: None,
    )
    workspace = result[0]
    if workspace is None:
        sys.exit(f"{key}: {result[1]}")
    (CACHE / f"workspace_BRCA_{key}.pkl").write_bytes(pickle.dumps(workspace))
    for name, branch in (("gene", workspace.gene_branch), ("dna", workspace.dna_branch)):
        m = branch.biological_mask
        print(f"{key:10s} {name}: inputs={m.shape[0]} hidden genes={m.shape[1]} edges={(m > 0).sum()} "
              f"median in-degree={np.median((m > 0).sum(0)):.0f} pathways={len(branch.pathways)} "
              f"GenePT dim={branch.embeddings.shape[1]}", flush=True)