#!/usr/bin/env python

# Written by Lukas Heumos and released under MIT license.

import argparse
import sys
import logging
from collections import defaultdict

# logging setup
console = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
console.setFormatter(formatter)
LOG = logging.getLogger("Postprocess Peptides MHCNuggets")
LOG.addHandler(console)
LOG.setLevel(logging.INFO)


def main():
    model = argparse.ArgumentParser(description="Postprocess Neoepitopes predicted by MHCNuggets")

    model.add_argument("-i", "--input", type=str, nargs="*", help="predicted class 2 peptides")

    model.add_argument("-p", "--peptides_seq_ID", type=str, help="peptides to seq_ID csv file")

    model.add_argument("-o", "--output", type=str, help="output file name")

    args = model.parse_args()

    peptide_to_geneID = defaultdict()
    with open(args.peptides_seq_ID, "r") as f:
        for line in f:
            split = line.split(",")
            peptide_to_geneID[split[0]] = split[1]

    for prediction in args.input:
        peptide_to_geneID_ic50 = defaultdict()

        with open(prediction, "r") as f:
            # skip header
            f.readline()
            for line in f:
                split = line.split(",")
                peptide = split[0].strip()
                geneID = peptide_to_geneID[peptide].strip()
                ic50 = split[1].strip()
                peptide_to_geneID_ic50[peptide] = (geneID, ic50)

            with open(args.output, "w") as f:
                f.write("peptide,geneID,ic50\n")
                for peptide, pair in peptide_to_geneID_ic50.items():
                    f.write(peptide + "," + pair[0] + "," + pair[1] + "\n")


if __name__ == "__main__":
    main()
