#!/usr/bin/env python

# native Python imports
import os.path
import shutil
import sys
import sqlite3

import annotations
import subprocess
from cluster_helper.cluster import cluster_view
import database as gemini_db
from gemini_load_chunk import GeminiLoader
import uuid
import time
import datetime


def load(parser, args):
    if (args.db is None or args.vcf is None):
        parser.print_help()
        exit("ERROR: load needs both a VCF file and a database file\n")

    annos = annotations.get_anno_files(args)
    # force skipping CADD and GERP if the data files have not been installed
    if args.skip_cadd is False:
        if 'cadd_score' not in annos:
            sys.stderr.write("\nCADD scores are not being loaded because the"
            " annotation file could not be found.\n"
            "`Run gemini update --dataonly --extra cadd_score`"
            " to install the annotation file.\n\n")
            args.skip_cadd = True
        else:
            sys.stderr.write("CADD scores are being loaded (to skip use:--skip-cadd).\n")
    if args.skip_gerp_bp is False:
        if 'gerp_bp' not in annos:
            sys.stderr.write("\nGERP per bp is not being loaded because the annotation file"
                        " could not be found.\n    Run `gemini update --dataonly --extra gerp_bp`"
                        " to install the annotation file.\n\n")
            args.skip_gerp_bp = True
        else:
            sys.stderr.write("GERP per bp is being loaded (to skip use:--skip-gerp-bp).\n")
    # collect of the the add'l annotation files
    annotations.load_annos(args)

    if args.scheduler:
        load_ipython(args)
    elif args.cores > 1:
        load_multicore(args)
    else:
        load_singlecore(args)

def load_singlecore(args):
    # create a new gemini loader and populate
    # the gemini db and files from the VCF
    gemini_loader = GeminiLoader(args)
    gemini_loader.store_resources()
    gemini_loader.store_version()
    gemini_loader.store_vcf_header()
    gemini_loader.populate_from_vcf()
    gemini_db.add_max_aaf(gemini_loader.c)

    if not args.skip_gene_tables and not args.test_mode:
        gemini_loader.update_gene_table()
    if not args.test_mode:
        gemini_loader.build_indices_and_disconnect()

    if not args.no_genotypes and not args.no_load_genotypes:
        gemini_loader.store_sample_gt_counts()


def load_multicore(args):
    grabix_file = bgzip(args.vcf)
    chunks = load_chunks_multicore(grabix_file, args)
    merge_chunks_multicore(chunks, args)

def load_ipython(args):
    grabix_file = bgzip(args.vcf)
    with cluster_view(*get_ipython_args(args)) as view:
        chunks = load_chunks_ipython(grabix_file, args, view)
        merge_chunks_ipython(chunks, args, view)

def merge_chunks(chunks, db, kwargs):
    cmd = get_merge_chunks_cmd(chunks, db, tempdir=kwargs.get("tempdir"),
                               vcf=kwargs.get("vcf"), anno_type=kwargs.get("anno_type"))
    print "Merging chunks."
    subprocess.check_call(cmd, shell=True)
    cleanup_temp_db_files(chunks)
    return db

def get_merge_chunks_cmd(chunks, db, tempdir=None, vcf=None, anno_type=None):
    chunk_names = ""
    for chunk in chunks:
        chunk_names += " --chunkdb  " + chunk

    tempdir_string, vcf_string, annotype_string = "", "", ""
    if tempdir is not None:
        tempdir_string = " --tempdir " + tempdir
    if vcf is not None:
        vcf_string = " --vcf " + vcf
    if anno_type is not None:
        annotype_string = " -t " + anno_type

    return ("gemini merge_chunks {chunk_names} {tempdir_string} "
            "{vcf_string} {annotype_string} --db {db}").format(**locals())


