import os, math
import numpy as np
import pandas as pd
pd.set_option('display.width', 120)
pd.set_option('max_colwidth', 90)
%matplotlib inline
import matplotlib.pyplot as plt
import epitopepredict as ep
from epitopepredict import sequtils, analysis, base, tepitope, plotting, utilities
from IPython.display import display, HTML
import seaborn as sns
from importlib import reload
sns.set_style("ticks", {'axes.facecolor': '#F7F7F7', 'axes.grid': False,
                        'legend.frameon':True, 'legend.fontsize':12})
sns.set_context("notebook", font_scale=1.4)

"""
Explanation: Predictions for IFN T cell responses in healthy individuals infected with Mtb in a South African cohort.
The authors used eleven representative Mtb antigens and a large panel of previously defined Mtb epitopes. 50– 75% of the total response in a cohort of 63 individuals. As expected, responses were highly heterogeneous, with responses to a total of 125 epitopes detected
Reference:
Lindestam Arlehamn, C.S. et al., 2016. A Quantitative Analysis of Complexity of Human Pathogen-Specific CD4 T Cell Responses in Healthy M. tuberculosis Infected South Africans. PLOS Pathogens, 12(7), p.e1005760. Available at: http://dx.plos.org/10.1371/journal.ppat.1005760.
End of explanation
"""


#sette mtb epitopes
sette = pd.read_csv('sette_SA_MTB_epitopes.csv')
proteins = sorted(sette.name.unique())
mtb = sequtils.genbank_to_dataframe(base.mtb_genome,cds=True)
mtb = mtb[mtb.locus_tag.isin(sette.name)]

df = pd.merge(sette, mtb[['locus_tag','translation']],left_on='name',right_on='locus_tag')
df['start'] = df.apply( lambda x: x.translation.find(x.sequence), 1 )
df['end'] = df.start+df.sequence.str.len()
#for i,r in df[df.start==-1].iterrows():
#    print (r.locus_tag, r.sequence, r.translation)
df=df.drop('translation',1)
df=df[df.start>=0]
prots = df.name.unique()
exp = df
print (df)

donoralleles = pd.read_csv('sa_mtb_donor_alleles.csv')
donoralleles = donoralleles.set_index('Donor ID')
donoralleles = donoralleles.apply(lambda x: 'HLA-'+x)
#print donoralleles

x = donoralleles.apply(lambda x : x.value_counts())
x = x.sum(1).sort_values(ascending=False)
x = x[x>=2]
drballeles = x[x.index.str.contains('DRB')]
drballeles = list(drballeles.index)
#print drballeles
dqalleles = x[x.index.str.contains('DQ')]
dqalleles = list(dqalleles.index)


"""
Explanation: get exp dataset
End of explanation
"""


Pt=ep.get_predictor('tepitope')
Pt.predict_sequences(mtb, alleles=drballeles,path='mtbsa_tepitope', overwrite=False, threads=10)
Pt.load(path='mtbsa_tepitope')
Pn=ep.get_predictor('netmhciipan')
Pn.predict_sequences(mtb, alleles=drballeles,path='mtbsa_netmhciipan', overwrite=False, threads=10)
Pn.load(path='mtbsa_netmhciipan')

print (proteins)
ax = plotting.plot_tracks([Pn,Pt],name=proteins[-3],
                     n=5, cutoff=10, kind='tracks',                   
                     legend=True, cutoff_method='rank',            
                     figsize=(14,10))


"""
Explanation: run predictions
End of explanation
"""


bn = Pn.get_binders(cutoff=5,cutoff_method='rank',limit=20)#, name='Rv3875')
bt = Pt.get_binders(cutoff=5,cutoff_method='rank',limit=20)#, name='Rv3875')
ax = utilities.venndiagram([bn.peptide,bt.peptide], ['netmhciipan','tepitope'])

"""
Explanation: overlap of top ranking binders
End of explanation
"""


pbn = Pn.promiscuous_binders(n=5, cutoff=10, cutoff_method='rank')
pbt = Pt.promiscuous_binders(n=5, cutoff=10, cutoff_method='rank')
ax = utilities.venndiagram([pbn.peptide,pbt.peptide], ['netmhciipan','tepitope'])

#compare to exp results?
reload(analysis)
P=Pt
labels=['netmhciipan','tepitope']
res=[]
for P in [Pn, Pt]:
    #binders by median rank    
    rb = P.promiscuous_binders(cutoff_method='rank',cutoff=30,n=5)
    #all binders by promiscuity, just use low perc to get all
    pball = P.promiscuous_binders(n=1, cutoff=0.4).sort_values(['alleles','median_rank'],ascending=False)
    #reload(analysis)
    a = analysis.prediction_coverage(exp, pball, key='sequence', perc=50)
    print (a)
    b = analysis.prediction_coverage(exp, rb, key='sequence', perc=50)
    res.append({'all':a,'ranked':b})
    
x = pd.DataFrame(res,index=['netmhciipan','tepitope'])
print (x)
x.plot(kind='barh')


import difflib

def get_hits(n):
    cutoffs = range(2,40,3)
    res = []
    for c in cutoffs:
        rbn = Pn.promiscuous_binders(cutoff_method='rank',cutoff=c,n=n)
        rbt = Pt.promiscuous_binders(cutoff_method='rank',cutoff=c,n=n)
        rboth = pd.concat([rbn,rbt]).drop_duplicates(['peptide','name'])
        #print (rbt[:15])
        #print (len(rbn),len(rbt),len(rboth))
        df=exp#[exp.name==name].copy()
        def find_matches(x, p):
            return len(difflib.get_close_matches(x.sequence, p, n=10, cutoff=.6))

        df.loc[:,'hits1'] = df.apply(lambda x: find_matches(x, rbn.peptide),1)        
        f1 = len(df[df.hits1>0])/len(df)
        df.loc[:,'hits2'] = df.apply(lambda x: find_matches(x, rbt.peptide),1)
        f2 = len(df[df.hits2>0])/len(df)
        #print (df[:10])
        b = (len(df[(df.hits1>0) | (df.hits2>0)]))/len(df)
        res.append({'cutoff':c,'len1':len(rbn),'len2':len(rbt),'lenb':len(rboth),
                    'netMHCIIpan':f1, 'tepitope':f2, 'both':b})
    res = pd.DataFrame(res)
    return res

print (len(proteins))
n=5
res = get_hits(n)
print (res)
ax=res.plot(x='len1',y='netMHCIIpan',lw=2)
res.plot(x='len2',y='tepitope',lw=2,ax=ax)
res.plot(x='lenb',y='both',lw=2,ax=ax)
plt.xlabel('predicted binders')
plt.tight_layout()
plt.savefig('mtb_sette_hits.png',dpi=150)

fig,axs=plt.subplots(3,3,figsize=(12,10))
axs=axs.flat
i=0
for n in range(5,15):
    res = get_hits(n)    
    res.plot(x='len1',y='netMHCIIpan',lw=2,ax=axs[i])
    res.plot(x='len2',y='tepitope',lw=2,ax=axs[i])
    i+=1

r=Pn.promiscuous_binders(cutoff_method='rank',cutoff=10,n=7,name='Rv3875')
r

"""
Explanation: overlap of promiscuous binders
End of explanation
"""


