#!/usr/bin/env python3
import pysam
import argparse
import os
import pickle
import shutil
import subprocess
import sys
import time
import math
import matplotlib
import numpy
import multiprocessing

matplotlib.use('PS')  # PS seems to be faster than Agg, even though .pngs are generated. but why?
import matplotlib.pyplot as plt


def detect_tiles(path):
    """
    Function that automatically detects the tiles to use for later HiLive use.
    :param path: str path to the bcdir
    :return: list of str for all existing tiles in bcdir in Lane001 C1.1 after completion of first cycle
    """

    while not os.path.isdir(os.path.join(path, 'L001', 'C2.1')):  # wait until first cycle is completed
        time.sleep(10)

    alltiles = os.listdir(os.path.join(path, 'L001', 'C1.1'))
    return [os.path.splitext(x.split('_')[-1])[0] for x in alltiles if x.endswith('.bcl')]  # return all tile numbers


def detect_lanes(path):
    """
    Function that automatically detects the lanes to use for later HiLive use.
    :param path: str path to bcdir
    :return: list of str for all existing lanes
    """

    while not os.path.isdir(os.path.join(path, 'L001')):
        time.sleep(10)

    return [x[-1] for x in os.listdir(path) if
            x.startswith('L00')]  # return the number n of folders with pattern 'L00n'

def assemble_hilive2_call(arguments, cyclelist):
    # optional arguments
    call = [arguments.hilive, '--num-threads {}'.format(arguments.num_threads)]
    if not arguments.lanes=="":
        call.append('--lanes {}'.format(arguments.lanes))
    if arguments.max_tile:
        call.append('--max-tile {}'.format(arguments.max_tile))
    else:
        call.append('--tiles {}'.format(arguments.tiles))
    if not arguments.config=="":
        call.append('--config {}'.format(arguments.config))
    call.append('--temp {}'.format(os.path.join(arguments.out_dir, 'temp')))
    call.append('--out-mode {}'.format(arguments.out_mode))
    call.append('--max-softclip-ratio {}'.format(arguments.max_softclip_ratio))
    call.append('--out-cycles {}'.format(arguments.out_cycles))
    call.append('--out-format {}'.format(arguments.out_format)) 

    # fixed arguments
    call.append('--align-mode {}'.format(arguments.align_mode))
    call.append('--bcl-dir {}'.format(arguments.bcl_dir))
    call.append('--index {}'.format(arguments.index))
    call.append('--reads {}'.format(arguments.reads))
    call.append('--out-dir {}'.format(arguments.out_dir))
    print(' '.join(call))
    return ' '.join(call)

def run_subprocess(command, capture_streams=False):
    """
    Function that launches a subprocess, waits for it to terminate and returns the returncode and if desired the output.
    :param command: str shell command to execute
    :param capture_streams: bool designating whether
    :return: tuple of (int returncode of the executed program, tuple of (stdout, stderr))
    """

    if not capture_streams:
        proc = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    else:
        proc = subprocess.Popen(command, shell=True)

    return proc


def unpickle_backgrounddict(backgroundfile):
    """
    Loads background dict from pickle-file
    :param backgroundfile: path to backgrounddict
    :return: dictionary of covered viral reference positions from human genome project
    """
    backgroundfile_open = open(backgroundfile, 'rb')
    backgrounddict = pickle.load(backgroundfile_open)
    backgroundfile_open.close()
    return backgrounddict


