import subprocess
import re


def netmhc_pep_prep(filepath, svfusions):
    """
    :param svfusions: a list of svfusion classes
    :param filepath: the output file
    :return: None
    """
    print("Preparing peptides for netMHC.")
    with open(filepath, 'w') as f:
        for svfusion in svfusions:
            for neoepitope in svfusion.neoepitopes:
                f.write(neoepitope + '\n')


def netmhc_run(netmhcpath, peppath, alleles, outpath):
    """
    :param netmhcpath: absolute path of the netmhc execution file
    :param peppath: input peptides file
    :param alleles: HLA alleles separated by ,
    :param outpath: outfile
    :return: None
    """
    print("Running netMHC.")
    cmd = netmhcpath \
        + ' -a ' + alleles \
        + ' -f ' + peppath \
        + ' -inptype 1' \
        + ' -BA' \
        + ' > ' + outpath
    p = subprocess.Popen(cmd, shell=True)
    p.wait()


def netmhc_reload(resultpath):
    """
    :param resultpath: the result of netmhc
    :return: a dictionary {neoepitope: {allele: [affinity, rank, FILTER]}}
    """
    pep_dic = {}
    with open(resultpath, 'r') as f:
        for line in f:
            if line.startswith("    1  HLA"):
                tmpline = re.split(r'\s+', line)
                allele = tmpline[2]
                pep = tmpline[3]
                affinity = float(tmpline[13])
                rank = float(tmpline[14])
                if pep not in pep_dic:
                    pep_dic[pep] = {}
                pep_dic[pep][allele] = [affinity, rank, 'FILTER']
    return pep_dic


def netmhc_filter(pep_dic, aff_thre, rank_thre):
    """
    :param pep_dic: the dictionary for neoepitopes generated by netmhc_load
    :param aff_thre: threshold for binding affinity
    :param rank_thre: threshold for rank percentile
    :return: a dictionary in which passed neoepitopes are labeled PASS
    """
    print("Filtering netMHC result by binding affinity.")
    for pep in pep_dic:
        for allele in pep_dic[pep]:
            if pep_dic[pep][allele][0] < aff_thre and pep_dic[pep][allele][1] < rank_thre:
                pep_dic[pep][allele][2] = 'PASS'
    return pep_dic


