"""Convert a CSV/TSV expression table into the sparse H5AD used by ODlab.

Install locally: pip install anndata pandas numpy scipy
Samples/spot rows: python prepare_expression.py source.csv bulk_expression.h5ad
Gene rows: python prepare_expression.py source.tsv bulk_expression.h5ad --gene-rows
Reference: add --annotations cells.csv and use bulk_reference.h5ad as output.
The annotations first column is cell_id, followed by cell_type and optional donor fields.
This example preserves values; it does not normalize or reverse a log transform.
"""
import argparse
from pathlib import Path

import anndata as ad
import numpy as np
import pandas as pd
from scipy.sparse import csr_matrix


def read_table(path):
    return pd.read_csv(path, sep="\t" if Path(path).suffix.lower() in {".tsv", ".txt"} else ",", index_col=0)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source")
    parser.add_argument("output")
    parser.add_argument("--gene-rows", action="store_true")
    parser.add_argument("--annotations")
    args = parser.parse_args()
    output = Path(args.output)
    if output.exists():
        parser.error("Output already exists; choose a new filename to protect existing data.")
    if output.suffix != ".h5ad":
        parser.error("Output filename must end in .h5ad.")
    frame = read_table(args.source)
    if args.gene_rows: frame = frame.T
    for axis, values in [("entity", frame.index), ("gene", frame.columns)]:
        values = values.astype(str)
        if not values.is_unique or any(not x.strip() or x != x.strip() for x in values):
            parser.error(f"{axis} IDs must be unique, nonempty and have no surrounding whitespace.")
    numeric = frame.apply(pd.to_numeric, errors="raise").to_numpy(dtype=np.float32)
    if not numeric.size or not (np.isfinite(numeric) & (numeric >= 0)).all():
        parser.error("Expression must be nonempty, finite and nonnegative.")
    obs = pd.DataFrame(index=frame.index.astype(str))
    var = pd.DataFrame(index=frame.columns.astype(str))
    obs.index.name = None
    var.index.name = None
    if args.annotations:
        annotations = read_table(args.annotations)
        annotations.index = annotations.index.astype(str)
        if not annotations.index.is_unique or "cell_type" not in annotations:
            parser.error("Reference annotations need unique cell IDs and a cell_type column.")
        annotations = annotations.reindex(obs.index)
        if annotations.cell_type.isna().any():
            parser.error("Every reference cell must have a matched cell_type annotation.")
        obs = annotations.copy()
        obs.index.name = None
        obs["cell_type"] = obs.cell_type.astype("category")
    result = ad.AnnData(X=csr_matrix(numeric), obs=obs, var=var)
    result.uns["preprocessing_note"] = "Values preserved; select a model compatible with the source expression layer."
    result.write_h5ad(output, compression="gzip")
    print(f"Saved {output}: {result.n_obs} entities x {result.n_vars} genes")


if __name__ == "__main__": main()