#!/usr/bin/env python

#-- Diego Montiel 2019
#-- Genetics Identificacion Group @ Erasmus MC  --
#-- YLeaf detection of Y-Haplogroups in Human genome DNA 
import time
import subprocess
import string
import random
import argparse, os
import pandas as pd
import numpy as np
from argparse   import ArgumentParser
from subprocess import Popen
import collections
import operator
import gc
pd.options.mode.chained_assignment = None  # default='warn'

def get_arguments():

    parser = ArgumentParser()    

    parser.add_argument("-fastq", "--Fastq",
            dest="Fastq", required=False, type=file_exists,
            help="Use raw FastQ files", metavar="PATH")

    parser.add_argument("-bam", "--Bamfile",
        dest="Bamfile", required=False, type=file_exists,
        help="input BAM file", metavar="PATH")    
    
    parser.add_argument("-ref", "--reference",  dest="reference",
            help="Build Genome Assembly \n [hg19-hg38]", metavar="STRING", required=True)    

    parser.add_argument("-out", "--output",
            dest="Outputfile", required=True,                        
            help="Folder name containing outputs", metavar="STRING")            
    
    parser.add_argument("-r", "--Reads_thresh",
            help="The minimum number of reads for each base",
            type=int, required=False,
            default=50)
    
    parser.add_argument("-q", "--Quality_thresh",
            help="Minimum quality for each read, integer between 10 and 39, inclusive \n [10-40]",
            type=int, required=True)
    
    parser.add_argument("-b", "--Base_majority",
            help="The minimum percentage of a base result for acceptance \n [50-99]",
            type=int, required=True)

    parser.add_argument("-t", "--Threads", dest="threads",
            help="Set number of additional threads to use [CPUs]",
            type=int,
            default=2)
            
    args = parser.parse_args()    
    return args

def get_frequency_table(mpileup):
    
    bases = ["A","T","G","C","+","-"]
    frequency_table = {}
    
    for i in mpileup.values:       
        fastadict = {"A":0,"T":0,"G":0,"C":0}                    
        sequence = i[9] #actual sequence                
        sequence = sequence.upper()         
        sequence = trimm_caret(sequence)                            
        sequence = sequence.replace("$", "")                   
        indel_pos = find_all_indels(sequence)
        ### Count number of indels
        indels = count_indels(sequence, indel_pos)        
        fastadict.update(indels)
        fastadict["-"] += sequence.count("*")
        ### Trimm Indels
        trimm_sequence = trimm_indels(sequence, indel_pos)        
        for seq in trimm_sequence:    
            if seq in fastadict:            
                fastadict[seq] +=1    
        frequency_table.update({i[3]:list(fastadict.values())})    
    df_frequency_table = pd.DataFrame.from_dict(frequency_table, orient='index')
    df_frequency_table.columns = bases      
    return df_frequency_table

def find_all_indels(s):
    find_all = lambda c,s: [x for x in range(c.find(s), len(c)) if c[x] == s]
    list_pos = []
    for i in find_all(s,"-"):
        list_pos.append(i)
    for i in find_all(s,"+"):
        list_pos.append(i)    
    return sorted(list_pos)

def count_indels(s, pos):    
    dict_indel = {"+":0,"-":0}    
    if pos == []:
        return dict_indel            
    if len(pos) > 0:
        for i in range(0,len(pos)): 
            try: # in case it is not a number but a base pair e.g. A
                dict_indel[s[pos[i]]] += int(s[pos[i]+1])                                                                        
            except ValueError:                
                dict_indel[s[pos[i]]] += 1
                continue                
    return dict_indel

def trimm_indels(s, pos):    
    ## Receives a sequence and trimms indels            
    if pos == []:
        return s
    u_sequence = ""  
    start =  pos[0]
    count = (start+1)    
    try: # in case it is not a number but a base pair e.g. A
        end = count+int(s[count])+1
    except ValueError:                
        end = start+1            
    u_sequence = s[:start]    
    if len(pos) > 1:
        for i in range(1,len(pos)):                      
            start = end                
            u_sequence += s[start:pos[i]]
            start = pos[i]
            count = (start+1)            
            try: # in case it is not a number but a base pair e.g. A
                end = count+int(s[count])+1  
            except ValueError:
                end = start+1                
            if pos[-1] == pos[i]:    
                #print(s[end:])
                u_sequence += s[end:]
    else:        
        u_sequence += s[end:]        
    return u_sequence

