#!/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 Neoepitopes 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 neoepitopes")

    model.add_argument("-n", "--neoepitopes", type=str, help="neoepitopes file")

    args = model.parse_args()

    neoepitope_to_geneID = defaultdict()
    with open(args.neoepitopes, "r") as f:
        # skip header
        f.readline()
        for line in f:
            split = line.split(",")
            neoepitope_to_geneID[split[0]] = split[1]

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

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

            with open(prediction + ".csv", "w") as f:
                f.write("peptide,geneID,ic50\n")
                for neoepitope, pair in neoepitope_to_geneID_ic50.items():
                    f.write(neoepitope + "," + pair[0] + "," + pair[1] + "\n")


if __name__ == "__main__":
    main()