def finalize_merged_db(tmp_db, db):
    ts = time.time()
    st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
    print st, "indexing final database."

    main_conn = sqlite3.connect(tmp_db)
    main_conn.isolation_level = None
    main_curr = main_conn.cursor()
    main_curr.execute('PRAGMA synchronous = OFF')
    main_curr.execute('PRAGMA journal_mode=MEMORY')

    gemini_db.add_max_aaf(main_curr)
    gemini_db.create_indices(main_curr)

    main_conn.commit()
    main_curr.close()

    shutil.move(tmp_db, db)

def merge_chunks_ipython(chunks, args, view):
    if len(chunks) == 1:
        finalize_merged_db(chunks[0], args.db)
        return args.db
    else:
        ts = time.time()
        st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
        print st, "merging", len(chunks), "chunks."
        sub_merges = get_chunks_to_merge(chunks)
        tmp_dbs = get_temp_dbs(len(sub_merges), os.getcwd())
        merge_args = {"tempdir": args.tempdir, "vcf": args.vcf, "anno_type": args.anno_type}
        view.map(merge_chunks, sub_merges, tmp_dbs, [merge_args] * len(sub_merges))
        merge_chunks_ipython(tmp_dbs, args, view)

def merge_chunks_multicore(chunks, args):
    if len(chunks) <= 1:
        finalize_merged_db(chunks[0], args.db)
        return args.db
    else:
        ts = time.time()
        st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
        print st, "merging", len(chunks), "chunks."
        procs = []
        sub_merges = get_chunks_to_merge(chunks)
        tmp_dbs = get_temp_dbs(len(sub_merges), os.path.dirname(sub_merges[0][0]))
        for sub_merge, tmp_db in zip(sub_merges, tmp_dbs):
            cmd = get_merge_chunks_cmd(sub_merge, tmp_db, tempdir=args.tempdir, vcf=args.vcf,
                                       anno_type=args.anno_type)
            procs.append(subprocess.Popen(cmd, shell=True))
        wait_until_finished(procs)
        cleanup_temp_db_files(chunks)
        merge_chunks_multicore(tmp_dbs, args)

def get_chunks_to_merge(chunks):
    sublist = list_to_sublists(chunks, 2)
    if len(sublist[-1]) > 1:
        return sublist
    else:
        sublist[-2].extend(sublist[-1])
        return sublist[:-1]

def list_to_sublists(l, n):
    """ convert list l to sublists of length n """
    return [l[i:i+n] for i in xrange(0, len(l), n)]

def get_temp_dbs(n, tmp_dir):
    return [os.path.join(tmp_dir, str(uuid.uuid4())) + ".db" for _ in xrange(n)]

def get_chunk_name(chunk):
    return "--chunkdb " + chunk

def load_chunks_multicore(grabix_file, args):
    cores = args.cores

    # specify the PED file if given one
    ped_file = ""
    if args.ped_file is not None:
        ped_file = "-p " + args.ped_file

    # specify the annotation type if given one
    anno_type = ""
    if args.anno_type is not None:
        anno_type = "-t " + args.anno_type

    tempdir = ""
    if args.tempdir is not None:
        if not os.path.exists(args.tempdir):
            os.makedirs(args.tempdir)
        tempdir = "--tempdir " + args.tempdir
    chunk_dir = args.tempdir + "/" if args.tempdir else ""

    no_genotypes = ""
    if args.no_genotypes is True:
        no_genotypes = "--no-genotypes"

    no_load_genotypes = ""
    if args.no_load_genotypes is True:
        no_load_genotypes = "--no-load-genotypes"

    skip_gerp_bp = ""
    if args.skip_gerp_bp is True:
        skip_gerp_bp = "--skip-gerp-bp"

    skip_gene_tables = ""
    if args.skip_gene_tables is True:
        skip_gene_tables = "--skip-gene-tables"

    skip_cadd = ""
    if args.skip_cadd is True:
        skip_cadd = "--skip-cadd"

    test_mode = ""
    if args.test_mode is True:
        test_mode = "--test-mode"

    passonly = ""
    if args.passonly is True:
        passonly = "--passonly"

    skip_info_string = ""
    if args.skip_info_string is True:
        skip_info_string = "--skip-info-string"

    submit_command = "{cmd}"
    vcf, _ = os.path.splitext(os.path.basename(grabix_file))
    chunk_steps = get_chunk_steps(grabix_file, args)
    chunk_vcfs = []
    chunk_dbs = []
    procs = []

    for chunk_num, chunk in chunk_steps:
        start, stop = chunk
        print "Loading chunk " + str(chunk_num) + "."
        gemini_load = gemini_pipe_load_cmd().format(**locals())
        if os.environ.get('GEMINI_DEBUG') == "TRUE":
            print >>sys.stderr, gemini_load
        procs.append(subprocess.Popen(submit_command.format(cmd=gemini_load),
                                      shell=True))

        chunk_vcf = chunk_dir + vcf + ".chunk" + str(chunk_num)
        chunk_vcfs.append(chunk_vcf)
        chunk_dbs.append(chunk_vcf + ".db")

    wait_until_finished(procs)
    print "Done loading {0} variants in {1} chunks.".format(stop, chunk_num+1)
    return chunk_dbs