def trimm_caret(s):       
    find_all = lambda c,s: [x for x in range(c.find(s), len(c)) if c[x] == s]
    list_pos = []
    for i in find_all(s,"^"):
        list_pos.append(i)    
    if list_pos == []:
        return s
    i = 0    
    start = 0
    end = 0
    sequence = ""
    while i<len(s):
        if s[i] == "^":        
            end = i
            sequence += (s[start:end])                    
            start = i+1
        elif i >= list_pos[-1]+1:
            sequence += (s[list_pos[-1]+1:])
            break
        i+=1        
    return sequence

def execute_log(command):
            
    proc = subprocess.Popen(command, stdout=subprocess.PIPE)    
    try:
        out, err = proc.communicate(input=input)
        retcode = proc.returncode        
        return str(out), err, retcode
    except e:
        print("%r %r returned %r" % (command, e))
        raise    

def file_exists(file):
    """
    'Type' for argparse - checks that file exists but does not open.
    """
    if not os.path.exists(file):
        # Argparse uses the ArgumentTypeError to give a rejection message like:
        # error: argument input: file does not exist
        raise argparse.ArgumentTypeError("{0} does not exist".format(file))
    return file
        
def execute_mpileup(header, bam_file, pileupfile, Quality_thresh, folder):
            
    cmd = "samtools mpileup -AQ{} -r {} {} > {}".format(Quality_thresh, header, bam_file, pileupfile)        
    subprocess.call(cmd, shell=True)                    
    
def chromosome_table(bam_file,bam_folder,file_name):
    
    output = bam_folder+'/'+file_name+'.chr'
    tmp_output = "tmp_bam.txt"

    f = open(tmp_output, "w")
    subprocess.call(["samtools", "idxstats",bam_file], stdout=f)
    df_chromosome = pd.read_csv(tmp_output, sep='\t', header=None)
    total_reads = sum(df_chromosome[2])
    df_chromosome["perc"] = (df_chromosome[2]/total_reads)*100
    df_chromosome = df_chromosome.round(decimals=2)
    df_chromosome['perc'] = df_chromosome['perc'].astype(str) + '%'
    df_chromosome = df_chromosome.drop(columns=[1,3])
    df_chromosome.columns = ['chr','reads','perc']    
    df_chromosome.to_csv(output, index=None, sep="\t")
    
    cmd = "rm "+tmp_output
    subprocess.call(cmd, shell=True)

    if 'Y' in df_chromosome["chr"].values:
        return "Y", total_reads    
    elif 'chrY' in df_chromosome["chr"].values:
        return "chrY", total_reads    
    
def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for x in range(size))

def check_if_folder(path,ext):
    
    list_files = []
    if os.path.isdir(path):
        dirpath = os.walk(path)
        for dirpath, dirnames, filenames in dirpath:
            for filename in [f for f in filenames if f.endswith(ext)]:
                files = os.path.join(dirpath, filename)
                list_files.append(files)
        return list_files
    else:
        return [path]

def get_folder_name(path_file):

    folder      = path_file.split('/')[-1]
    folder_name = os.path.splitext(folder)[0]        
    return folder_name

def create_tmp_dirs(folder):

    flag = True
    if os.path.isdir(folder):    
        while(flag):
                cmd = 'mkdir '+folder
                subprocess.call(cmd, shell=True)                
                flag = False
                return True   
    else:
        cmd = 'mkdir '+folder
        subprocess.call(cmd, shell=True)        
        return True
    
