#!/usr/bin/python3
# Author: Temesgen H. Dadi
# Freie University of Berlin
# AG Algorithmische Bioinformatik
# July 2014

__author__ = 'Temesgen H. Dadi (temehi@gmail.com)[Original Version], Enrico Seiler[Adaptions for Donald]'
__version__ = '1.0.1 Donald'
__date__ = 'July 2016'
__maintainer__ = 'enrico.seiler@fu-berlin.de'

import sys, math, time
import os
from math import floor
import pkgutil
import pkg_resources
import threading
import argparse

import gps as gps
import taxonomy as taxonomy


parser = argparse.ArgumentParser(description =
    '{0}\n{1}\n{0}\nVersion: {2}\nDate: {3}\nAuthors: {4}\nMaintainer: {5}\n{6}'.format(
        '-------------------------------', 'MicrobeGPS-cmd 1.01. Donald',
        __version__, __date__, __author__, __maintainer__,
        'A comand line tool for MicrobeGPS by Martin S. Lindner for the use with Donald.'),
    formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument('input_alignment', help = 'The file path containing the alignment file (sam, bam, cram).')
parser.add_argument('output_path', help = 'The file path of output file.')
parser.add_argument('--gi_mapping_file', default = '<default>', help = 'The file path containing the mapping of gi/accession to organisms.')
reg_param = parser.add_argument_group('Raw filter parameters', 'Parameters used on the raw filter step of MicrobeGPS')
reg_param.add_argument('--min_support', default = 50, type = int, help = 'Minimum Genome Support discards all genomes with less than \n a specified number of reads.(default: 50)')
reg_param.add_argument('--max_read_matches', default = 80, type = int, help = 'Max. Read Matches discards all reads matching to more \n than a specified number of genomes. These reads are \n considered as uninformative.(default: 80)')
reg_param.add_argument('--max_mapping_error', default = 1.00, type = float, help = 'Max. Read Mapping Error discards all read mappings with \n error above a specified value. The error is the fraction of not \n matching positions to the total read length.(default: 1)')
reg_param.add_argument('--mapping_error_quantile', default = 1.00, type = float, help = 'Mapping Error Quantile keeps the specified quantile of the best \n matching reads. Default: 1 keeps all reads.(default: 1)')
reg_param = parser.add_argument_group('Quality Filtering parameters ', 'Parameters used on the Quality Filtering step of MicrobeGPS')
reg_param.add_argument('--min_unique_reads', default = 20, type = int, help = 'Minimum Genome Support discards all genomes with less\n than a specified number of reads.(default: 50)')
reg_param.add_argument('--max_homogeneity', default = 0.6, type = float, help = 'Minimum Genome Support discards all genomes with less \n than a specified number of reads.(default: 0.600)')
reg_param = parser.add_argument_group('Calculate Candidates parameters', 'Parameters used on the Calculate Candidates step of MicrobeGPS')
reg_param.add_argument('--min_validity', default = 0.001, type = float, help = 'Discards all organisms below the specified validity \n threshold.(default: 0.001)')
reg_param.add_argument('--coverage_similarity', default = 0.20, type = float, help = 'sets the required relative coverage difference between \n all matches for a read.(default: 0.20)')
reg_param.add_argument('--fraction_shared_usr', default = 0.20, type = float, help = 'Msets for a reference the minimum required fraction of \n Unique Source Reads (USR) to be grouped with another reference.\n (default: 0.20)')
reg_param.add_argument('--fraction_shared_all', default = 0.60, type = float, help = 'sets for a reference the minimum required fraction of \n reads to be grouped with another reference.(default: 0.60)')
args = parser.parse_args()

t_start = time.time()

########## LOAD DATA ##########
print('Started loading data.\n--- This may take a while.')
alignment_file = args.input_alignment
try:
    target_table,read_table = gps.read_alignment_file_pysam(alignment_file)
    #target_table,read_table = gps.read_sam_files_no_pysam(alignment_file)
except Exception as e:
    print(str(e))
    raise
else:
    print('Done loading data!')

t = time.time() - t_start
t_min = floor(t/60)
t_sec = int(t - t_min*60)
print('---Found {} reads mapping to {} targets.\nFinished Loading Data. ({} min, {} sec)'.format(len(read_table),len(target_table), t_min, t_sec))

########## RAW FILTERING ##########
print('Started raw filtering.')
try:
    gps.filter_raw(target_table,read_table,
            max_matches= args.max_read_matches,
            min_support = args.min_support,
            max_error = args.max_mapping_error,
            qual_percentile =args.mapping_error_quantile)
    gps.calculate_mapping_statistics(target_table)
except Exception as e:
    print(str(e))
    raise
else:
    print('Done raw filtering!')

########## CALCULATE REFERENCE TABLE ##########
print('Calculating Reference Table.')
try:
    if args.gi_mapping_file == '<default>':
        gi_map = open(os.path.join(os.path.dirname(sys.argv[0]), 'data/bact.catalog'))
    elif args.gi_mapping_file.endswith('.gz'):
        gi_map = gzip.open(args.gi_mapping_file, 'r')
    else:
        gi_map = open(args.rt,'r')

    names_map_file = open(os.path.join(os.path.dirname(sys.argv[0]), 'data/names.dmp'))
    names_map = taxonomy.parse_names_dmp(names_map_file)
    ref_table = gps.get_reference_table_NCBI(target_table,read_table, gi_map)
except Exception as e:
    print(str(e))
    raise
else:
    print('Done Calculating Reference Table!')

# ########## QUALITY FILTERING ##########
print('Started quality filtering.')
try:
    def custom_filt1(ref):
        if not ref.unique >= args.min_unique_reads:
            return False
        if not ref.cov_homog <= args.max_homogeneity:
            return False
        return True
    gps.filter_ref_table(ref_table,read_table,custom_filt1)
except Exception as e:
    print(str(e))
    raise
else:
    print('Done quality filtering')

# ########## CALCULATE CANDIDATES ##########
print('Started calculating candidates.')
try:
    from multiprocessing.pool import ThreadPool
    pool = ThreadPool(5)
    pool.map(gps.calculate_valcov_one,iter(ref_table.values()))
    pool.close()
    def custom_filt2(ref):
        if not ref.validity >= args.min_validity or not ref.coverage > 0:
            return False
        return True
    gps.filter_ref_table(ref_table,read_table,custom_filt2)
    usr_table = gps.extract_USR(ref_table,read_table,args.coverage_similarity)
    usr_mat,n2i,i2n = gps.get_read_matrix(usr_table,ref_table)
    all_mat,n2ia,i2na = gps.get_read_matrix(read_table,ref_table)
    resort = [n2ia[i2n[i]] for i in range(len(n2i))]
    all_mat = all_mat[:,resort]
    group_dict = gps.create_groups_dc(usr_mat,all_mat,n2i,args.fraction_shared_usr)
    candidates = gps.enrich_groups(group_dict,ref_table,read_table,i2n)
except Exception as e:
    print(str(e))
    raise
print('--- Found {} candidates.\nFinished candidate list.'.format(len(candidates)))

########## SAVE RESULTS ##########

# # # # #  Save to CSV # # # # #
fname = args.output_path
print('Writing to csv file: {}'.format(fname))
total_unique_reads = sum(ref_table[key].unique for key in ref_table)
outf = open(fname, 'w')
outf.write('\t'.join(['Candidate', 'Name', 'Abundance', 'Num. Reads', 'Unique Reads', 'Coverage', 'Validity', 'Homogeneity', 'Mapping Error']))
outf.write('\n')
i = 1
for g,grp in enumerate(candidates):
    n_mem = len(grp.members)
    sorted_members = sorted(grp.members, key=lambda x: -grp.members[x].reads)
    j = 1
    for m in sorted_members:
        mbr = grp.members[m]
        coverage = mbr.coverage if not math.isnan(mbr.coverage) else 0.
        validity = mbr.validity if not math.isnan(mbr.validity) else 0.

        outf.write('\t'.join([str(i)+'.'+str(j), m, str(float(mbr.reads)/len(read_table)), str(mbr.reads), str(mbr.unique), str(coverage), str(validity), str(mbr.cov_homog), str(mbr.map_qual)]))
        outf.write('\n')
        j += 1
    i += 1
outf.close()
print('Finished writing to csv file: {}'.format(fname))

t = time.time() - t_start
t_min = floor(t/60)
t_sec = int(t - t_min*60)
print('Calculation took {} min {} sec'.format(t_min,t_sec))
