"""
Script for training a protein predictor
"""

import os
import sys
import logging
import argparse
import copy
import functools
import itertools
import collections
from typing import *
import json

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

import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
import skorch
import skorch.helper

torch.backends.cudnn.deterministic = True  # For reproducibility
torch.backends.cudnn.benchmark = False

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)
MODELS_DIR = os.path.join(SRC_DIR, "models")
assert os.path.isdir(MODELS_DIR)
sys.path.append(MODELS_DIR)

import sc_data_loaders
import autoencoders
import loss_functions
import model_utils
from protein_utils import LOSS_DICT, OPTIM_DICT, ACT_DICT
import utils

from train_model import plot_loss_history

logging.basicConfig(level=logging.INFO)


def load_rna_files(
    rna_counts_fnames: List[str], model_dir: str, transpose: bool = True
) -> ad.AnnData:
    """Load the RNA files in, filling in unmeasured genes as necessary"""
    # Find the genes that the model understands
    rna_genes_list_fname = os.path.join(model_dir, "rna_genes.txt")
    assert os.path.isfile(
        rna_genes_list_fname
    ), f"Cannot find RNA genes file: {rna_genes_list_fname}"
    learned_rna_genes = utils.read_delimited_file(rna_genes_list_fname)
    assert isinstance(learned_rna_genes, list)
    assert utils.is_all_unique(
        learned_rna_genes
    ), "Learned genes list contains duplicates"

    temp_ad = utils.sc_read_multi_files(
        rna_counts_fnames,
        feature_type="Gene Expression",
        transpose=transpose,
        join="outer",
    )
    logging.info(f"Read input RNA files for {temp_ad.shape}")
    temp_ad.X = utils.ensure_arr(temp_ad.X)

    # Filter for mouse genes and remove human/mouse prefix
    temp_ad.var_names_make_unique()
    kept_var_names = [
        vname for vname in temp_ad.var_names if not vname.startswith("MOUSE_")
    ]
    if len(kept_var_names) != temp_ad.n_vars:
        temp_ad = temp_ad[:, kept_var_names]
    temp_ad.var = pd.DataFrame(index=[v.strip("HUMAN_") for v in kept_var_names])

    # Expand adata to span all genes
    # Initiating as a sparse matrix doesn't allow vectorized building
    intersected_genes = set(temp_ad.var_names).intersection(learned_rna_genes)
    assert intersected_genes, "No overlap between learned and input genes!"
    expanded_mat = np.zeros((temp_ad.n_obs, len(learned_rna_genes)))
    skip_count = 0
    for gene in intersected_genes:
        dest_idx = learned_rna_genes.index(gene)
        src_idx = temp_ad.var_names.get_loc(gene)
        if not isinstance(src_idx, int):
            logging.warn(f"Got multiple source matches for {gene}, skipping")
            skip_count += 1
            continue
        v = utils.ensure_arr(temp_ad.X[:, src_idx]).flatten()
        expanded_mat[:, dest_idx] = v
    if skip_count:
        logging.warning(
            f"Skipped {skip_count}/{len(intersected_genes)} genes due to multiple matches"
        )
    expanded_mat = sparse.csr_matrix(expanded_mat)  # Compress
    retval = ad.AnnData(
        expanded_mat, obs=temp_ad.obs, var=pd.DataFrame(index=learned_rna_genes)
    )
    return retval


