#!/usr/bin/env python

import argparse

## At the moment this script is used to generate the coverage of each MAG in a single sample, but could be "improved" to include a list of samples etc"
parser = argparse.ArgumentParser(description='''tabulate the bases mapped for the MAGs of interest as generated by "scafstats" flag in the bbmap command''')
parser.add_argument('--mappingfile', help='the resulting file specified by scafstats=<file> in the bbmap command, the pythonic field 8 represents the bases recruited to the MAG scaffold')
parser.add_argument('--out', help='the file to write the awesome results including the MAG name and the sum of bases recruited to the mag')
parser.add_argument('--ids', help='the list of scaffold ids - required because if there are no reads recruited, then a MAG might not be included in the output')
args=parser.parse_args()

# The outfile to put results
outfile = open(args.out, 'w')
outfile.write("a"+str(args.mappingfile)+'\t'+"stat"+'\n')

# The empty list to hold all the mag ids associated with those found in the "ids" file that you provided as input
all_MAGs = {}
count = 0
for line in open(args.ids, 'r'):
    x = line.strip().split(":")
    all_MAGs[x[0]] = "n_"+str(count)
    count += 1
for key in all_MAGs.keys():
    print(key)

# A dictionary to hold the contigs and their stats in the mapping file
all_data_dict = {}
# This populates the dictionary of each scaffold and the stats generated by bbmap
for line in open(args.mappingfile, 'r'):
    x = line.strip().split('\t')
    if "Bin" in x[0]:
        all_data_dict[x[0]] = x[0:len(x)]

#A dictionary to store the results prior to writing
mags_and_stats_dict = {}
#For each mag in the collection, sum the total lenght of all reads mapped to each contig
for MAG in all_MAGs:
    ssum=[]
    for key in all_data_dict.keys():
        if MAG == key.split(":")[0]:
            print(MAG,all_data_dict[key])
            ssum.append(int(all_data_dict[key][8]))

    mags_and_stats_dict[MAG] = sum(ssum)


print("this is the total number of MAGs in you set", len(all_MAGs.keys()))
print("number of mags found in mapping file",len(mags_and_stats_dict.keys()))

for MAG in all_MAGs.keys():
    if MAG not in mags_and_stats_dict.keys():
        print("this MAG was not found in mapping file", MAG)
        mags_and_stats_dict[MAG] = 0

for MAG in mags_and_stats_dict.keys():
    outfile.write(MAG+'\t'+str(mags_and_stats_dict[MAG])+'\n')

 
