"""
Code for evaluating a model's ability to generalize to cells that it wasn't trained on.
Can only be used to evalute within a species.
Generates raw predictions of data modality transfer, and optionally, plots.
"""

import os
import sys
from typing import *
import functools
import logging
import argparse
import copy

import scipy

import anndata as ad
import scanpy as sc

import torch
import skorch

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 loss_functions
import model_utils
import plot_utils
import adata_utils
import utils
from models import autoencoders

DATA_DIR = os.path.join(os.path.dirname(SRC_DIR), "data")
assert os.path.isdir(DATA_DIR)

logging.basicConfig(level=logging.INFO)

DATASET_NAME = ""


def do_evaluation_rna_from_rna(
    spliced_net,
    sc_dual_full_dataset,
    gene_names: str,
    atac_names: str,
    outdir: str,
    ext: str,
    marker_genes: List[str],
    prefix: str = "",
):
    """
    Evaluate the given network on the dataset
    """
    # Do inference and plotting
    ### RNA > RNA
    logging.info("Inferring RNA from RNA...")
    sc_rna_full_preds = spliced_net.translate_1_to_1(sc_dual_full_dataset)
    sc_rna_full_preds_anndata = sc.AnnData(
        sc_rna_full_preds,
        obs=sc_dual_full_dataset.dataset_x.data_raw.obs,
    )
    sc_rna_full_preds_anndata.var_names = gene_names

    logging.info("Writing RNA from RNA")
    sc_rna_full_preds_anndata.write(
        os.path.join(outdir, f"{prefix}_rna_rna_adata.h5ad".strip("_"))
    )
    if hasattr(sc_dual_full_dataset.dataset_x, "size_norm_counts") and ext is not None:
        logging.info("Plotting RNA from RNA")
        plot_utils.plot_scatter_with_r(
            sc_dual_full_dataset.dataset_x.size_norm_counts.X,
            sc_rna_full_preds,
            one_to_one=True,
            logscale=True,
            density_heatmap=True,
            title=f"{DATASET_NAME} RNA > RNA".strip(),
            fname=os.path.join(outdir, f"{prefix}_rna_rna_log.{ext}".strip("_")),
        )


def do_evaluation_atac_from_rna(
    spliced_net,
    sc_dual_full_dataset,
    gene_names: str,
    atac_names: str,
    outdir: str,
    ext: str,
    marker_genes: List[str],
    prefix: str = "",
):
    ### RNA > ATAC
    logging.info("Inferring ATAC from RNA")
    sc_rna_atac_full_preds = spliced_net.translate_1_to_2(sc_dual_full_dataset)
    sc_rna_atac_full_preds_anndata = sc.AnnData(
        scipy.sparse.csr_matrix(sc_rna_atac_full_preds),
        obs=sc_dual_full_dataset.dataset_x.data_raw.obs,
    )
    sc_rna_atac_full_preds_anndata.var_names = atac_names
    logging.info("Writing ATAC from RNA")
    sc_rna_atac_full_preds_anndata.write(
        os.path.join(outdir, f"{prefix}_rna_atac_adata.h5ad".strip("_"))
    )

    if hasattr(sc_dual_full_dataset.dataset_y, "data_raw") and ext is not None:
        logging.info("Plotting ATAC from RNA")
        plot_utils.plot_auroc(
            utils.ensure_arr(sc_dual_full_dataset.dataset_y.data_raw.X).flatten(),
            utils.ensure_arr(sc_rna_atac_full_preds).flatten(),
            title_prefix=f"{DATASET_NAME} RNA > ATAC".strip(),
            fname=os.path.join(outdir, f"{prefix}_rna_atac_auroc.{ext}".strip("_")),
        )
        # plot_utils.plot_auprc(
        #     utils.ensure_arr(sc_dual_full_dataset.dataset_y.data_raw.X).flatten(),
        #     utils.ensure_arr(sc_rna_atac_full_preds),
        #     title_prefix=f"{DATASET_NAME} RNA > ATAC".strip(),
        #     fname=os.path.join(outdir, f"{prefix}_rna_atac_auprc.{ext}".strip("_")),
        # )


