"""
Script for calculating gene activity scores
"""

import os
import sys
import logging
import argparse

import numpy as np
import pandas as pd
import anndata as ad
import scanpy as sc
import scipy

SRC_DIR = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "babel",
)
assert os.path.isdir(SRC_DIR)
sys.path.append(SRC_DIR)
import sc_data_loaders
import adata_utils
import metrics
import atac_utils
import utils

logging.basicConfig(level=logging.INFO)


ANNOTATION_DICT = {
    "hg19": sc_data_loaders.HG19_GTF,
    "hg38": sc_data_loaders.HG38_GTF,
}


def build_parser():
    """Build parser for running as a script"""
    parser = argparse.ArgumentParser(
        usage="Convert ATAC h5ad to infered gene activity scores",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "input_h5ad",
        type=str,
        nargs="*",
        help="ATAC h5ad to infer gene activity scores for",
    )
    parser.add_argument(
        "output_h5ad", type=str, help="h5ad file to write gene activity scores to"
    )
    parser.add_argument(
        "--genome",
        "-g",
        choices=ANNOTATION_DICT.keys(),
        default="hg38",
        help="Genome annotation to use",
    )
    parser.add_argument("--raw", action="store_true", help="Use raw atribaute of input")
    parser.add_argument(
        "--sizenorm",
        action="store_true",
        help="Normalize gene activity scores by span of gene",
    )
    parser.add_argument(
        "--naive",
        action="store_true",
        help="Use naive method instead of archr-derived method",
    )
    return parser


def main():
    """On the fly testing"""
    parser = build_parser()
    args = parser.parse_args()
    logging.info(f"Inputs: {args.input_h5ad}")

    # Note that all inputs should have the same suffix
    for input_h5ad in args.input_h5ad:
        assert input_h5ad.endswith(".h5ad") or input_h5ad.endswith(".h5")
    assert args.output_h5ad.endswith(".h5ad")
    logging.info(f"Calculating gene activity scores for: {args.input_h5ad}")

    reader_func = (
        ad.read_h5ad
        if args.input_h5ad[0].endswith(".h5ad")
        else lambda x: utils.sc_read_10x_h5_ft_type(x, "Peaks")
    )
    # adata = reader_func(args.input_h5ad)
    adata = utils.sc_read_multi_files(args.input_h5ad, reader=reader_func)
    logging.info(f"Received input of {adata.shape}")
    if args.naive:
        gene_act_adata = atac_utils.gene_activity_matrix_from_adata(
            adata if not args.raw else adata.raw,
            size_norm=args.sizenorm,
            annotation=ANNOTATION_DICT[args.genome],
        )
    else:
        logging.warning("Using ArchR based method, size norm is always on")
        gene_act_adata = atac_utils.archr_gene_activity_matrix_from_adata(
            adata if not args.raw else adata.raw,
            annotation=ANNOTATION_DICT[args.genome],
        )
    logging.info(f"Writing gene activity scores to {args.output_h5ad}")
    gene_act_adata.write(args.output_h5ad)


if __name__ == "__main__":
    main()