#!/usr/bin/env python
import argparse
import logging
import pandas as pd
import numpy as np
import sys

from collections import defaultdict
from mhcflurry import Class1AffinityPredictor

#logging setup
console = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console.setFormatter(formatter)
LOG = logging.getLogger("MHCFlurry Predict mztab")
LOG.addHandler(console)
LOG.setLevel(logging.INFO)

def parse_mztab(identified_peptides_file):
    """
    parses an mztab file and returns all identified proteins
    :param identified_peptides_file: path to the mztab file
    :return: identified proteins
    """
    mztab = open(identified_peptides_file)
    mztab_read = mztab.readlines()
    mztab.close()

    seq_geneIDs = defaultdict(str)
    for line in mztab_read:
        if line.startswith("PEP"):
            content = line.split('\t')
            seq = content[1]
            geneID = content[2]
            if not 'U' in seq and not 'X' in seq and not 'Z' in seq and not 'J' in seq and not 'B' in seq:
                seq_geneIDs[seq] = geneID

    return seq_geneIDs

#List of alleles supported by mhcflurry
supported_alleles = "A*01:01,A*02:01,A*02:02,A*02:03,A*02:05,A*02:06,A*02:07,A*02:11,A*02:12,A*02:16,A*02:17,A*02:19," \
                    "A*02:50,A*03:01,A*11:01,A*23:01,A*24:02,A*24:03,A*25:01,A*26:01,A*26:02,A*26:03,A*29:02,A*30:01," \
                    "A*30:02,A*31:01,A*32:01,A*33:01,A*66:01,A*68:01,A*68:02,A*68:23,A*69:01,A*80:01,B*07:01,B*07:02," \
                    "B*08:01,B*08:02,B*08:03,B*14:02,B*15:01,B*15:02,B*15:03,B*15:09,B*15:17,B*18:01,B*27:02,B*27:03," \
                    "B*27:04,B*27:05,B*27:06,B*35:01,B*35:03,B*37:01,B*38:01,B*39:01,B*39:06,B*40:01,B*40:02,B*42:01," \
                    "B*44:02,B*44:03,B*45:01,B*46:01,B*48:01,B*51:01,B*53:01,B*54:01,B*57:01,B*58:01,B*83:01,C*03:03," \
                    "C*04:01,C*05:01,C*06:02,C*07:02,C*08:02,C*12:03,C*14:02,C*15:02".split(",")

#read provided allotypes
alleles=sys.argv[-3].split(";")

#extract and verify alleles
unsupported_alleles=[a for a in alleles if a not in supported_alleles]
alleles=[a for a in alleles if a in supported_alleles]

if unsupported_alleles:
    for allele in unsupported_alleles:
        LOG.warning("Allele: " + allele + " is not supported by MHCFlurry!")
if not alleles:
    LOG.warning("Submitted alleles are not supported or formatting of input.tsv is not correct!")

# read identified peptides
seqs_to_geneID = parse_mztab(sys.argv[-2])

if len(seqs_to_geneID) > 0:
    # call mhcflurry
    for allele in alleles:
        predictor = Class1AffinityPredictor.load()
        df_pred = predictor.predict_to_dataframe(allele=allele, peptides=seqs_to_geneID.keys())
        df_pred.insert(1, 'geneID', pd.Series(np.array(seqs_to_geneID.values())))
        df_pred.to_csv(allele + '_' + sys.argv[-1])
else:
    op=open(sys.argv[-1],'w')
    op.close()