def do_evaluation_atac_from_atac(
    spliced_net,
    sc_dual_full_dataset,
    gene_names: str,
    atac_names: str,
    outdir: str,
    ext: str,
    marker_genes: List[str],
    prefix: str = "",
):
    ### ATAC > ATAC
    logging.info("Inferring ATAC from ATAC")
    sc_atac_full_preds = spliced_net.translate_2_to_2(sc_dual_full_dataset)
    sc_atac_full_preds_anndata = sc.AnnData(
        sc_atac_full_preds,
        obs=sc_dual_full_dataset.dataset_y.data_raw.obs.copy(deep=True),
    )
    sc_atac_full_preds_anndata.var_names = atac_names
    logging.info("Writing ATAC from ATAC")

    # Infer marker bins
    # logging.info("Getting marker bins for ATAC from ATAC")
    # plot_utils.preprocess_anndata(sc_atac_full_preds_anndata)
    # adata_utils.find_marker_genes(sc_atac_full_preds_anndata)
    # inferred_marker_bins = adata_utils.flatten_marker_genes(
    #     sc_atac_full_preds_anndata.uns["rank_genes_leiden"]
    # )
    # logging.info(f"Found {len(inferred_marker_bins)} marker bins for ATAC from ATAC")
    # with open(
    #     os.path.join(outdir, f"{prefix}_atac_atac_marker_bins.txt".strip("_")), "w"
    # ) as sink:
    #     sink.write("\n".join(inferred_marker_bins) + "\n")

    sc_atac_full_preds_anndata.write(
        os.path.join(outdir, f"{prefix}_atac_atac_adata.h5ad".strip("_"))
    )
    if hasattr(sc_dual_full_dataset.dataset_y, "data_raw") and ext is not None:
        logging.info("Plotting ATAC from ATAC")
        plot_utils.plot_auroc(
            utils.ensure_arr(sc_dual_full_dataset.dataset_y.data_raw.X).flatten(),
            utils.ensure_arr(sc_atac_full_preds).flatten(),
            title_prefix=f"{DATASET_NAME} ATAC > ATAC".strip(),
            fname=os.path.join(outdir, f"{prefix}_atac_atac_auroc.{ext}".strip("_")),
        )
        # plot_utils.plot_auprc(
        #     utils.ensure_arr(sc_dual_full_dataset.dataset_y.data_raw.X).flatten(),
        #     utils.ensure_arr(sc_atac_full_preds).flatten(),
        #     title_prefix=f"{DATASET_NAME} ATAC > ATAC".strip(),
        #     fname=os.path.join(outdir, f"{prefix}_atac_atac_auprc.{ext}".strip("_")),
        # )

    # Remove some objects to free memory
    del sc_atac_full_preds
    del sc_atac_full_preds_anndata


