"""

Script for calculating pairwise distances between 2 adatas

"""

import os

import sys

import logging

import argparse

import json

import numpy as np

from evaluate_bulk_rna_concordance import load_file_flex_format

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 adata_utils

import metrics

logging.basicConfig(level=logging.INFO)

def build_parser():

    """Build CLI parser"""

    parser = argparse.ArgumentParser(

        usage=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter

    )

    parser.add_argument("adata1", type=str, help="First adata object to compare")

    parser.add_argument("adata2", type=str, help="Second adata object to compare")

    parser.add_argument(

        "--output", "-o", type=str, default="", help="Json file to write matches to"

    )

    parser.add_argument(

        "--method", "-m", type=str, choices=["euclidean", "cosine"], default="euclidean"

    )

    parser.add_argument(

        "--log",

        "-l",

        action="store_true",

        help="Log transform before computing distance",

    )

    parser.add_argument(

        "--numtop",

        "-n",

        type=int,

        default=0,

        help="Number of top matches to report. 0 indicates reporting all in descending order",

    )

    return parser

def main():

    """Run script"""

    parser = build_parser()

    args = parser.parse_args()

    # Compute distances

    x = load_file_flex_format(args.adata1)

    logging.info(f"Loaded in {args.adata1} for {x.shape}")

    y = load_file_flex_format(args.adata2)

    logging.info(f"Loaded in {args.adata2} for {y.shape}")

    # Log, because often times in this project the output is actually linear space

    # and comparing expression is typically done in log space

    if args.log:

        logging.info("Log transforming inputs")

        x.X = np.log1p(x.X)

        y.X = np.log1p(y.X)

    pairwise_dist = adata_utils.evaluate_pairwise_cell_distance(

        x, y, method=args.method

    )

    if args.numtop == 0:

        args.numtop = y.n_obs

    # Figure out the top few matches per cell

    matches = {}

    for i, row in pairwise_dist.iterrows():

        sorted_idx = np.argsort(row.values)

        cell_matches = pairwise_dist.columns[sorted_idx[: args.numtop]]

        matches[i] = list(cell_matches)

    if args.output:

        assert args.output.endswith(".json")

        with open(args.output, "w") as sink:

            logging.info(f"Writing matches to {args.output}")

            json.dump(matches, sink, indent=4)

    # Report Top N accuracy if relevant

    if x.n_obs == y.n_obs and np.all(x.obs_names == y.obs_names):

        n = args.numtop if args.numtop < y.n_obs else 10

        acc = metrics.top_n_accuracy(matches.values(), matches.keys(), n=n)

        logging.info(f"Top {n} accuracy: {acc:.4f}")

if __name__ == "__main__":

    main()