def parse_bamfile(bamfile, backgrounddict):
    """
    Function that parses the necessary information from generated .bamfile
    :param bamfile: path to file of type .sam - must contain header
    :param backgrounddict: dict of background abundances
    :return: tuple of dicts: referencedict, uniquebasesdict, coveragedict, readcountdict
    :return: referencedict: key = str refname, val = int reflength
    :return: uniquebasesdict: key = str refname, val = int ref positions only covered in foreground
    :return: coveragedict: key = str refname, val = dict; key = reference positions, val = coverage
    :return: readcountdict: key = str refname, val = int number of mapped reads

    """
    readcountdict = {}  # dict with read counts for tree
    coveragedict = {}  # dict with complete coverages for plotting
    uniquebasesdict = {}  # dict containing binary information if a position is covered
    bamfile = pysam.AlignmentFile(bamfile, 'rb')
    referencedict = dict(zip(bamfile.references, bamfile.lengths))
    for read in bamfile:
        alignment_score = int(read.tags[0][1])

        ref = read.reference_name
        hit_range = read.get_reference_positions()
        for hitpos in hit_range:
            try:
                if hitpos not in backgrounddict[ref]:  # if a position is also hit in BG: Ignore it
                    try:
                        uniquebasesdict[ref] |= {
                            hitpos}  # add hitpos to unique set of positions only found in FG
                    except KeyError:
                        uniquebasesdict[ref] = set()
                        uniquebasesdict[ref] |= {hitpos}
            except KeyError:  # if backgrounddict does not contain reference
                try:
                    uniquebasesdict[ref] |= {hitpos}
                except KeyError:
                    uniquebasesdict[ref] = set()
                    uniquebasesdict[ref] |= {hitpos}
            try:
                coveragedict[ref][hitpos] += 1
            except KeyError:
                try:
                    coveragedict[ref][hitpos] = 1
                except KeyError:
                    coveragedict[ref] = {}
                    coveragedict[ref][hitpos] = 1

        try:
            readcountdict[ref] += 1  # count reads mapping on reference
        except KeyError:
            readcountdict[ref] = 1

    for ref in uniquebasesdict.keys():  # flatten dict of covered bases, count covered bases and return int(len(set))
        uniquebasesdict[ref] = len(uniquebasesdict[ref])

    return referencedict, uniquebasesdict, coveragedict, readcountdict


def parse_samfile(samfile, backgrounddict):
    """
    Function that parses the necessary information from generated .samfiles
    :param samfile: path to file of type .sam - must contain header
    :param backgrounddict: dict of background abundances
    :return: tuple of dicts: referencedict, uniquebasesdict, coveragedict, readcountdict
    :return: referencedict: key = str refname, val = int reflength
    :return: uniquebasesdict: key = str refname, val = int ref positions only covered in foreground
    :return: coveragedict: key = str refname, val = dict; key = reference positions, val = coverage
    :return: readcountdict: key = str refname, val = int number of mapped reads

    """
    readcountdict = {}  # dict with read counts for tree
    coveragedict = {}  # dict with complete coverages for plotting
    uniquebasesdict = {}  # dict containing binary information if a position is covered
    samfile = pysam.AlignmentFile(samfile, 'r')
    referencedict = dict(zip(samfile.references, samfile.lengths))
    for read in samfile:
        alignment_score = int(read.tags[0][1])

        ref = read.reference_name
        hit_range = read.get_reference_positions()
        for hitpos in hit_range:
            try:
                if hitpos not in backgrounddict[ref]:  # if a position is also hit in BG: Ignore it
                    try:
                        uniquebasesdict[ref] |= {
                            hitpos}  # add hitpos to unique set of positions only found in FG
                    except KeyError:
                        uniquebasesdict[ref] = set()
                        uniquebasesdict[ref] |= {hitpos}
            except KeyError:  # if backgrounddict does not contain reference
                try:
                    uniquebasesdict[ref] |= {hitpos}
                except KeyError:
                    uniquebasesdict[ref] = set()
                    uniquebasesdict[ref] |= {hitpos}
            try:
                coveragedict[ref][hitpos] += 1
            except KeyError:
                try:
                    coveragedict[ref][hitpos] = 1
                except KeyError:
                    coveragedict[ref] = {}
                    coveragedict[ref][hitpos] = 1

        try:
            readcountdict[ref] += 1  # count reads mapping on reference
        except KeyError:
            readcountdict[ref] = 1

    for ref in uniquebasesdict.keys():  # flatten dict of covered bases, count covered bases and return int(len(set))
        uniquebasesdict[ref] = len(uniquebasesdict[ref])

    return referencedict, uniquebasesdict, coveragedict, readcountdict