def extract_haplogroups(path_Markerfile, Reads_thresh, Base_majority, 
                        path_Pileupfile, log_output, fmf_output, Outputfile):    

    print("Extracting haplogroups...")
    Markerfile = pd.read_csv(path_Markerfile, header=None, sep="\t")
    Markerfile.columns = ["chr", "marker_name", "haplogroup", "pos", "mutation", "anc", "der"]
    Markerfile = Markerfile.drop_duplicates(subset='pos', keep='first', inplace=False)    

    Pileupfile = pd.read_csv(path_Pileupfile, header=None, sep="\t", dtype = {0:str,1:int,2:str,3:int,4:str,5:str})
    Pileupfile.columns = ['chr', 'pos', 'refbase', 'reads', 'align', 'quality']
    
    log_output_list = []
    log_output_list.append("Total of reads: "+str(len(Pileupfile))) #total of reads

    intersect_pos = np.intersect1d(Pileupfile['pos'], Markerfile['pos'])
    Markerfile = Markerfile.loc[Markerfile['pos'].isin(intersect_pos)]
    Markerfile = Markerfile.sort_values(by=['pos'])
    Pileupfile = Pileupfile.loc[Pileupfile['pos'].isin(intersect_pos)]

    Pileupfile = Pileupfile.drop(['chr'], axis=1)
    df = pd.merge(Markerfile, Pileupfile, on='pos')

    del [[Pileupfile,Markerfile]]
    gc.collect()
    Pileupfile=pd.DataFrame()
    Markerfile=pd.DataFrame()

    log_output_list.append("Valid markers: "+str(len(df))) #valid markers provided

    index_belowzero = df[df["reads"] == 0].index
    df_belowzero = df[df.index.isin(index_belowzero)]
    df_belowzero = df_belowzero.drop(['refbase','align','quality'], axis=1)
    df_belowzero["called_perc"] = "NA"
    df_belowzero["called_base"] = "NA"
    df_belowzero["state"] = "NA"
    df_belowzero["Description"] = "Position with zero reads"

    df = df[~df.index.isin(index_belowzero)]

    df_freq_table = get_frequency_table(df)
    df_freq_table = df_freq_table.drop(['+','-'], axis=1)
    df = df.drop(['refbase','align','quality'], axis=1)

    list_col_indices = np.argmax(df_freq_table.values, axis=1)
    called_base = df_freq_table.columns[list_col_indices]
    total_count_bases = np.sum(df_freq_table.values, axis=1)
    max_count_bases = np.max(df_freq_table, axis=1)
    called_perc = round((max_count_bases/total_count_bases)*100,1)
    
    bool_anc = np.equal(np.array(called_base), df["anc"].values)
    bool_der = np.equal(np.array(called_base), df["der"].values)

    bool_list_anc = np.where(bool_anc,'A','D')
    bool_list_anc = bool_list_anc.astype('object')
    bool_list_der = np.where(bool_der,'D','A')
    bool_list_der = bool_list_der.astype('object')
    bool_list_state = np.equal(bool_list_anc, bool_list_der)

    df["called_perc"] = np.array(called_perc, dtype=int)
    df["called_base"] = called_base
    df["state"] = bool_list_anc
    df["bool_state"] = bool_list_state

    df_discordantgenotype = df[~bool_list_state]
    df_discordantgenotype = df_discordantgenotype.drop(["bool_state"], axis=1)
    df_discordantgenotype["state"] = "NA"
    df_discordantgenotype["Description"] = "Discordant genotype"

    columns_fmf = df_discordantgenotype.columns

    ## read threshold
    df_readsthreshold = df[df["reads"] < Reads_thresh]
    df_readsthreshold["Description"] = "Below read threshold"

    ## filter by base percentage 
    df_basemajority = df[df["called_perc"] < Base_majority]
    df_basemajority["Description"] = "Below base majority"

    index_to_remove = np.concatenate((df_basemajority.index,df_readsthreshold.index, 
                                      df_discordantgenotype.index), axis=None)

    df_fmf = pd.concat([df_belowzero,df_readsthreshold, df_basemajority, df_discordantgenotype], axis=0, sort=True)    
    df_fmf = df_fmf[columns_fmf]

    df_out = df[~df.index.isin(index_to_remove)]
    df_out = df_out.drop(["bool_state"], axis=1)
    df_out = df_out.sort_values(by=['haplogroup'], ascending=True)

    log_output_list.append("Markers with zero reads: "+str(len(df_belowzero))) 
    log_output_list.append("Markers below the read threshold {"+str(Reads_thresh)+"}: "+str(len(df_readsthreshold))) 
    log_output_list.append("Markers below the base majority threshold {"+str(Base_majority)+"}: "+str(len(df_basemajority))) 
    log_output_list.append("Markers with discordant genotype: "+str(len(df_discordantgenotype))) 
    log_output_list.append("Markers without haplogroup information: "+str(len(df_fmf))) 
    log_output_list.append("Markers with haplogroup information: "+str(len(df_out))) 

    with open(log_output, "a") as log:
        for marker in log_output_list:
            log.write(marker)
            log.write("\n")

    del [[df_basemajority,df_belowzero, df_discordantgenotype, df_readsthreshold, df_freq_table, df]]
    gc.collect()
    df_basemajority=pd.DataFrame()
    df_belowzero=pd.DataFrame()
    df_discordantgenotype=pd.DataFrame()
    df_readsthreshold=pd.DataFrame()
    df_freq_table=pd.DataFrame()
    df = pd.DataFrame()
        
    df_out = df_out[["chr","pos","marker_name","haplogroup","mutation","anc","der","reads","called_perc","called_base","state"]]
    df_fmf.to_csv(fmf_output, sep="\t", index=False)
    df_out.to_csv(Outputfile, sep="\t", index=False)

