# vireoSNP - donor deconvolution for multiplexed scRNA-seq data
# Author: Yuanhua Huang
# Date: 24-06-2019

import os
import sys
import time
import subprocess
import numpy as np
import multiprocessing
from scipy.io import mmread
from optparse import OptionParser, OptionGroup

from .version import __version__
from .utils.vireo_base import match, optimal_match
from .utils.vireo_wrap import vireo_wrap

from .plot.base_plot import plot_GT
from .utils.io_utils import match_donor_VCF
from .utils.io_utils import write_donor_id, read_cellSNP, read_vartrix
from .utils.vcf_utils import load_VCF, write_VCF, parse_donor_GPb
from .utils.vcf_utils import read_sparse_GeneINFO, GenoINFO_maker


START_TIME = time.time()

def show_progress(RV=None):
    return RV

def main():
    # import warnings
    # warnings.filterwarnings('error')

    # parse command line options
    parser = OptionParser()
    parser.add_option("--cellData", "-c", dest="cell_data", default=None,
        help=("The cell genotype file in VCF format or cellSNP folder with "
              "sparse matrices."))
    parser.add_option("--nDonor", "-N", type="int", dest="n_donor",
        default=None, help=("Number of donors to demultiplex; can be larger "
        "than provided in donor_file"))
    parser.add_option("--outDir", "-o", dest="out_dir", default=None,
        help=("Dirtectory for output files [default: $cellFilePath/vireo]"))

    group0 = OptionGroup(parser, "Optional input files")
    group0.add_option("--vartrixData", dest="vartrix_data", default=None,
        help=("The cell genotype files in vartrix outputs (three/four files, "
              "comma separated): alt.mtx,ref.mtx,barcodes.tsv,SNPs.vcf.gz. "
              "This will suppress cellData argument."))
    group0.add_option("--donorFile", "-d", dest="donor_file", default=None,
        help=("The donor genotype file in VCF format. Please filter the sample "
        "and region with bcftools -s and -R first!"))
    group0.add_option("--genoTag", "-t", dest="geno_tag", default='PL',
        help=("The tag for donor genotype: GT, GP, PL [default: %default]"))

    group1 = OptionGroup(parser, "Optional arguments")
    group1.add_option("--noDoublet", dest="no_doublet", action="store_true",
        default=False, help="If use, not checking doublets.")
    group1.add_option("--nInit", "-M", type="int", dest="n_init", default=50,
        help=("Number of random initializations, when GT needs to learn "
        "[default: %default]"))
    group1.add_option("--extraDonor", type=int, dest="n_extra_donor",
        default=0, help=("Number of extra donor in pre-cluster, when GT "
        "needs to learn [default: %default]"))
    group1.add_option("--extraDonorMode", dest="extra_donor_mode",
        default="distance", help=("Method for searching from extra donors. "
        "size: n_cell per donor; distance: GT distance between donors "
        "[default: %default]"))
    group1.add_option("--forceLearnGT", dest="force_learnGT", default=False,
        action="store_true", help="If use, treat donor GT as prior only.")
    group1.add_option("--ASEmode", dest="ASE_mode", default=False,
        action="store_true", help="If use, turn on SNP specific allelic ratio.")
    group1.add_option("--noPlot", dest="no_plot", default=False,
        action="store_true", help="If use, turn off plotting GT distance.")
    group1.add_option("--randSeed", type="int", dest="rand_seed", default=None,
        help="Seed for random initialization [default: %default]")
    group1.add_option("--cellRange", type="str", dest="cell_range", default=None,
        help="Range of cells to process, eg. 0-10000 [default: all]")
    group1.add_option("--callAmbientRNAs", dest="check_ambient", default=False,
        action="store_true", help=("If use, detect ambient RNAs in each cell "
        "(under development)"))
    group1.add_option("--nproc", "-p", type="int", dest="nproc", default=1,
        help=("Number of subprocesses for computing - this sacrifices memory "
        "for speedups [default: %default]"))

    parser.add_option_group(group0)
    parser.add_option_group(group1)
    (options, args) = parser.parse_args()

    if len(sys.argv[1:]) == 0:
        print("Welcome to vireoSNP v%s!\n" %(__version__))
        print("use -h or --help for help on argument.")
        sys.exit(1)

    ## out directory
    if options.out_dir is None:
        print("Warning: no outDir provided, we use $cellFilePath/vireo.")
        out_dir = os.path.dirname(os.path.abspath(options.cell_file)) + "/vireo"
    elif os.path.dirname(options.out_dir) == "":
        out_dir= "./" + options.out_dir
    else:
        out_dir = options.out_dir
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    ## input data (VCF.gz or a folder with sparse matrices)
    if options.cell_data is None and options.vartrix_data is None:
        print("Error: need cell data in vcf file, or cellSNP output folder, or "
              "vartrix's alt.mtx,ref.mtx,barcodes.tsv.")
        sys.exit(1)
    elif options.vartrix_data is not None:
        print("[vireo] Loading vartrix files ...")
        vartrix_files = options.vartrix_data.split(",")
        if len(vartrix_files) < 3 or len(vartrix_files) > 4:
            print("Error: vartrixData requires 3 or 4 comma separated files")
            sys.exit(1)
        elif len(vartrix_files) == 3:
            vartrix_files.append(None)

        cell_dat = read_vartrix(vartrix_files[0], vartrix_files[1],
                                vartrix_files[2], vartrix_files[3])
    elif os.path.isdir(os.path.abspath(options.cell_data)):
        print("[vireo] Loading cell folder ...")
        cell_dat = read_cellSNP(options.cell_data)
    else:
        print("[vireo] Loading cell VCF file ...")
        cell_vcf = load_VCF(options.cell_data, biallelic_only=True)
        cell_dat = read_sparse_GeneINFO(cell_vcf['GenoINFO'], keys=['AD', 'DP'])
        for _key in ['samples', 'variants', 'FixedINFO', 'contigs', 'comments']:
            cell_dat[_key] = cell_vcf[_key]

    ## subset input cell data if necessary
    if options.cell_range is not None:
        cellRange = options.cell_range.split("-")
        cellRange[0] = int(cellRange[0])
        cellRange[1] = int(cellRange[1])
        cell_dat['AD'] = cell_dat['AD'][:,cellRange[0]:cellRange[1]]
        cell_dat['DP'] = cell_dat['DP'][:,cellRange[0]:cellRange[1]]
        cell_dat['samples'] = cell_dat['samples'][cellRange[0]:cellRange[1]]


    ## input donor genotype
    n_donor = options.n_donor
    if options.donor_file is not None:
        if "variants" not in cell_dat.keys():
            print("No variants information is loaded, please provide base.vcf.gz")
            sys.exit(1)

        print("[vireo] Loading donor VCF file ...")
        donor_vcf = load_VCF(options.donor_file, biallelic_only=True,
                             sparse=False, format_list=[options.geno_tag])
        
        if (donor_vcf['n_SNP_tagged'][0] < 
            (0.1 * len(donor_vcf['GenoINFO'][options.geno_tag]))):
            print("[vireo] No " + options.geno_tag + " tag in donor genotype; "
                "please try another tag for genotype, e.g., GT")
            print("        %s" %options.donor_file)
            sys.exit(1)

        cell_dat, donor_vcf = match_donor_VCF(cell_dat, donor_vcf)
        donor_GPb = parse_donor_GPb(donor_vcf['GenoINFO'][options.geno_tag],
            options.geno_tag)

        if n_donor is None or n_donor == donor_GPb.shape[1]:
            n_donor = donor_GPb.shape[1]
            donor_names = donor_vcf['samples']
            learn_GT = False
        elif n_donor < donor_GPb.shape[1]:
            learn_GT = False
            donor_names = ['donor%d' %x for x in range(n_donor)]
        else:
            learn_GT = True
            donor_names = (donor_vcf['samples'] +
                ['donor%d' %x for x in range(donor_GPb.shape[1], n_donor)])
    else:
        learn_GT = True
        donor_GPb = None
        donor_names = ['donor%d' %x for x in range(n_donor)]

    n_vars = np.array(np.sum(cell_dat['DP'] > 0, axis=0)).reshape(-1)

    if options.force_learnGT:
        learn_GT = True

    # extra donor for initial search, only for learn_GT
    n_extra_donor = 0
    if learn_GT:
        if options.n_extra_donor is None or options.n_extra_donor == "None":
            n_extra_donor = int(round(np.sqrt(n_donor)))
        else:
            n_extra_donor = options.n_extra_donor

    # number of initials, only for learn_GT
    n_init = options.n_init if learn_GT else 1

    check_doublet = options.no_doublet == False

    ## run vireo model (try multiple initializations)
    print("[vireo] Demultiplex %d cells to %d donors with %d variants." %(
        cell_dat['AD'].shape[1], n_donor, cell_dat['AD'].shape[0]))
    res_vireo = vireo_wrap(cell_dat['AD'], cell_dat['DP'], n_donor=n_donor,
        GT_prior=donor_GPb, learn_GT=learn_GT, n_init=n_init,
        n_extra_donor=n_extra_donor, extra_donor_mode=options.extra_donor_mode,
        check_doublet=check_doublet, random_seed=options.rand_seed,
        ASE_mode=options.ASE_mode, check_ambient=options.check_ambient, 
        nproc=options.nproc)

    if (n_donor is not None and
        donor_GPb is not None and n_donor < donor_GPb.shape[1]):
        idx = optimal_match(res_vireo['GT_prob'], donor_GPb)[1]
        donor_names = [donor_vcf['samples'][x] for x in idx]

    ## save donor id for each cell
    write_donor_id(out_dir, donor_names, cell_dat['samples'], n_vars, res_vireo)

    if options.no_plot == False and options.vartrix_data is None:
        idx = np.array(np.sum(cell_dat['DP'], axis=1) > (3*n_donor)).reshape(-1)
        if learn_GT and donor_GPb is not None:
            plot_GT(out_dir, res_vireo['GT_prob'][idx, :, :], donor_names,
                    donor_GPb[idx, :, :], donor_vcf['samples'])
        else:
            plot_GT(out_dir, res_vireo['GT_prob'][idx, :, :], donor_names)

    # ## save inferred donor genotype
    if learn_GT and 'variants' in cell_dat.keys():
        donor_vcf_out = cell_dat
        donor_vcf_out['samples'] = donor_names
        donor_vcf_out['GenoINFO'] = GenoINFO_maker(res_vireo['GT_prob'],
            cell_dat['AD'] * res_vireo['ID_prob'],
            cell_dat['DP'] * res_vireo['ID_prob'])
        write_VCF(out_dir + "/GT_donors.vireo.vcf.gz", donor_vcf_out)

    run_time = time.time() - START_TIME
    print("[vireo] All done: %d min %.1f sec" %(int(run_time / 60),
                                                    run_time % 60))
    print()


if __name__ == "__main__":
    main()