def do_evaluation_rna_from_atac(
    spliced_net,
    sc_dual_full_dataset,
    gene_names: str,
    atac_names: str,
    outdir: str,
    ext: str,
    marker_genes: List[str],
    prefix: str = "",
):
    ### ATAC > RNA
    logging.info("Inferring RNA from ATAC")
    sc_atac_rna_full_preds = spliced_net.translate_2_to_1(sc_dual_full_dataset)
    # Seurat expects everything to be sparse
    # https://github.com/satijalab/seurat/issues/2228
    sc_atac_rna_full_preds_anndata = sc.AnnData(
        sc_atac_rna_full_preds,
        obs=sc_dual_full_dataset.dataset_y.data_raw.obs.copy(deep=True),
    )
    sc_atac_rna_full_preds_anndata.var_names = gene_names
    logging.info("Writing RNA from ATAC")

    # Seurat also expects the raw attribute to be populated
    sc_atac_rna_full_preds_anndata.raw = sc_atac_rna_full_preds_anndata.copy()
    sc_atac_rna_full_preds_anndata.write(
        os.path.join(outdir, f"{prefix}_atac_rna_adata.h5ad".strip("_"))
    )
    # sc_atac_rna_full_preds_anndata.write_csvs(
    #     os.path.join(outdir, f"{prefix}_atac_rna_constituent_csv".strip("_")),
    #     skip_data=False,
    # )
    # sc_atac_rna_full_preds_anndata.to_df().to_csv(
    #     os.path.join(outdir, f"{prefix}_atac_rna_table.csv".strip("_"))
    # )

    # If there eixsts a ground truth RNA, do RNA plotting
    if hasattr(sc_dual_full_dataset.dataset_x, "size_norm_counts") and ext is not None:
        logging.info("Plotting RNA from ATAC")
        plot_utils.plot_scatter_with_r(
            sc_dual_full_dataset.dataset_x.size_norm_counts.X,
            sc_atac_rna_full_preds,
            one_to_one=True,
            logscale=True,
            density_heatmap=True,
            title=f"{DATASET_NAME} ATAC > RNA".strip(),
            fname=os.path.join(outdir, f"{prefix}_atac_rna_log.{ext}".strip("_")),
        )

    # Remove objects to free memory
    del sc_atac_rna_full_preds
    del sc_atac_rna_full_preds_anndata


def do_latent_evaluation(
    spliced_net, sc_dual_full_dataset, outdir: str, prefix: str = ""
):
    """
    Pull out latent space and write to file
    """
    logging.info("Inferring latent representations")
    encoded_from_rna, encoded_from_atac = spliced_net.get_encoded_layer(
        sc_dual_full_dataset
    )

    if hasattr(sc_dual_full_dataset.dataset_x, "data_raw"):
        encoded_from_rna_adata = sc.AnnData(
            encoded_from_rna,
            obs=sc_dual_full_dataset.dataset_x.data_raw.obs.copy(deep=True),
        )
        encoded_from_rna_adata.write(
            os.path.join(outdir, f"{prefix}_rna_encoded_adata.h5ad".strip("_"))
        )
    if hasattr(sc_dual_full_dataset.dataset_y, "data_raw"):
        encoded_from_atac_adata = sc.AnnData(
            encoded_from_atac,
            obs=sc_dual_full_dataset.dataset_y.data_raw.obs.copy(deep=True),
        )
        encoded_from_atac_adata.write(
            os.path.join(outdir, f"{prefix}_atac_encoded_adata.h5ad".strip("_"))
        )


def infer_reader(fname: str, mode: str = "atac") -> Callable:
    """Given a filename, infer the correct reader to use"""
    assert mode in ["atac", "rna"], f"Unrecognized mode: {mode}"
    if fname.endswith(".h5"):
        if mode == "atac":
            return functools.partial(utils.sc_read_10x_h5_ft_type, ft_type="Peaks")
        else:
            return utils.sc_read_10x_h5_ft_type
    elif fname.endswith(".h5ad"):
        return ad.read_h5ad
    else:
        raise ValueError(f"Unrecognized extension: {fname}")