def samtools(threads, folder, folder_name, bam_file, Quality_thresh):
    
    #Change this file to Concatenate from the parameter you give    
    args.Markerfile =  'bin/yleaf/Position_files/'+args.reference+'.txt'                    
        
    start_time = time.time()    
    if not os.path.exists(bam_file+'.bai'): 
        
        bam_file_order = folder+"/"+folder_name+".order.bam"                        
        cmd = "samtools sort -m 2G -@ {} {} > {}".format(threads, bam_file, bam_file_order)        
        print("\tSorting Bam file...")        
        subprocess.call(cmd, shell=True)
        cmd = "samtools index -@ {} {}".format(threads, bam_file_order)        
        subprocess.call(cmd, shell=True)                
        bam_file = bam_file_order
                    
    file_name  = folder_name
    Outputfile = folder+"/"+folder_name+".out"    
    log_output = folder+"/"+folder_name+".log"
    fmf_output = folder+"/"+folder_name+".fmf"
    pileupfile = folder+"/"+folder_name+".pu" 

    header,total_reads = chromosome_table(bam_file,folder,file_name)
    execute_mpileup(header, bam_file, pileupfile, Quality_thresh, folder)                      
    print("--- %.2f seconds in run PileUp ---" % (time.time() - start_time))    
    
    start_time = time.time()            
    extract_haplogroups(args.Markerfile, args.Reads_thresh, args.Base_majority, 
                            pileupfile, log_output, fmf_output, Outputfile)
    
    cmd = "rm {};".format(pileupfile)
    subprocess.call(cmd, shell=True)

    print("--- %.2f seconds in extracting haplogroups --- " % (time.time() - start_time) )
    print("--- %.2f seconds to run Yleaf  ---" % (time.time() - whole_time))
    return Outputfile

def logo():
    print(r"""
    
           |
          /|\          
         /\|/\    
        \\\|///   
         \\|//  
          |||   
          |||    
          |||    

        """)

def identify_haplogroup(path_file, output):
    
    cmd = "python3 bin/yleaf/predict_haplogroup.py -input {} -out {}".format(path_file, output)
    print(cmd)
    subprocess.call(cmd, shell=True)                    
    
if __name__ == "__main__":
            
    whole_time = time.time()
    print("\tErasmus MC Department of Genetic Identification \n\n\tYleaf: software tool for human Y-chromosomal \n\tphylogenetic analysis and haplogroup inference v2.0\n")
    logo()
    args = get_arguments()    
    app_folder = os.path.dirname(os.path.realpath(__file__))    
    sam_file    = ''
    folder_name = ''                    
    out_path    = args.Outputfile   
    threads     = args.threads 
    cwd         = os.getcwd()        
    if os.path.isabs(out_path):
        out_folder = out_path
    else:
        if cwd == "/":
            out_folder = out_path
        else:
            out_folder = cwd+"/"+out_path        
    if create_tmp_dirs(out_folder):        
        if args.Fastq:                
                files = check_if_folder(args.Fastq,'.fastq')
                for path_file in files:            
                    print("Starting...")
                    print(path_file)
                    bam_file = path_file
                    folder_name = get_folder_name(path_file)
                    folder = os.path.join(app_folder,out_folder,folder_name)                                                
                    if create_tmp_dirs(folder):                                            
                        start_time = time.time()
                        sam_file = folder+"/"+folder_name+".sam"               
                        reference = app_folder+"/"+args.reference +"/"+args.reference+".fa"
                        fastq_cmd = "bwa mem -t {} {} {} > {}".format(threads, reference, path_file, sam_file)
                        print(fastq_cmd)
                        subprocess.call(fastq_cmd, shell=True)
                        print("--- %s seconds in Indexing reads to reference ---" % (time.time()-start_time))                
                        start_time = time.time()
                        bam_file = folder+"/"+folder_name+".bam"             
                        cmd = "samtools view -@ {} -bS {} | samtools sort -@ {} -m 2G -o {}".format(threads, sam_file, threads, bam_file)
                        print(cmd)
                        subprocess.call(cmd, shell=True)
                        print("--- %s seconds in convertin Sam to Bam ---" % (time.time()-start_time))                                
                        cmd = "samtools index -@ {} {}".format(threads, bam_file)
                        subprocess.call(cmd, shell=True)        
                        output_file = samtools(threads, folder, folder_name, bam_file, args.Quality_thresh)                                    
                hg_out = out_folder+"/"+out_path+".hg"
                identify_haplogroup(out_folder, hg_out)                                                                        
        elif args.Bamfile:                
                files = check_if_folder(args.Bamfile,'.bam')
                for path_file in files:            
                    print("Starting...")
                    print(path_file)
                    bam_file = path_file
                    folder_name = get_folder_name(path_file)
                    folder = os.path.join(app_folder,out_folder,folder_name)                            
                    if create_tmp_dirs(folder):                                            
                        output_file = samtools(threads, folder, folder_name, bam_file, args.Quality_thresh)                        
                hg_out = out_folder+"/"+out_path+".hg"
                identify_haplogroup(out_folder, hg_out)                                                                        
    else:
        print("--- Yleaf failed! please check inputs... ---")