def load_chunks_ipython(grabix_file, args, view):
    # specify the PED file if given one
    ped_file = ""
    if args.ped_file is not None:
        ped_file = "-p " + args.ped_file

    # specify the annotation type if given one
    anno_type = ""
    if args.anno_type is not None:
        anno_type = "-t " + args.anno_type

    tempdir = ""
    if args.tempdir is not None:
        if not os.path.exists(args.tempdir):
            os.makedirs(args.tempdir)
        tempdir = "--tempdir " + args.tempdir
    chunk_dir = args.tempdir + "/" if args.tempdir else ""

    no_genotypes = ""
    if args.no_genotypes is True:
        no_genotypes = "--no-genotypes"

    no_load_genotypes = ""
    if args.no_load_genotypes is True:
        no_load_genotypes = "--no-load-genotypes"

    skip_gerp_bp = ""
    if args.skip_gerp_bp is True:
        skip_gerp_bp = "--skip-gerp-bp"

    skip_gene_tables = ""
    if args.skip_gene_tables is True:
        skip_gene_tables = "--skip-gene-tables"

    skip_cadd = ""
    if args.skip_cadd is True:
        skip_cadd = "--skip-cadd"

    test_mode = ""
    if args.test_mode is True:
        test_mode = "--test-mode"

    passonly = ""
    if args.passonly is True:
        passonly = "--passonly"

    skip_info_string = ""
    if args.skip_info_string is True:
        skip_info_string = "--skip-info-string"

    vcf, _ = os.path.splitext(os.path.basename(grabix_file))
    chunk_steps = get_chunk_steps(grabix_file, args)
    total_chunks = len(chunk_steps)
    scheduler, queue, cores = get_ipython_args(args)
    load_args = {"ped_file": ped_file,
                 "anno_type": anno_type,
                 "tempdir": tempdir,
                 "chunk_dir": chunk_dir,
                 "vcf": vcf,
                 "grabix_file": grabix_file,
                 "no_genotypes": no_genotypes,
                 "no_load_genotypes": no_load_genotypes,
                 "skip_gerp_bp": skip_gerp_bp,
                 "skip_gene_tables": skip_gene_tables,
                 "skip_cadd": skip_cadd,
                 "test_mode": test_mode,
                 "passonly": passonly,
                 "skip_info_string": skip_info_string}
    chunk_dbs = view.map(load_chunk, chunk_steps, [load_args] * total_chunks)

    print "Done loading variants in {0} chunks.".format(total_chunks)
    return chunk_dbs

def load_chunk(chunk_step, kwargs):
    chunk_num, chunk = chunk_step
    start, stop = chunk
    args = combine_dicts(locals(), kwargs)
    gemini_load = gemini_pipe_load_cmd().format(**args)
    subprocess.check_call(gemini_load, shell=True)
    chunk_db = kwargs["chunk_dir"] + args["vcf"] + ".chunk" + str(chunk_num) + ".db"
    return chunk_db

