#!/usr/bin/env python3
# mypy: ignore-errors

import argparse
import logging
import os
import sys
import tempfile

import geneformer
import torch
from datasets import disable_progress_bar
from transformers import BertConfig

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(module)s [%(levelname)s] %(message)s")
logger = logging.getLogger(os.path.basename(__file__))
disable_progress_bar()


def main(argv):
    assert torch.cuda.is_available(), "CUDA is not available"
    args = parse_arguments(argv)

    tiledbsoma_context = None
    if args.tiledbsoma:
        # prep tiledbsoma (and fail fast if there's a problem)
        import boto3
        import tiledb
        import tiledbsoma

        logger.info(f"loaded tiledbsoma=={tiledbsoma.__version__} tiledb=={tiledb.__version__}")

        aws_region = "us-west-2"
        try:
            aws_region = boto3.Session().region_name
        except Exception:  # noqa: BLE001
            pass
        tiledbsoma_context = tiledbsoma.options.SOMATileDBContext(
            tiledb_config={
                "py.init_buffer_bytes": 4 * 1024**3,
                "soma.init_buffer_bytes": 4 * 1024**3,
                "vfs.s3.region": aws_region,
            }
        )

        with tiledbsoma.SparseNDArray.open(args.outfile, "r", context=tiledbsoma_context):
            # TODO: verify schema compatibility
            pass

    num_classes = 0
    if args.model_type != "Pretrained":
        # classifier model: detect the number of classes from the model configuration (without
        # having to load in the whole model redundantly with EmbExtractor below)
        num_classes = BertConfig.from_pretrained(args.model).num_labels
        logger.info(f"detected {num_classes} labels in {args.model_type} model {args.model}")

    with tempfile.TemporaryDirectory() as scratch_dir:
        # prepare the dataset, taking only one shard of it if so instructed
        logger.info("Extracting embeddings...")
        emb_kwargs = {
            "model_type": args.model_type,
            "num_classes": num_classes,
            "max_ncells": None,
            "emb_mode": args.emb_mode,
            "emb_layer": args.emb_layer,
            "emb_label": args.features,
            "forward_batch_size": args.batch_size,
        }
        if args.token_dictionary:
            emb_kwargs["token_dictionary_file"] = args.token_dictionary
        extractor = geneformer.EmbExtractor(**emb_kwargs)
        # NOTE: EmbExtractor will use only one GPU
        #       see https://huggingface.co/ctheodoris/Geneformer/blob/main/geneformer/emb_extractor.py
        embs_df = extractor.extract_embs(
            model_directory=args.model,
            input_data_file=args.dataset,
            # the method always writes out a .csv file which we discard (since it also returns the
            # embeddings data frame)
            output_directory=scratch_dir,
            output_prefix="embs",
        )

        if args.tiledbsoma:
            import numpy as np
            import pyarrow as pa
            from scipy.sparse import coo_matrix

            # NOTE: embs_df has columns -named- 0, 1, ..., 511 as well as the requested features.
            embedding_dim = embs_df.shape[1] - len(args.features)
            logger.info(f"writing to tiledbsoma.SparseNDArray at {args.outfile}, embedding_dim={embedding_dim}...")
            with tiledbsoma.SparseNDArray.open(args.outfile, "w", context=tiledbsoma_context) as array:
                dim0 = np.repeat(embs_df["soma_joinid"].values, embedding_dim)
                dim1 = np.tile(np.arange(embedding_dim), len(embs_df))
                data = embs_df.loc[:, range(embedding_dim)].values.flatten()
                array.write(pa.SparseCOOTensor.from_scipy(coo_matrix((data, (dim0, dim1)))))
        else:
            logger.info(f"writing {args.outfile}...")
            # reorder embs_df columns and write to TSV
            cols = embs_df.columns.tolist()
            emb_cols = [col for col in cols if isinstance(col, int)]
            anno_cols = [col for col in cols if not isinstance(col, int)]
            embs_df = embs_df[anno_cols + emb_cols].set_index("soma_joinid").sort_index()
            embs_df.to_csv(args.outfile, sep="\t", header=True, index=True, index_label="soma_joinid")

        logger.info("SUCCESS")


def parse_arguments(argv):
    parser = argparse.ArgumentParser(description="Generate fine-tuned Geneformer embeddings for given cells dataset")
    parser.add_argument("model", type=str, help="path to model")
    parser.add_argument("dataset", type=str, help="saved cell Dataset path")
    parser.add_argument("outfile", type=str, help="output filename")
    parser.add_argument(
        "--model-type",
        type=str,
        choices=("CellClassifier", "Pretrained"),
        default="CellClassifier",
        help="model type (Pretrained or default CellClassifier)",
    )
    parser.add_argument(
        "--emb-mode", type=str, choices=("cell", "gene", "cls"), default="cell", help="embedding mode (default cell)"
    )
    parser.add_argument(
        "--emb-layer", type=int, choices=(-1, 0), default=-1, help="desired embedding layer (0 or default -1)"
    )
    parser.add_argument("--token-dictionary", type=str, default=None, help="custom token dictionary pickle file")
    parser.add_argument(
        "--features",
        type=str,
        default="soma_joinid,cell_type,cell_type_ontology_term_id,cell_subclass,cell_subclass_ontology_term_id",
        help="dataset features to copy into output dataframe (comma-separated)",
    )
    parser.add_argument("--batch-size", type=int, default=16, help="batch size")
    parser.add_argument(
        "--tiledbsoma",
        action="store_true",
        help="outfile is URI to an existing tiledbsoma.SparseNDArray to write into (instead of TSV file)",
    )

    args = parser.parse_args(argv[1:])

    if args.features:
        args.features = [s.strip() for s in args.features.split(",")]
    else:
        args.features = []
    if "soma_joinid" not in args.features:
        args.features.append("soma_joinid")

    logger.info("arguments: " + str(vars(args)))
    return args


if __name__ == "__main__":
    sys.exit(main(sys.argv))
