#!/usr/bin/env python3

"""
This simple script combines multiple abundance files, generated by multiple executions of jgi_summarize_bam_contig_depths
on some subset of the total bam files.
This script will error if the columns between abundance files are named the same.
This script will error if the order or name or length of the contig rows is different between the abundance files,
as it is required that contigs are identical and against the same assembly fasta

Each file must have the same format and be tab separated:
contigName	contigLen	totalAvgDepth	contigs-1000.fastq.bam	contigs-1000.fastq.bam-var ...
NODE_1_length_5374_cov_8.558988	5404	14.2158	14.2158	16.817 ...

"""

import sys

def merge_files(outputFile, inputFiles):
    isTitle = True
    out = open(outputFile, "w")
    inputs = []
    for inputFile in inputFiles:
        inputs.append(open(inputFile, "r"))
    isEOF = False 
    headers = dict()
    while not isEOF:
        filen = 0
        contig = None
        contigLen = None
        total = 0.0
        vals = []
        for input_fh in inputs:
            filen += 1
            inputLine = input_fh.readline()
            inputLine = inputLine.rstrip('\n')
            if not inputLine:
                if filen == 1:
                    isEOF = True
                elif not isEOF:
                    raise(ValueError(f"input #{filen}({inputFiles[filen-1]}) has fewer lines than the first!"))
                continue
            if isEOF:
                raise(ValueError(f"input #{filen}({inputFiles[filen-1]}) has more lines than the first!"))
            input = inputLine.split('\t')
            if len(input) < 5 or len(input) % 2 != 1:
                raise(ValueError(f"Expected at least 5 tab-separated columns in input #{filen}({inputFiles[filen-1]}) and an odd number, got {len(input)} :'{inputLine}'"))
            if contig is None:
                contig = input[0]
                contigLen = input[1]
                if isTitle:
                    total = input[2]
            elif contig != input[0] or contigLen != input[1]:
                raise(ValueError(f"Invalid contig ({contig}) or contigLen ({contigLen}) in input #{filen}({inputFiles[filen-1]}): '{input}'"))
            if isTitle and total != input[2]:
                raise(ValueError(f"Invalid abundance file in input #{filen}({inputFiles[filen-1]}) expected 3rd column to be '{total}' not '{input[2]}'"))
            col=3
             
            while col < len(input)-1:
                if not isTitle:
                    total += float(input[col])
                else:
                    if input[col] in headers:
                        raise(ValueError(f"Detected duplicate column {input[col]} in input #{filen}({inputFiles[filen-1]})"))
                    headers[input[col]] = 1
                vals.append(input[col])
                vals.append(input[col+1]) 
                col+=2
        if contig is not None:
            cols = [contig, contigLen, "totalAvgDepth" if isTitle else total ]
            cols.extend(vals)     
            out.write("\t".join(map(str,cols)))
            out.write("\n")
        isTitle = False

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: merge_abundances.py output.txt input.txt ...")
        exit(1) 
    merge_files(sys.argv[1], sys.argv[2:])