def build_parser():
    parser = argparse.ArgumentParser(
        usage=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        nargs="*",
        required=False,
        default=[
            os.path.join(model_utils.MODEL_CACHE_DIR, "cv_logsplit_01_model_only")
        ],
        help="Checkpoint directory to load model from. If not given, automatically download and use a human pretrained model",
    )
    parser.add_argument("--prefix", type=str, default="net_", help="Checkpoint prefix")
    parser.add_argument("--data", required=True, nargs="*", help="Data files")
    parser.add_argument(
        "--dataname", default="", help="Name of dataset to include in plot titles"
    )
    parser.add_argument(
        "--outdir", type=str, required=True, help="Output directory for files and plots"
    )
    parser.add_argument(
        "--genes",
        type=str,
        default="",
        help="Genes that the model uses (inferred based on checkpoint dir if not given)",
    )
    parser.add_argument(
        "--bins",
        type=str,
        default="",
        help="ATAC bins that the model uses (inferred based on checkpoint dir if not given)",
    )
    parser.add_argument(
        "--liftHg19toHg38",
        action="store_true",
        help="Liftover input ATAC bins from hg19 to hg38",
    )
    parser.add_argument("--device", type=str, default="0", help="Device to use")
    parser.add_argument(
        "--ext",
        type=str,
        default="pdf",
        choices=["pdf", "png", "jpg"],
        help="File format to use for plotting",
    )
    parser.add_argument(
        "--noplot", action="store_true", help="Disable plotting, writing output only"
    )
    parser.add_argument(
        "--transonly",
        action="store_true",
        help="Disable doing same-modality inference",
    )
    parser.add_argument(
        "--skiprnasource", action="store_true", help="Skip analysis starting from RNA"
    )
    parser.add_argument(
        "--skipatacsource", action="store_true", help="Skip analysis starting from ATAC"
    )
    parser.add_argument(
        "--nofilter",
        action="store_true",
        help="Whether or not to perform filtering (note that we always discard cells with no expressed genes)",
    )
    return parser


def load_rna_files_for_eval(
    data, checkpoint: str, rna_genes_list_fname: str = "", no_filter: bool = False
):
    """ """
    if not rna_genes_list_fname:
        rna_genes_list_fname = os.path.join(checkpoint, "rna_genes.txt")
    assert os.path.isfile(
        rna_genes_list_fname
    ), f"Cannot find RNA genes file: {rna_genes_list_fname}"
    rna_genes = utils.read_delimited_file(rna_genes_list_fname)
    rna_data_kwargs = copy.copy(sc_data_loaders.TENX_PBMC_RNA_DATA_KWARGS)
    if no_filter:
        rna_data_kwargs = {
            k: v for k, v in rna_data_kwargs.items() if not k.startswith("filt_")
        }
        # Always discard cells with no expressed genes
        rna_data_kwargs["filt_cell_min_genes"] = 1
    rna_data_kwargs["fname"] = data
    reader_func = functools.partial(
        utils.sc_read_multi_files,
        reader=lambda x: sc_data_loaders.repool_genes(
            utils.get_ad_reader(x, ft_type="Gene Expression")(x), rna_genes
        ),
    )
    rna_data_kwargs["reader"] = reader_func
    try:
        logging.info(f"Building RNA dataset with parameters: {rna_data_kwargs}")
        sc_rna_full_dataset = sc_data_loaders.SingleCellDataset(
            mode="skip",
            **rna_data_kwargs,
        )
        assert all(
            [x == y for x, y in zip(rna_genes, sc_rna_full_dataset.data_raw.var_names)]
        ), "Mismatched genes"
        _temp = sc_rna_full_dataset[0]  # Try that query works
        # adata_utils.find_marker_genes(sc_rna_full_dataset.data_raw, n_genes=25)
        # marker_genes = adata_utils.flatten_marker_genes(
        #     sc_rna_full_dataset.data_raw.uns["rank_genes_leiden"]
        # )
        marker_genes = []
        # Write out the truth
    except (AssertionError, IndexError) as e:
        logging.warning(f"Error when reading RNA gene expression data from {data}: {e}")
        logging.warning("Ignoring RNA data")
        # Update length later
        sc_rna_full_dataset = sc_data_loaders.DummyDataset(
            shape=len(rna_genes), length=-1
        )
        marker_genes = []
    return sc_rna_full_dataset, rna_genes, marker_genes