def parse_phage_file(phage_file):
    """
    Generates BSL-Value-Dictionary from textfile
    :param bsl_file: path to BSL List in .txt-format
    :return: dictionary of bsl scores to trivial virus names. May contain different virus names than expected
    """
    with open(phage_file) as phage_file:
        phage_list = []
        for line in phage_file:
            li = line.rstrip()
            phage_list.append(str(li).lower()) # switch names to lower case
    return phage_list


def parse_bsl_file(bsl_file):
    """
    Generates BSL-Value-Dictionary from textfile
    :param bsl_file: path to BSL List in .txt-format
    :return: dictionary of bsl scores to trivial virus names. May contain different virus names than expected
    """
    with open(bsl_file) as bsl_file:
        bsl_dict = {}
        for line in bsl_file:
            li = line.rstrip()
            l = li.split('\t')
            bsl_dict[str(l[0]).lower()] = l[1]  # switch names to lower case
    return bsl_dict


def bsl_leveller(bsl_dict, phage_list, familyname, speciesname):
    """
    Parses the BSL-scores of species from the provided BSL-dict.
    :param bsl_dict: dictionary of bsl scores to trivial virus names. May contain different virus names than expected,
    mapping is not perfect
    :param speciesname: name of the virus found in the reference database. Not guaranteed to be in same format as
    BSL-dict
    :return: int [1:4] BSL-Score, always lowest (=1) if not defined
    """
    lower_speciesname = speciesname.lower()
    if 'phage' in lower_speciesname:
        return 1
    if familyname.lower() in phage_list:
        return 1
    try:
        return [int(bsl_dict[s]) for s in bsl_dict.keys() if
                lower_speciesname in s or s.split('viruses')[0] in lower_speciesname][
                   0] + 1
        # either the found virus appears in bsl_dict or the bsl_dict key without viruses appears in found virus.
        # Note: Some groups of viruses are defined on higher level (e.g. adenoviruses on family level).
    except IndexError:
        return 2


def copy_tree_structure(program_path, resultpath):
    """
    Copies files used for html presentation of results to the resultfolder
    :param program_path: str path to program
    :param resultpath: str path to result
    :return: nothing
    """
    shutil.copytree(os.path.join(program_path, 'Collapsible_Tree_with_Amounts'),
                    os.path.join(resultpath, 'Collapsible_Tree_with_Amounts'))