def build_parser():
    """Build CLI parser"""
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--rnaCounts",
        type=str,
        nargs="*",
        required=True,
        help="file containing raw RNA counts",
    )
    parser.add_argument(
        "--proteinCounts",
        type=str,
        nargs="*",
        required=True,
        help="file containing raw protein counts",
    )
    parser.add_argument(
        "--encoder", required=True, type=str, help="Model folder to find encoder"
    )
    parser.add_argument(
        "--outdir",
        type=str,
        default=os.getcwd(),
        help="Output directory for model, defaults to current dir",
    )
    parser.add_argument(
        "--clusterres",
        type=float,
        default=1.5,
        help="Cluster resolution for train/valid/test splits",
    )
    parser.add_argument(
        "--validcluster", type=int, default=0, help="Cluster ID to use as valid cluster"
    )
    parser.add_argument(
        "--testcluster", type=int, default=1, help="Cluster ID to use as test cluster"
    )
    parser.add_argument(
        "--preprocessonly",
        action="store_true",
        help="Preprocess data only, do not train model",
    )
    parser.add_argument(
        "--act",
        type=str,
        choices=ACT_DICT.keys(),
        default="prelu",
        help="Activation function",
    )
    parser.add_argument(
        "--loss", type=str, choices=LOSS_DICT.keys(), default="L1", help="Loss"
    )
    parser.add_argument(
        "--optim", type=str, choices=OPTIM_DICT.keys(), default="adam", help="Optimizer"
    )
    parser.add_argument("--interdim", type=int, default=64, help="Intermediate dim")
    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
    parser.add_argument("--bs", type=int, default=512, help="Batch size")
    parser.add_argument(
        "--epochs", type=int, default=600, help="Maximum number of epochs to train"
    )
    parser.add_argument(
        "--notrans",
        action="store_true",
        help="Do not transpose (already in row obs form)",
    )
    parser.add_argument("--device", default=0, type=int, help="Device for training")
    return parser


