#!/usr/bin/env python

import sys
import argparse
import os.path
import pandas as pd
import csv
import gzip
import statistics

from Bio import SeqIO


def parse_args(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('-b', '--bins'         , required=True, nargs="+", metavar='FILE'                             , help="Bins: FASTA containing all contigs.")
    parser.add_argument('-d', '--depths'       , required=True           , metavar='FILE'                             , help="(Compressed) TSV file containing contig depths for each sample: contigName, contigLen, totalAvgDepth, sample1_avgDepth, sample1_var [, sample2_avgDepth, sample2_var, ...].")
    parser.add_argument('-a', '--assembler'    , required=True                           , type=str                   , help="Assembler name.")
    parser.add_argument('-i', '--id'           , required=True                           , type=str                   , help="Sample or group id.")
    parser.add_argument('-m', '--binner'       , required=True                           , type=str                   , help="Binning method.")
    return parser.parse_args(args)
# Processing contig depths for each binner again, i.e. not the most efficient way, but ok

def main(args=None):
    args = parse_args(args)

    # load contig depths for all samples into dict (could use pandas as well)
    sample_names = []
    dict_contig_depths = {}
    with gzip.open(args.depths, "rt") as infile:
        reader = csv.reader(infile, delimiter = "\t")
        # process header
        header = next(reader)
        for sample in range(int((len(header)-3)/2)):
            col_name = header[3+2*sample]
            # retrieve sample name: "<assembler>-<id>-<other sample_name>.bam"
            sample_name = col_name[len(args.assembler)+1+len(args.id)+1:-4]
            sample_names.append(sample_name)
        # process contig depths
        for row in reader:
            contig_depths = []
            for sample in range(int((len(row)-3)/2)):
                contig_depths.append(float(row[3+2*sample]))
            dict_contig_depths[str(row[0])] = contig_depths

    # Initialize output files
    n_samples = len(sample_names)
    with open(args.assembler + "-" + args.binner + "-" + args.id + "-binDepths.tsv", 'w') as outfile:
        print("bin", '\t'.join(sample_names), sep='\t', file=outfile)

    # for each bin, access contig depths and compute mean bin depth (for all samples)
    for file in args.bins:
        all_depths = [[] for i in range(n_samples)]

        if file.endswith('.gz'):
            with gzip.open(file, 'rt') as infile:
                for rec in SeqIO.parse(infile,'fasta'):
                    contig_depths = dict_contig_depths[rec.id]
                    for sample in range(n_samples):
                        all_depths[sample].append(contig_depths[sample])
        else:
            with open(file, "rt") as infile:
                for rec in SeqIO.parse(infile,'fasta'):
                    contig_depths = dict_contig_depths[rec.id]
                    for sample in range(n_samples):
                        all_depths[sample].append(contig_depths[sample])

        binname = os.path.basename(file)
        with open(args.assembler + "-" + args.binner + "-" + args.id + "-binDepths.tsv", 'a') as outfile:
            print(binname, '\t'.join(str(statistics.median(sample_depths)) for sample_depths in all_depths), sep='\t', file=outfile)


if __name__ == "__main__":
    sys.exit(main())
