"""Produce a valid PRIMO submission in one command, then swap in your own model.

Downloads every public dataset, embeds each one, and writes the single file the
Submit tab expects. The embedding here is deliberately dumb -- log2(CPM+1) then
PCA -- because the point is the plumbing, not the score: replace ``embed`` with
your encoder and nothing else changes.

    pip install anndata scikit-learn pandas pyyaml huggingface_hub
    python quickstart.py --out submission.parquet

Standalone on purpose: no import from this Space and none from our monorepo, so
it keeps working if you copy the file into your own project.
"""

import argparse
from pathlib import Path

import anndata as ad
import numpy as np
import pandas as pd
import yaml
from huggingface_hub import snapshot_download
from sklearn.decomposition import PCA

PUBLIC_REPO = "PRIMOmics/primo"
MANIFEST_FILENAME = "datasets.yaml"

DATASET_ID = "dataset_id"
SAMPLE_ID = "sample_id"

TARGET_SUM = 1_000_000
N_COMPONENTS = 50
RANDOM_STATE = 0


def embed(adata: ad.AnnData) -> np.ndarray:
    """One dataset's raw counts -> one vector per patient. Replace me.

    Whatever you return, the contract is the same: one row per sample, in
    ``adata.obs_names`` order, all finite. The embedding width is yours to pick
    and may differ from one dataset to the next.
    """
    x = adata.X
    x = x.toarray() if hasattr(x, "toarray") else np.asarray(x)
    x = x.astype(float)
    counts = x.sum(axis=1, keepdims=True)
    x = np.log2(x / np.where(counts == 0, 1, counts) * TARGET_SUM + 1)
    k = min(N_COMPONENTS, x.shape[0] - 1, x.shape[1])
    return PCA(n_components=k, random_state=RANDOM_STATE).fit_transform(x)


def download(token: str | None) -> Path:
    """Pull the public benchmark (manifest + every ``expression.h5ad``)."""
    return Path(snapshot_download(PUBLIC_REPO, repo_type="dataset", token=token))


def dataset_ids(root: Path) -> list[str]:
    """The opaque ids to embed, read off the public manifest."""
    manifest = yaml.safe_load((root / MANIFEST_FILENAME).read_text())
    entries = manifest.get("datasets", []) if isinstance(manifest, dict) else manifest
    return [str(entry["id"]) for entry in entries]


def build(root: Path) -> pd.DataFrame:
    """Embed every dataset into the one frame the Submit tab expects.

    Datasets of different widths stack into one table; the extra columns of a
    narrower dataset stay empty and the evaluator drops them per dataset, so each
    dataset keeps its own embedding size.
    """
    blocks = []
    for dataset_id in dataset_ids(root):
        adata = ad.read_h5ad(root / dataset_id / "expression.h5ad")
        vectors = embed(adata)
        print(f"{dataset_id}: {adata.n_obs} samples -> {vectors.shape[1]} dims")
        block = pd.DataFrame(
            vectors, columns=[f"e{i}" for i in range(vectors.shape[1])]
        )
        block.insert(0, SAMPLE_ID, adata.obs_names.to_numpy())
        block.insert(0, DATASET_ID, dataset_id)
        blocks.append(block)
    return pd.concat(blocks, ignore_index=True)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--out", type=Path, default=Path("submission.parquet"))
    parser.add_argument("--token", default=None, help="HF token, if you need one.")
    args = parser.parse_args()

    submission = build(download(args.token))
    if args.out.suffix == ".csv":
        submission.to_csv(args.out, index=False)
    else:
        submission.to_parquet(args.out, index=False)
    print(f"\nWrote {args.out}: {len(submission)} rows. Upload it on the Submit tab.")


if __name__ == "__main__":
    main()