# mypy: ignore-errors
"""Read cell embeddings from a h5ad generated by `transcriptformer inference` and write them into
an existing tiledbsoma.SparseNDArray.
"""

import argparse
import logging
import os

import anndata as ad
import numpy as np
import pyarrow as pa
import tiledb
import tiledbsoma
from scipy.sparse import coo_matrix


def main():
    parser = argparse.ArgumentParser(description="Put TranscriptFormer embeddings into SOMASparseNDArray.")
    parser.add_argument(
        "h5ad",
        type=str,
    )
    parser.add_argument(
        "array",
        type=str,
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="[%(asctime)s][put_embeddings][%(levelname)s] - %(message)s",
    )
    logging.info(f"loaded tiledbsoma=={tiledbsoma.__version__} tiledb=={tiledb.__version__}")

    # read the embeddings from the h5ad file
    adata = ad.read_h5ad(args.h5ad)
    soma_joinid = adata.obs["soma_joinid"].values
    embeddings = adata.obsm["embeddings"]
    n = embeddings.shape[0]
    embedding_dim = embeddings.shape[1]
    logging.info(f"read {embeddings.shape} embeddings from {args.h5ad}")

    # write them to the tiledbsoma.SparseNDArray
    tiledbsoma_context = tiledbsoma.options.SOMATileDBContext(
        tiledb_config={
            "py.init_buffer_bytes": 1024**3,
            "soma.init_buffer_bytes": 1024**3,
            "vfs.s3.region": os.environ["AWS_DEFAULT_REGION"],
        }
    )
    with tiledbsoma.SparseNDArray.open(args.array, "w", context=tiledbsoma_context) as array:
        dim0 = np.repeat(soma_joinid, embedding_dim)
        dim1 = np.tile(np.arange(embedding_dim), n)
        data = embeddings.flatten()
        array.write(pa.SparseCOOTensor.from_scipy(coo_matrix((data, (dim0, dim1)))))
    logging.info(f"wrote embeddings to tiledbsoma.SparseNDArray {args.array}")


if __name__ == "__main__":
    main()