def coverage_plot(inputdata):
    """
    Plots foreground coverage and background abundance
    :param inputdata: tuple of all input data
    [0]length: int length of reference
    [1]coveragedict: dict; key = str refname, val = dict; key = reference positions, val = coverage
    [2]backgroundcov: dict of covered viral reference positions from human genome project
    [3]refname: str name of reference
    [4]resultpath: str path for results
    :return: nothing
    """
    length = inputdata[0]
    coveragedict = inputdata[1]
    backgroundcov = inputdata[2]
    refname = inputdata[3]
    resultpath = inputdata[4]
    backcovlist = []
    forecovlist = []
    for i in range(length):
        try:
            backcovlist.append(-numpy.log10(backgroundcov[
                                                i] + 1))  # +1 to avoid logarithmic zeros.
        except KeyError:
            backcovlist.append(0)

    for i in range(length):
        try:
            forecovlist.append(numpy.log10(coveragedict[i] + 1))
        except KeyError:
            forecovlist.append(0)

    fig, ax = plt.subplots()

    try:
        plt.title(refname.split('||')[-1].split(':')[2].split('|')[0].replace('_', ' '))
    except IndexError:
        plt.title(refname)  # maybe not necessary anymore
    plt.xlabel('reference position')
    plt.ylabel('background abundance          foreground coverage')
    ax.set_xlim([0, length])
    try:
        limit = max(max(backgroundcov.values()), max(
            coveragedict.values())) + 1
        # sets limit to highest value of fore- and backgrounddict, +1 because we don't want bars to touch the border
    except ValueError:
        limit = max(coveragedict.values()) + 1
    maximized_limit = (len(str(int(limit))))  # calculates the next full power of 10
    ax.set_ylim([-maximized_limit, maximized_limit])
    minor_ticks_list = [math.log(j * 10 ** i, 10) for j in range(2, 10) for i in
                        range(0, maximized_limit)]  # generates list of form log10([2,3,4,5,6,7,8,9,20,30,40...])
    minor_ticks = sorted(
        [-x for x in minor_ticks_list] + minor_ticks_list)  # concatenates sorted list of negative and positive ticks
    ax.set_yticks(minor_ticks, minor=True)
    labels = [10 ** i for i in range(maximized_limit, 0, -1)] + [0] + [10 ** i for i in range(1,
                                                                                              maximized_limit + 1)]
    # generates list of form [...100,10,0,10,100...]
    major_ticks = range(-maximized_limit, maximized_limit + 1)
    ax.set_yticks(major_ticks, minor=False)
    ax.set_yticklabels(labels)
    ax.fill_between(range(length), 0, forecovlist, facecolor='green', edgecolor='green')
    ax.fill_between(range(length), 0, backcovlist, facecolor='red', edgecolor='red')
    ax.axhline(y=0, color='black', linewidth=1.5)  # this is just a fix to hide red line slightly above 0
    plt.savefig(
        os.path.join(resultpath, '{}.png'.format(str(refname).split('|')[1])))  # name plot by GI number
    plt.close()


def print_treedata(readcountdict, uniquebasesdict, backgrounddict, bsl_dict, phage_list, resultfile):
    """
    Writes results into csv file for D3
    :param readcountdict: dict of int coverage of int position of str reference in foreground
    :param uniquebasesdict: dict of int coverage of str reference found in foreground only
    :param backgrounddict: dict of int coverage of str reference found in background
    :param bsl_dict: dict of int bsl-level of str reference identifier (not always exactly matching str reference)
    :param resultfile: str path to file to write results to
    :return: nothing
    """
    outfile = open(resultfile, 'w')
    outfile.write('Category,Level1,Level2,Level3,Level4,AllHits,BSLlvl,,UnambigBases,WeightedScore,Total\n')
    linelist = []
    for key, value in readcountdict.items():
        giname = key.split('|')[1]  # set to 0 for old database format!
        speciesnum = value
        if key not in uniquebasesdict:
            weighted_score = 0
        else:
            if key in backgrounddict and key in uniquebasesdict:
                weighted_score = uniquebasesdict[key] / len(
                    backgrounddict[key]) * (math.log(speciesnum)+1) # ratio of number of unique fore- and background bases
            else:
                weighted_score = uniquebasesdict[key] * (math.log(speciesnum)+1) # add pseudocount to circumvent 0-scores for one read
        try:
            basesnum = uniquebasesdict[key]
        except KeyError:  # references without unique foreground bases are missing in uniquebasesdict
            basesnum = 0
        key = key.replace(',', ';')  # remove commas from species tax name to avoid crashing the csv tree
        if '|species:' not in key:
            speciesname = 'unassigned species'
            speciesname_no_number = speciesname
        else:
            speciesname = key.split('|species:')[1].split('|')[0]
            speciesname_no_number = speciesname.split(':')[1]
        if '|genus:' not in key:
            genusname = 'unassigned genus'
        else:
            genusname = key.split('|genus:')[1].split('|')[0]
        if '|family:' not in key:
            familyname = 'unassigned family'
            familyname_no_number = familyname
        else:
            familyname = key.split('|family:')[1].split('|')[0]
            familyname_no_number = familyname.split(':')[1]
        linelist.append('---Root,' + familyname + ',' + genusname + ',' + speciesname + ',' + giname + ',' + str(
            speciesnum) + ',' + str(bsl_leveller(bsl_dict, phage_list, familyname_no_number, speciesname_no_number)) + ',,' + str(basesnum) + ',' + str(
            int(weighted_score)) + ',0.1,\n')
    linelist.sort()
    for each in linelist:
        outfile.write(each)
    outfile.close()