def main():
    """Train a protein predictor"""
    parser = build_parser()
    args = parser.parse_args()

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

    # Specify output log file
    logger = logging.getLogger()
    fh = logging.FileHandler(os.path.join(args.outdir, "training.log"))
    fh.setLevel(logging.INFO)
    logger.addHandler(fh)

    # Log parameters
    for arg in vars(args):
        logging.info(f"Parameter {arg}: {getattr(args, arg)}")
    with open(os.path.join(args.outdir, "params.json"), "w") as sink:
        json.dump(vars(args), sink, indent=4)

    # Load the model
    pretrained_net = model_utils.load_model(args.encoder, device=args.device)

    # Load in some files
    rna_genes = utils.read_delimited_file(os.path.join(args.encoder, "rna_genes.txt"))
    atac_bins = utils.read_delimited_file(os.path.join(args.encoder, "atac_bins.txt"))

    # Read in the RNA
    rna_data_kwargs = copy.copy(sc_data_loaders.TENX_PBMC_RNA_DATA_KWARGS)
    rna_data_kwargs["cluster_res"] = args.clusterres
    rna_data_kwargs["fname"] = args.rnaCounts
    rna_data_kwargs["reader"] = lambda x: load_rna_files(
        x, args.encoder, transpose=not args.notrans
    )

    # Construct data folds
    full_sc_rna_dataset = sc_data_loaders.SingleCellDataset(
        valid_cluster_id=args.validcluster,
        test_cluster_id=args.testcluster,
        **rna_data_kwargs,
    )
    full_sc_rna_dataset.data_raw.write_h5ad(os.path.join(args.outdir, "full_rna.h5ad"))

    train_valid_test_dsets = []
    for mode in ["all", "train", "valid", "test"]:
        logging.info(f"Constructing {mode} dataset")
        sc_rna_dataset = sc_data_loaders.SingleCellDatasetSplit(
            full_sc_rna_dataset, split=mode
        )
        sc_rna_dataset.data_raw.write_h5ad(
            os.path.join(args.outdir, f"{mode}_rna.h5ad")
        )  # Write RNA input
        sc_atac_dummy_dataset = sc_data_loaders.DummyDataset(
            shape=len(atac_bins), length=len(sc_rna_dataset)
        )
        # RNA and fake ATAC
        sc_dual_dataset = sc_data_loaders.PairedDataset(
            sc_rna_dataset,
            sc_atac_dummy_dataset,
            flat_mode=True,
        )
        # encoded(RNA) as "x" and RNA + fake ATAC as "y"
        sc_rna_encoded_dataset = sc_data_loaders.EncodedDataset(
            sc_dual_dataset, model=pretrained_net, input_mode="RNA"
        )
        sc_rna_encoded_dataset.encoded.write_h5ad(
            os.path.join(args.outdir, f"{mode}_encoded.h5ad")
        )
        sc_protein_dataset = sc_data_loaders.SingleCellProteinDataset(
            args.proteinCounts,
            obs_names=sc_rna_dataset.obs_names,
            transpose=not args.notrans,
        )
        sc_protein_dataset.data_raw.write_h5ad(
            os.path.join(args.outdir, f"{mode}_protein.h5ad")
        )  # Write protein
        # x = 16 dimensional encoded layer, y = 25 dimensional protein array
        sc_rna_protein_dataset = sc_data_loaders.SplicedDataset(
            sc_rna_encoded_dataset, sc_protein_dataset
        )
        _temp = sc_rna_protein_dataset[0]  # ensure calling works
        train_valid_test_dsets.append(sc_rna_protein_dataset)

    # Unpack and do sanity checks
    _, sc_rna_prot_train, sc_rna_prot_valid, sc_rna_prot_test = train_valid_test_dsets
    x, y, z = sc_rna_prot_train[0], sc_rna_prot_valid[0], sc_rna_prot_test[0]
    assert (
        x[0].shape == y[0].shape == z[0].shape
    ), f"Got mismatched shapes: {x[0].shape} {y[0].shape} {z[0].shape}"
    assert (
        x[1].shape == y[1].shape == z[1].shape
    ), f"Got mismatched shapes: {x[1].shape} {y[1].shape} {z[1].shape}"

    protein_markers = list(sc_protein_dataset.data_raw.var_names)
    with open(os.path.join(args.outdir, "protein_proteins.txt"), "w") as sink:
        sink.write("\n".join(protein_markers) + "\n")
    assert len(
        utils.read_delimited_file(os.path.join(args.outdir, "protein_proteins.txt"))
    ) == len(protein_markers)
    logging.info(f"Predicting on {len(protein_markers)} proteins")

    if args.preprocessonly:
        return

    protein_decoder_skorch = skorch.NeuralNet(
        module=autoencoders.Decoder,
        module__num_units=16,
        module__intermediate_dim=args.interdim,
        module__num_outputs=len(protein_markers),
        module__activation=ACT_DICT[args.act],
        module__final_activation=nn.Identity(),
        # module__final_activation=nn.Linear(
        #     len(protein_markers), len(protein_markers), bias=True
        # ),  # Paper uses identity activation instead
        lr=args.lr,
        criterion=LOSS_DICT[args.loss],  # Other works use L1 loss
        optimizer=OPTIM_DICT[args.optim],
        batch_size=args.bs,
        max_epochs=args.epochs,
        callbacks=[
            skorch.callbacks.EarlyStopping(patience=15),
            skorch.callbacks.LRScheduler(
                policy=torch.optim.lr_scheduler.ReduceLROnPlateau,
                patience=5,
                factor=0.1,
                min_lr=1e-6,
                # **model_utils.REDUCE_LR_ON_PLATEAU_PARAMS,
            ),
            skorch.callbacks.GradientNormClipping(gradient_clip_value=5),
            skorch.callbacks.Checkpoint(
                dirname=args.outdir,
                fn_prefix="net_",
                monitor="valid_loss_best",
            ),
        ],
        train_split=skorch.helper.predefined_split(sc_rna_prot_valid),
        iterator_train__num_workers=8,
        iterator_valid__num_workers=8,
        device=utils.get_device(args.device),
    )
    protein_decoder_skorch.fit(sc_rna_prot_train, y=None)

    # Plot the loss history
    fig = plot_loss_history(
        protein_decoder_skorch.history, os.path.join(args.outdir, "loss.pdf")
    )


if __name__ == "__main__":
    main()