"""
Script for predicting protein expression
"""

import os
import sys
import logging
import argparse

import numpy as np
import pandas as pd

import torch
import torch.nn as nn
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)
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
import protein_utils
import utils

from predict_model import (
    load_atac_files_for_eval,
    load_rna_files_for_eval,
)


def build_parser():
    """Build commandline parser"""
    parser = argparse.ArgumentParser(
        usage=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--babel", type=str, required=True, help="Path to babel model")
    parser.add_argument(
        "--protmodel", type=str, required=True, help="Path to latent-to-protein model"
    )
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument("--atac", type=str, nargs="*", help="Input ATAC")
    input_group.add_argument("--rna", type=str, nargs="*", help="Input RNA")
    parser.add_argument(
        "--liftHg19toHg38",
        action="store_true",
        help="Liftover input ATAC bins from hg19 to hg38 (only used for ATAC input)",
    )
    parser.add_argument(
        "-o", "--output", required=True, type=str, help="csv file to output"
    )
    parser.add_argument("--device", default=1, type=int, help="Device for training")
    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()
    assert args.output.endswith(".csv")

    # Specify output log file
    logger = logging.getLogger()
    fh = logging.FileHandler(args.output + ".log")
    fh.setLevel(logging.INFO)
    logger.addHandler(fh)

    # Log parameters
    for arg in vars(args):
        logging.info(f"Parameter {arg}: {getattr(args, arg)}")

    # Load the model
    babel = model_utils.load_model(args.babel, device=args.device)
    # Load in some related files
    rna_genes = utils.read_delimited_file(os.path.join(args.babel, "rna_genes.txt"))
    atac_bins = utils.read_delimited_file(os.path.join(args.babel, "atac_bins.txt"))

    # Load in the protein accesory model
    babel_prot_acc_model = protein_utils.load_protein_accessory_model(args.protmodel)
    proteins = utils.read_delimited_file(
        os.path.join(args.protmodel, "protein_proteins.txt")
    )

    # Get the encoded layer based on input
    if args.rna:
        (
            sc_rna_dset,
            _rna_genes,
            _marker_genes,
            _housekeeper_genes,
        ) = load_rna_files_for_eval(args.rna, checkpoint=args.babel, no_filter=True)
        sc_atac_dummy_dset = sc_data_loaders.DummyDataset(
            shape=len(atac_bins), length=len(sc_rna_dset)
        )
        sc_dual_dataset = sc_data_loaders.PairedDataset(
            sc_rna_dset,
            sc_atac_dummy_dset,
            flat_mode=True,
        )
        sc_dual_encoded_dataset = sc_data_loaders.EncodedDataset(
            sc_dual_dataset, model=babel, input_mode="RNA"
        )
        cell_barcodes = list(sc_rna_dset.data_raw.obs_names)
        encoded = sc_dual_encoded_dataset.encoded
    else:
        sc_atac_dset, _loaded_atac_bins = load_atac_files_for_eval(
            args.atac, checkpoint=args.babel, lift_hg19_to_hg39=args.liftHg19toHg38
        )
        sc_rna_dummy_dset = sc_data_loaders.DummyDataset(
            shape=len(rna_genes), length=len(sc_atac_dset)
        )
        sc_dual_dataset = sc_data_loaders.PairedDataset(
            sc_rna_dummy_dset, sc_atac_dset, flat_mode=True
        )
        sc_dual_encoded_dataset = sc_data_loaders.EncodedDataset(
            sc_dual_dataset, model=babel, input_mode="ATAC"
        )
        cell_barcodes = list(sc_atac_dset.data_raw.obs_names)
        encoded = sc_dual_encoded_dataset.encoded

    # Array of preds
    prot_preds = babel_prot_acc_model.predict(encoded.X)
    prot_preds_df = pd.DataFrame(
        prot_preds,
        index=cell_barcodes,
        columns=proteins,
    )
    prot_preds_df.to_csv(args.output)


if __name__ == "__main__":
    main()