def wait_until_finished(procs):
    """Wait for parts to finish and ensure a clean finish for each.
    """
    pids = [p.wait() for p in procs]
    if len([p for p in pids if p != 0]) > 0:
        raise ValueError("Processing failed on GEMINI chunk load")

def cleanup_temp_db_files(chunk_dbs):
    for chunk_db in chunk_dbs:
        if os.path.exists(chunk_db):
            os.remove(chunk_db)

def gemini_pipe_load_cmd():
    grabix_cmd = "grabix grab {grabix_file} {start} {stop}"
    gemini_load_cmd = ("gemini load_chunk -v - {anno_type} {ped_file}"
                       " {no_load_genotypes} {no_genotypes}"
                       " {skip_gerp_bp} {skip_gene_tables} {skip_cadd}"
                       " {passonly} {skip_info_string} {test_mode} {tempdir}"
                       " -o {start} {chunk_dir}{vcf}.chunk{chunk_num}.db")
    return " | ".join([grabix_cmd, gemini_load_cmd])

def get_chunk_steps(grabix_file, args):
    index_file = grabix_index(grabix_file)
    num_lines = get_num_lines(index_file)
    args.cores = min(int(args.cores), num_lines)
    chunk_size = int(num_lines) / args.cores
    print "Breaking {0} into {1} chunks.".format(grabix_file, args.cores)

    starts = []
    stops = []
    for chunk in range(args.cores):
        start = (chunk * chunk_size) + 1
        stop  = start + chunk_size - 1
        # make sure the last chunk covers the remaining lines
        if chunk >= (args.cores - 1) and stop < num_lines:
            stop = num_lines
        starts.append(start)
        stops.append(stop)
    return list(enumerate(zip(starts, stops)))

def get_num_lines(index_file):
    with open(index_file) as index_handle:
        index_handle.next()
        num_lines = int(index_handle.next().strip())
    print "Loading %d variants." % (num_lines)
    return num_lines

def grabix_index(fname):
    if not which("grabix"):
        print_cmd_not_found_and_exit("grabix")
    index_file = fname + ".gbi"
    if file_exists(index_file) and os.path.getmtime(index_file) > os.path.getmtime(fname):
        return index_file
    print "Indexing {0} with grabix.".format(fname)
    subprocess.check_call("grabix index {fname}".format(fname=fname), shell=True)
    return index_file

def bgzip(fname):

    if not which("bgzip"):
        print_cmd_not_found_and_exit("bgzip")

    if is_gz_file(fname):
        return fname

    vcf_time = os.path.getmtime(fname)
    bgzip_file = fname + ".gz"

    if not file_exists(bgzip_file) or \
       (file_exists(bgzip_file) and os.path.getmtime(bgzip_file) < vcf_time):
        print "Bgzipping {0} into {1}.".format(fname, fname + ".gz")
        subprocess.check_call("bgzip -c {fname} > \
                              {fname}.gz".format(fname=fname),
                              shell=True)
    elif file_exists(bgzip_file) and os.path.getmtime(bgzip_file) > vcf_time:
        print "Loading with existing bgzip ({0}) version of {1}.".format(fname + ".gz", fname)

    return bgzip_file


def is_gz_file(fname):
    _, ext = os.path.splitext(fname)
    if ext == ".gz":
        return True
    else:
        return False

def get_submit_command(args):
    return "{cmd}"


def file_exists(fname):
    """Check if a file exists and is non-empty.
    """
    return os.path.exists(fname) and os.path.getsize(fname) > 0

def which(program):
    """ returns the path to an executable or None if it can't be found
     http://stackoverflow.com/questions/377017/test-if-executable-exists-in-python
     """

    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

def combine_dicts(d1, d2):
    return dict(d1.items() + d2.items())

def get_ipython_args(args):
    return (args.scheduler, args.queue, args.cores)

def print_cmd_not_found_and_exit(cmd):
    sys.stderr.write("Cannot find {cmd}, install it or put it in your "
                     "path.".format(cmd=cmd))
    exit(1)

def use_scheduler(args):
    return bool(args.scheduler)