def parse_command_line_options(args):
    """
    Parses arguments given in command line
    :param args: command line arguments
    :return: string
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('-b', '--bcl_dir', help='Illumina BaseCalls directory'
        ' which contains the sequence information of the reads.', 
        required=True)
    parser.add_argument('-i', '--index', help='Reference index directory for '
        'HiLive2.', required=True)
    parser.add_argument('-r', '--reads', help='Length and types of the read '
        'segments. Each segment is either a read (R) or a barcode (B). Please'
        ' give the segments in the correct order as they are produced by the '
        'sequencing machine. [Example: 101R,8B,8B,101R]', required=True)
    parser.add_argument('-O', '--out_cycles', help='Cycles to create output '
        'for.', required=True)
    parser.add_argument('-o', '--out_dir', help='Output directory.', 
        required=True)
    parser.add_argument('-R', '--max_softclip_ratio', help=' Maximal relative'
        ' length of the front softclip (only relevant during output) '
        '[Default: 0.2]', default=float(0.2), required=False)
    parser.add_argument('-H', '--hilive', help='Path to hilive executable', 
        default="hilive", required=False)
    parser.add_argument('-n', '--num_threads', help='Number of threads', 
        default=1, required=False)
    parser.add_argument('-m', '--align_mode', help='Alignment mode to balance'
        ' speed and accuracy [very-fast|fast|balanced|accurate|very-accurate].', 
        default="balanced", required=False)
    parser.add_argument('-M', '--out_mode', help='The output mode. [Default: '
        'ANYBEST]. [ALL|A]: Report all found alignments. [BESTN#|N#]: Report '
        'the # best found alignments. [ALLBEST|H]: Report all found '
        'alignments with the best score. [ANYBEST|B]: Report one best '
        'alignment. [UNIQUE|U]: Report only unique alignments.', default="B", 
        required=False)
    parser.add_argument('-f', '--out_format', help='Set the output format '
        '[Default: BAM]. [SAM]: SAM output. [BAM]: BAM output.', default='BAM',
        required=False)
    parser.add_argument('-v', '--verbose', help='Print output of HiLive. '
        'Default: False', default=False, required=False, action='store_true')
    parser.add_argument('-l', '--lanes', help='Lanes to be used for mapping, '
        '[Default = automatic detection]', required=False)
    parser.add_argument('-t', '--tiles', help='Number of tiles [Default: '
        'automatic detection]', required=False)
    parser.add_argument('-T', '--max_tile', help='Specify the highest tile '
        'number. The tile numbers will be computed by this number, '
        'considering the correct surface count, swath count and tile count '
        'for Illumina sequencing.', required=False)
    parser.add_argument('-c', '--config', help='Path to a config file for HiLive2.'
        ' This enables to use advanced alignment settings that can not be set with '
        ' the available PathoLive parameters.', default="", required=False)
    parser.add_argument('-S', '--skip_mapping', help=argparse.SUPPRESS, 
        default=False, required=False, action='store_true')
    return parser.parse_args(args)

def main():
    start = time.time()
    # command line parser
    print('parsing command line')
    arguments = parse_command_line_options(sys.argv[1:])
    program_path = os.path.dirname(os.path.realpath(__file__))
    if not os.path.exists(os.path.join(arguments.out_dir)):
        os.makedirs(os.path.join(arguments.out_dir))
    cyclelist = sorted(set([int(x) for x in arguments.out_cycles.split(",")]))

    if not arguments.tiles:
        print('detecting tiles')
        arguments.tiles = ",".join(detect_tiles(arguments.bcl_dir))
    if not arguments.lanes:
        print('detecting lanes')
        arguments.lanes = ",".join(detect_lanes(arguments.bcl_dir))
    
    if not arguments.skip_mapping:
        # detect tiles to map
        while not os.path.isdir(os.path.join(arguments.bcl_dir, 'L00' + arguments.lanes.split(",")[0],
                                             'C2.1')):  # wait until first cycle is completed
            print('waiting for first cycle to finish')
            time.sleep(30)



        # run hilive
        print('starting hilive')
        call = assemble_hilive2_call(arguments, cyclelist)
        #call = assemble_hilive_call(arguments, cyclelist)
        run_subprocess(call, arguments.verbose)

    # parse background coverage file
    print('parsing background data')
    backgrounddict = unpickle_backgrounddict(
        backgroundfile=os.path.join(program_path, 'prelim_data', 'background_coverages.pickle'))

    # parse bsl list
    print('parsing bsl list')
    bsl_dict = parse_bsl_file(
        bsl_file=os.path.join(program_path, 'prelim_data', 'belgian_bsl_scores_always_highest.txt'))

    # parse phage list
    print('parsing phage list')
    phage_list = parse_phage_file(
        phage_file=os.path.join(program_path, 'prelim_data', 'phages.txt'))

    # start waiting for results from HiLive
    print('waiting for results from Mapper')

    for cycle in cyclelist:
        while not os.path.isfile(
                os.path.join(arguments.out_dir,
                             'hilive_out_cycle{}_undetermined.sam'.format(cycle))) and not os.path.isfile(
                os.path.join(arguments.out_dir, 'hilive_out_cycle{}_undetermined.bam'.format(cycle))):
            time.sleep(2)
        if arguments.out_format=='SAM':
            print('parsing samfile for cycle {}'.format(cycle))
            referencedict, uniquebasesdict, coveragedict, readcountdict = parse_samfile(
                os.path.join(arguments.out_dir, 'hilive_out_cycle{}_undetermined.sam'.format(cycle)),
                backgrounddict)
        elif not arguments.out_format=='SAM':
            print('parsing bamfile for cycle {}'.format(cycle))
            referencedict, uniquebasesdict, coveragedict, readcountdict = parse_bamfile(
                os.path.join(arguments.out_dir, 'hilive_out_cycle{}_undetermined.bam'.format(cycle)),
                backgrounddict)

        # generate data structure for result presentation
        print('generating data for tree of results for cycle {}'.format(cycle))
        try:
            copy_tree_structure(os.path.join(program_path, 'prelim_data'),
                                os.path.join(arguments.out_dir, 'output_cycle{}'.format(cycle)))
        except FileExistsError:
            shutil.rmtree(os.path.join(arguments.out_dir, 'output_cycle{}'.format(cycle)))
            copy_tree_structure(os.path.join(program_path, 'prelim_data'),
                                os.path.join(arguments.out_dir, 'output_cycle{}'.format(cycle)))

        # generate tree
        print('writing data for tree')
        print_treedata(readcountdict, uniquebasesdict, backgrounddict, bsl_dict, phage_list,
                       resultfile=os.path.join(arguments.out_dir, 'output_cycle{}'.format(cycle),
                                               'Collapsible_Tree_with_Amounts',
                                               'data',
                                               'treedata.csv'))
        print(time.time() - start)
        start = time.time()

        # generate plots
        print('generating plots')
        worker_args = []
        for refname in readcountdict.keys():
            try:
                worker_args.append((referencedict[refname], coveragedict[refname], backgrounddict[refname], refname,
                                    os.path.join(arguments.out_dir, 'output_cycle{}'.format(cycle),
                                                 'Collapsible_Tree_with_Amounts',
                                                 'images')))
            except KeyError:
                worker_args.append((referencedict[refname], coveragedict[refname], {}, refname,
                                    os.path.join(arguments.out_dir, 'output_cycle{}'.format(cycle),
                                                 'Collapsible_Tree_with_Amounts',
                                                 'images')))
        with multiprocessing.Pool(processes=int(arguments.num_threads)) as pool:
            pool.map(coverage_plot, worker_args)


if __name__ == '__main__':
    main()
