"""

Plot the variance of a gene versus how accurately we can predict it

"""

import os

import sys

import argparse

import logging

import itertools

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 interpretation

import adata_utils

import plot_utils

import utils

from evaluate_bulk_rna_concordance import load_file_flex_format

REF_MARKER_GENES = {

    "PBMC": set(

        itertools.chain.from_iterable(interpretation.PBMC_MARKER_GENES.values())

    ),

    "PBMC_Seurat": set(

        itertools.chain.from_iterable(interpretation.SEURAT_PBMC_MARKER_GENES.values())

    ),

    "Housekeeper": utils.read_delimited_file(

        os.path.join(os.path.dirname(SRC_DIR), "data", "housekeeper_genes.txt")

    ),

}

def build_parser():

    """Build basic CLI parser"""

    parser = argparse.ArgumentParser(

        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter

    )

    parser.add_argument("preds", type=str, help="File with predicted expression")

    parser.add_argument("truth", type=str, help="File with ground truth expression")

    parser.add_argument("plotname", type=str, help="File to write plot to")

    parser.add_argument(

        "--genelist",

        "-g",

        type=str,

        default="",

        help="File to write outliers in expained variance",

    )

    parser.add_argument(

        "--highlight",

        choices=REF_MARKER_GENES.keys(),

        nargs="*",

        default=["Housekeeper"],

        help="HGenes to highlight",

    )

    parser.add_argument(

        "--linear", action="store_true", help="Plot in linear space instead of log"

    )

    parser.add_argument(

        "--unconstriained", action="store_true", help="Do not constrain axes"

    )

    parser.add_argument("--outliers", action="store_true", help="Label outliers")

    return parser

def main():

    """Run script"""

    parser = build_parser()

    args = parser.parse_args()

    truth = load_file_flex_format(args.truth)

    truth.X = utils.ensure_arr(truth.X)

    logging.info(f"Loaded truth {args.truth}: {truth.shape}")

    preds = load_file_flex_format(args.preds)

    preds.X = utils.ensure_arr(preds.X)

    logging.info(f"Loaded preds {args.preds}: {preds.shape}")

    common_genes = sorted(list(set(truth.var_names).intersection(preds.var_names)))

    logging.info(f"Shared genes: {len(common_genes)}")

    common_obs = sorted(list(set(truth.obs_names).intersection(preds.obs_names)))

    # All obs naames should intersect between preds and truth

    assert len(common_obs) == len(truth.obs_names) == len(preds.obs_names)

    plot_utils.plot_var_vs_explained_var(

        truth,

        preds,

        highlight_genes={k: REF_MARKER_GENES[k] for k in args.highlight},

        logscale=not args.linear,

        constrain_y_axis=not args.unconstriained,

        label_outliers=args.outliers,

        fname=args.plotname,

        fname_gene_list=args.genelist,

    )

if __name__ == "__main__":

    main()