def load_atac_files_for_eval(
    data: List[str],
    checkpoint: str,
    atac_bins_list_fname: str = "",
    lift_hg19_to_hg39: bool = False,
    predefined_split=None,
):
    """Load the ATAC files for evaluation"""
    if not atac_bins_list_fname:
        atac_bins_list_fname = os.path.join(checkpoint, "atac_bins.txt")
        logging.info(f"Auto-set atac bins fname to {atac_bins_list_fname}")
    assert os.path.isfile(
        atac_bins_list_fname
    ), f"Cannot find ATAC bins file: {atac_bins_list_fname}"
    atac_bins = utils.read_delimited_file(
        atac_bins_list_fname
    )  # These are the bins we are using (i.e. the bins the model was trained on)
    atac_data_kwargs = copy.copy(sc_data_loaders.TENX_PBMC_ATAC_DATA_KWARGS)
    atac_data_kwargs["fname"] = data
    atac_data_kwargs["cluster_res"] = 0  # Disable clustering
    filt_atac_keys = [k for k in atac_data_kwargs.keys() if k.startswith("filt")]
    for k in filt_atac_keys:  # Reset filtering
        atac_data_kwargs[k] = None
    atac_data_kwargs["pool_genomic_interval"] = atac_bins
    if not lift_hg19_to_hg39:
        atac_data_kwargs["reader"] = functools.partial(
            utils.sc_read_multi_files,
            reader=lambda x: sc_data_loaders.repool_atac_bins(
                infer_reader(data[0], mode="atac")(x),
                atac_bins,
            ),
        )
    else:  # Requires liftover
        # Read, liftover, then repool
        atac_data_kwargs["reader"] = functools.partial(
            utils.sc_read_multi_files,
            reader=lambda x: sc_data_loaders.repool_atac_bins(
                sc_data_loaders.liftover_atac_adata(
                    # utils.sc_read_10x_h5_ft_type(x, "Peaks")
                    infer_reader(data[0], mode="atac")(x)
                ),
                atac_bins,
            ),
        )

    try:
        sc_atac_full_dataset = sc_data_loaders.SingleCellDataset(
            mode="skip",
            predefined_split=predefined_split if predefined_split else None,
            **atac_data_kwargs,
        )
        _temp = sc_atac_full_dataset[0]  # Try that query works
        assert all(
            [x == y for x, y in zip(atac_bins, sc_atac_full_dataset.data_raw.var_names)]
        )
    except AssertionError as err:
        logging.warning(f"Error when reading ATAC data from {data}: {err}")
        logging.warning("Ignoring ATAC data, returning dummy dataset instead")
        sc_atac_full_dataset = sc_data_loaders.DummyDataset(
            shape=len(atac_bins), length=-1
        )
    return sc_atac_full_dataset, atac_bins


