"""
Code for merging two h5 objects
"""

import os
import sys
import argparse
import logging

import anndata as ad
import scanpy as sc

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 utils


def build_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("input1", type=str)
    parser.add_argument("input2", type=str)
    parser.add_argument("output", type=str)
    parser.add_argument(
        "--mode",
        "-m",
        type=str,
        default="ATAC",
        choices=["ATAC", "RNA"],
        help="Merging RNA or ATAC data",
    )
    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()

    if args.mode == "RNA":
        raise NotImplemented
    else:
        x = utils.sc_read_10x_h5_ft_type(args.input1, ft_type="Peaks")
        y = utils.sc_read_10x_h5_ft_type(args.input2, ft_type="Peaks")

        atac_bins = sc_data_loaders.harmonize_atac_intervals(x.var_names, y.var_names)
        logging.info(f"Harmonized to {len(atac_bins)} intervals")

        x_repool = sc_data_loaders.repool_atac_bins(x, atac_bins)
        y_repool = sc_data_loaders.repool_atac_bins(y, atac_bins)

        result = x_repool.concatenate(y_repool)

    logging.info(f"Shape after concat: {result.shape}")
    if "." in os.path.basename(args.output):
        _root, ext = os.path.splitext(args.output)
        if ext == ".h5ad":
            result.write_h5ad(args.output)
        else:
            raise NotImplementedError(f"Unrecognized file extension: {ext}")
    else:
        adata_utils.write_adata_as_10x_dir(result, args.output)


if __name__ == "__main__":
    main()