#!/usr/bin/env python3
"""
Concatenate all TRAL output files (generated by run_tral.py) from a directory into one file.

Author: Max Verbiest
Contact: max.verbiest@zhaw.ch
"""

import argparse
import os


def concatenate_files(file_list, output_file):
    """Concatenates all files from file_list into one big file: output_file

    file_list (list[str]):
                list of file handles of run_tral.py output files to be concatenated
    output_file (str):
                file where concatenated information will be deposited
    """

    # get header
    with open(file_list[0], "r") as f:
        header_line = f.readline().strip()
        if not header_line.startswith("begin"):
            raise ValueError("File layout different than expected!")
        header_line = "ID\t{}\n".format(header_line)

    with open(output_file, "w") as o:
        o.write(header_line)
        for file_name in file_list:
            handle = file_name.split("/")[-1]
            prot_id = handle.split(".")[0]
            with open(file_name, "r") as f:
                for line in f:
                    if not line.startswith("begin"):
                        if line.endswith("\n"):
                            o.write("{}\t{}".format(prot_id, line))
                        else:
                            o.write("{}\t{}\n".format(prot_id, line))


def parser():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--dir", "-d", type=str, required=True, help="Directory with tral output files."
    )

    return parser.parse_args()


def main():
    args = parser()
    # get handles of all .tsv files from a directory
    filenames = [os.path.join(args.dir, file) for file in os.listdir(args.dir) if file.endswith(".tsv")]
    # make output file name, will be deposited in the same dir as input files
    output_file = "{}/merged.tsv".format(args.dir)
    concatenate_files(filenames, output_file)


if __name__ == "__main__":
    main()