def main():
    parser = build_parser()
    args = parser.parse_args()
    logging.info(f"Evaluating: {' '.join(args.data)}")

    global DATASET_NAME
    DATASET_NAME = args.dataname

    # Create output directory
    if not os.path.isdir(args.outdir):
        os.makedirs(args.outdir)

    # Set up logging
    logger = logging.getLogger()
    fh = logging.FileHandler(os.path.join(args.outdir, "logging.log"), "w")
    fh.setLevel(logging.INFO)
    logger.addHandler(fh)

    if args.checkpoint[0] == os.path.join(
        model_utils.MODEL_CACHE_DIR, "cv_logsplit_01_model_only"
    ):
        _ = model_utils.load_model()  # Downloads if not downloaded
    (sc_rna_full_dataset, rna_genes, marker_genes,) = load_rna_files_for_eval(
        args.data, args.checkpoint[0], args.genes, no_filter=args.nofilter
    )

    if hasattr(sc_rna_full_dataset, "size_norm_counts"):
        logging.info("Writing truth RNA size normalized counts")
        sc_rna_full_dataset.size_norm_counts.write_h5ad(
            os.path.join(args.outdir, "truth_rna.h5ad")
        )

    sc_atac_full_dataset, atac_bins = load_atac_files_for_eval(
        args.data,
        args.checkpoint[0],
        args.bins,
        args.liftHg19toHg38,
        sc_rna_full_dataset if hasattr(sc_rna_full_dataset, "data_raw") else None,
    )
    # Write out the truth
    if hasattr(sc_atac_full_dataset, "data_raw"):
        logging.info("Writing truth ATAC binary counts")
        sc_atac_full_dataset.data_raw.write_h5ad(
            os.path.join(args.outdir, "truth_atac.h5ad")
        )

    if isinstance(sc_rna_full_dataset, sc_data_loaders.DummyDataset) and isinstance(
        sc_atac_full_dataset, sc_data_loaders.DummyDataset
    ):
        raise ValueError("Cannot proceed with two dummy datasets for both RNA and ATAC")
    # Update the RNA counts if we do not actually have RNA data
    if isinstance(sc_rna_full_dataset, sc_data_loaders.DummyDataset) and not isinstance(
        sc_atac_full_dataset, sc_data_loaders.DummyDataset
    ):
        sc_rna_full_dataset.length = len(sc_atac_full_dataset)
    elif isinstance(
        sc_atac_full_dataset, sc_data_loaders.DummyDataset
    ) and not isinstance(sc_rna_full_dataset, sc_data_loaders.DummyDataset):
        sc_atac_full_dataset.length = len(sc_rna_full_dataset)

    # Build the dual combined dataset
    sc_dual_full_dataset = sc_data_loaders.PairedDataset(
        sc_rna_full_dataset,
        sc_atac_full_dataset,
        flat_mode=True,
    )

    # Write some basic outputs related to variable and obs names
    with open(os.path.join(args.outdir, "rna_genes.txt"), "w") as sink:
        sink.write("\n".join(rna_genes) + "\n")
    with open(os.path.join(args.outdir, "atac_bins.txt"), "w") as sink:
        sink.write("\n".join(atac_bins) + "\n")
    with open(os.path.join(args.outdir, "obs_names.txt"), "w") as sink:
        sink.write("\n".join(sc_dual_full_dataset.obs_names))

    for i, ckpt in enumerate(args.checkpoint):
        # Dynamically determine the model we are looking at based on name
        checkpoint_basename = os.path.basename(ckpt)
        if checkpoint_basename.startswith("naive"):
            logging.info(f"Inferred model to be naive")
            model_class = autoencoders.NaiveSplicedAutoEncoder
        else:
            logging.info(f"Inferred model to be normal (non-naive)")
            model_class = autoencoders.AssymSplicedAutoEncoder

        prefix = "" if len(args.checkpoint) == 1 else f"model_{checkpoint_basename}"
        spliced_net = model_utils.load_model(
            ckpt,
            prefix=args.prefix,
            device=args.device,
        )

        do_latent_evaluation(
            spliced_net=spliced_net,
            sc_dual_full_dataset=sc_dual_full_dataset,
            outdir=args.outdir,
            prefix=prefix,
        )

        if (
            isinstance(sc_rna_full_dataset, sc_data_loaders.SingleCellDataset)
            and not args.skiprnasource
        ):
            if not args.transonly:
                do_evaluation_rna_from_rna(
                    spliced_net,
                    sc_dual_full_dataset,
                    rna_genes,
                    atac_bins,
                    args.outdir,
                    None if args.noplot else args.ext,
                    marker_genes,
                    prefix=prefix,
                )
            do_evaluation_atac_from_rna(
                spliced_net,
                sc_dual_full_dataset,
                rna_genes,
                atac_bins,
                args.outdir,
                None if args.noplot else args.ext,
                marker_genes,
                prefix=prefix,
            )
        if (
            isinstance(sc_atac_full_dataset, sc_data_loaders.SingleCellDataset)
            and not args.skipatacsource
        ):
            do_evaluation_rna_from_atac(
                spliced_net,
                sc_dual_full_dataset,
                rna_genes,
                atac_bins,
                args.outdir,
                None if args.noplot else args.ext,
                marker_genes,
                prefix=prefix,
            )
            if not args.transonly:
                do_evaluation_atac_from_atac(
                    spliced_net,
                    sc_dual_full_dataset,
                    rna_genes,
                    atac_bins,
                    args.outdir,
                    None if args.noplot else args.ext,
                    marker_genes,
                    prefix=prefix,
                )
        del spliced_net


if __name__ == "__main__":
    main()