import NotebookImport
from Imports import *

cancers = {c: run.load_cancer(c) for c in run.cancers}

from Data.ProcessClinicalDataPortal import update_clinical_object

clinical = {c: cancer.load_clinical() for c, cancer in cancers.iteritems()}
for cancer,clin in clinical.iteritems():
    try:
        path  = OUT_PATH + '/Followup/' + cancer + '/'
        clinical[cancer] = update_clinical_object(clin, path)
    except:
        print cancer
        do_nothing = True

surv = pd.DataFrame({c: v.survival.survival for c,v in clinical.iteritems()})
surv = surv.stack()

codes = pd.Series({s[0]: s[1] for s in surv[:,'days'].index})
codes = codes.groupby(level=0).first()
codes.name = 'codes'

surv.index = surv.index.droplevel(2)
surv = surv.groupby(level=[0,1]).first()
surv = surv.ix[ti(surv[:,'days'] >= 7)]

surv_5y = pd.DataFrame({c: v.survival.survival_5y for c,v in clinical.iteritems()})
surv_5y = surv_5y.stack()
surv_5y.index = surv_5y.index.droplevel(2)
surv_5y = surv_5y.groupby(level=[0,1]).first()

#bad = ((surv_5y.unstack()['days'] < 8) * (surv_5y.unstack()['event'] == 1)) == False
#surv_5y = surv_5y.unstack().ix[bad].stack()
surv_5y = surv_5y.ix[ti(surv[:,'days'] >= 7)]

for c in clinical.values():
    c = c.artificially_censor(10)
surv_10y = pd.DataFrame({c: v.survival.survival_10y for c,v in clinical.iteritems()})
surv_10y = surv_10y.stack()
surv_10y.index = surv_10y.index.droplevel(2)
surv_10y = surv_10y.groupby(level=[0,1]).first()
surv_10y = surv_10y.ix[ti(surv[:,'days'] >= 7)]

for c in clinical.values():
    c = c.artificially_censor(3)
surv_3y = pd.DataFrame({c: v.survival.survival_3y for c,v in clinical.iteritems()})
surv_3y = surv_3y.stack()
surv_3y.index = surv_3y.index.droplevel(2)
surv_3y = surv_3y.groupby(level=[0,1]).first()
surv_3y = surv_3y.ix[ti(surv[:,'days'] >= 7)]

all_mut = pd.read_csv(OUT_PATH + '/MAFs/mega_maf.csv', index_col=0)
all_mut = all_mut[all_mut.Tumor_Sample_Barcode.apply(lambda s: s[13:16]) == '01A']

non_coding = ['Silent','RNA','IGR',"5'Flank", "3'UTR",'Intron',"5'UTR"]
all_mut = all_mut[all_mut.Variant_Classification.isin(non_coding)==False]
all_mut.Tumor_Sample_Barcode = all_mut.Tumor_Sample_Barcode.map(lambda s: s[:12])
all_mut = all_mut.groupby(['Tumor_Sample_Barcode', 'Hugo_Symbol']).size()
all_mut = all_mut.unstack().T.fillna(0)

mut = {c: cancer.load_data('Mutation') for c, cancer in cancers.iteritems() 
       if run.sample_matrix.MAF.ix[c] > 0}
hit_mat = pd.concat([m.df.unstack() for c,m in mut.iteritems()], axis=1).fillna(0)
hit_mat = hit_mat.groupby(axis=1, level=0).first()
hit_mat = pd.concat([hit_mat, all_mut], axis=1)
hit_mat = hit_mat.groupby(axis=1, level=0).sum() > 0

cn_all = {}
for c in codes.unique():
    try:
        cn_all[c] = FH.get_gistic_gene_matrix(run.data_path, c)
    except:
        print c
cn_all = pd.concat(cn_all.values(), axis=1)
cn_all = cn_all.groupby(axis=1, level=0).first()

non_embargo = ['AML','BRCA','KIRC','COAD','READ','COADREAD','SKCM','GBM',
               'HNSC','LUAD','LUSC','OV','STAD','THCA','UCEC']
#codes = codes[codes.isin(non_embargo)]

c2 = codes[codes.isin(ti(codes.value_counts() > 30))]
survival_and_stats(c2, surv_10y, upper_lim=10, figsize=(10,12))

fig, axs = subplots(2,1, figsize=(8,8))
ax = axs[0]

t = get_surv_fit(surv, codes)
t = t.sort([('5y Survival','Surv')])
t['Stats'].plot(kind='bar', ax=ax)
prettify_ax(ax)

ax2 = axs[1]
t = get_surv_fit(surv, codes)
tt = t['5y Survival'].sort('Surv')
b = (tt['Surv']).plot(kind='bar', ax=ax2, color='grey',
     yerr=[tt.Surv-tt.Lower, tt.Upper-tt.Surv], ecolor='black')
ax2.set_ylabel('5Y Survival')
ax2.set_yticks([0, .5, 1.])
prettify_ax(ax2)
fig.tight_layout()

fys = get_surv_fit(surv, codes)['5y Survival']['Surv'].order()
nn  = codes.isin(true_index((fys > .1) & (fys < .83)))
fys

age = pd.concat([c.clinical.daystobirth for c in clinical.values()])
age = -1*np.floor(age / 365.25)
a2 = a2 = pd.concat([c.clinical.age for c in clinical.values()])
age = age.combine_first(a2)
age = age.groupby(level=0).min()
age = age.fillna(age.median())
age.name = 'age'

survival_and_stats((age / 15).round(), surv_5y)

survival_and_stats(age.dropna() >= 85, surv_5y)

get_cox_ph(surv_5y, age >= 85, print_desc=True);

old = age >= 75
old.name = 'age_over_75'

survival_and_stats(old.ix[ti(age < 85)], surv_5y)

exp(.773), exp(.773)-(exp(.773-.139))

year = pd.concat([c.clinical.yearofinitialpathologicdiagnosis for c in clinical.values()])
year = year.groupby(level=0).first().replace('[Discrepancy]', nan).astype(float).dropna()

survival_and_stats(year <= 2000, surv_5y)

get_cox_ph_ms(surv_5y, year < 2000, [codes])

r_status = pd.concat([c.clinical.residualtumor for c in clinical.values() if 'residualtumor' 
                      in c.clinical])
r_status = r_status.replace(['[Not Evaluated]', '[Unknown]', 'RX'])
r_status = r_status.str.lower()
r_status = r_status.groupby(level=0).first()

survival_and_stats(r_status.dropna(), surv_5y)

get_cox_ph(surv_5y, r_status=='r2', print_desc=True);

exp(1.05), exp(1.05)-(exp(1.05-.18))

margin_status = pd.concat([c.clinical.marginstatus for c in clinical.values() if 'marginstatus' in c.clinical])
margin_status = margin_status.groupby(level=0).first()

survival_and_stats(margin_status.dropna(), surv_5y)

s = {}
for can,c in clinical.iteritems():
    if 'pathologicstage' in c.stage:
        s[can] = c.stage.pathologicstage
    elif 'clinicalstage' in c.stage:
        s[can] = c.stage.clinicalstage
stage = pd.concat(s.values())
stage = stage.groupby(level=0).first()
stage = stage.dropna().map(lambda s: s.replace('A','').replace('B','').replace('C',''))
stage = stage.dropna().map(lambda s: s.replace('1','').replace('2',''))
stage = stage.replace('stage iic','stage ii')
stage = stage.ix[age.index].fillna('missing')
stage = stage.str.lower()
stage = stage.replace(['[unknown]','stage x','stage tis','i or ii nos', 'stage 0'], nan).dropna()
stage.name = 'stage'

survival_and_stats(stage.ix[true_index(nn)].dropna(), surv_5y)

get_cox_ph(surv_5y, stage == 'stage iv', print_desc=True);

exp(.712), exp(.712)-(exp(.712-.0625))

mets = pd.concat([c.stage.pathologicm for can,c in clinical.iteritems() 
                  if 'pathologicm' in c.stage])
mets = mets.groupby(level=0).first()

mets = pd.concat([c.stage.pathologicm for can,c in clinical.iteritems() if 'pathologicm' in c.stage])
mets = mets.groupby(level=0).first()
mets = mets.dropna().map(lambda s: s.replace('a','').replace('b','').replace('c',''))
mets = mets.replace(['[Unknown]','MX'], nan).dropna()
mets = mets=='M1'
mets.name = 'metastasis'
survival_and_stats(mets, surv_5y)

survival_and_stats(combine(mets, stage=='stage iv'), surv_5y)

new_p53_calls = pd.read_csv('../Extra_Data/p53_calls_pancancer.csv', header=None, index_col=0, 
                            squeeze=True)
p53_new = (hit_mat.ix['TP53'].combine_first(new_p53_calls)) > 0

survival_and_stats(combine(new_p53_calls, hit_mat.ix['TP53']), surv_5y)

filters = pd.concat([age > 85, mets, r_status=='r2', stage=='stage iv', nn==False, 
                     codes.isin(['ACC','ESCA'])], 1)
keepers = true_index((filters > 0).sum(1) == 0)
keepers = keepers.intersection(surv_5y.unstack().index)
keepers = keepers.intersection(p53_new.index).intersection(cn_all.columns)
keepers = keepers.intersection(codes.index)

stage = stage.ix[keepers.union(stage.index)].fillna('missing')

len(keepers)

survival_and_stats(codes.ix[keepers].order(), surv_10y, upper_lim=10, figsize=(10,10))
plt.savefig(FIGDIR + 'pancan_supp.pdf', transparent=True)

p53_all = p53_new > 0
p53_mut = p53_all.ix[keepers]
p53_all.name = 'TP53'
p53_mut.name = 'TP53'
survival_and_stats(p53_mut, surv_5y)

get_cox_ph_ms(surv_5y, p53_mut, [codes, age, stage], interactions=None)

cn_all2 = cn_all.ix[[i for i in cn_all.index if i[2] in rna.df.index]]
pct_del = (cn_all2 < 0).sum() / (1.*len(cn_all2))
pct_amp = (cn_all2 > 0).sum() / (1.*len(cn_all2))
pct_altered = pct_amp + pct_del
pct_altered.name = 'CIN_pct_altered'
pct_del.name = 'pct_del'

survival_and_stats(to_quants(pct_altered, q=.33), surv)

del_3p_all = cn_all.ix['3p14.2'].median().round()
del_3p_all = del_3p_all < 0
del_3p_all = del_3p_all.map({True:-1, False:0})
del_3p = del_3p_all.ix[keepers].dropna()
del_3p.name = 'del_3p'
del_3p_all.name = 'del_3p'
survival_and_stats(del_3p, surv_5y)

get_cox_ph_ms(surv_5y, del_3p < 0, [codes, age, stage, pct_altered], interactions=None)

old.name = 'age_over_75'

fig, axs = subplots(3, 1, figsize=(3.5,4.5), sharey=True, sharex=True)

fmla = 'Surv(days, event) ~ del_3p + pct_del + age + codes + age_over_75'
fmla2 = 'Surv(days, event) ~ pct_del + age + codes + age_over_75'
k2 = keepers.diff(codes[codes.isin(['HNSC','LUSC'])].index)
pts = [k2, ti(p53_mut>0).intersection(k2), 
       ti(p53_mut==0).intersection(k2)]

color_list = ['grey', '#a1d99b','#9ecae1']

for i, pt in enumerate(pts):
    ax = axs[i]
    m1 = get_cox_ph(surv_5y, covariates=[del_3p.ix[pt].dropna() < 0, pct_del, stage, age, old, codes], 
                    formula=fmla, print_desc=False, interactions=False);
    m2 = get_cox_ph(surv_5y, covariates=[del_3p.ix[pt].dropna() < 0, pct_del, stage, age, old, codes], 
                    formula=fmla2, print_desc=False, interactions=False);
    print LR_test(m1, m2)
    ci = convert_robj(robjects.r.summary(m1)[7])
    haz = ci['exp(coef)']
    for j,h in enumerate(haz):
        ax.scatter(h, j, marker='s', s=100, color=color_list[i], 
                   edgecolors=['black'], zorder=10)
        ax.plot(*zip(*((ci.iloc[j]['lower .95'],j), (ci.iloc[j]['upper .95'],j))), 
                lw=3, ls='-', marker='o', dash_joinstyle='bevel', color=color_list[i])
    ax.axvline(1, ls='--', color='black')
    ax.set_xscale('log')
    ax.set_xbound(.8,1.5)
    ax.set_ybound(-.5,len(ci.index) - .5)
    ax.set_xticks([1])
    ax.set_xticklabels([1])
    ax.set_yticks(range(len(ci.index)))
    ax.set_yticklabels(ci.index)
    ax.set_ybound(-.5, 1.5)
    prettify_ax(ax)
axs[2].set_xlabel('Hazard Ratio')
fig.tight_layout()

combo = combine(del_3p < 0, p53_mut>0)
combo = combo.ix[keepers].dropna()
combo.name = 'TP53 / 3p14'
len(combo), len(combo.ix[codes[codes != 'HNSC'].index].dropna())

venn_pandas(del_3p_all < 0, p53_all>0)

len(combine(del_3p_all.ix[del_3p_all.index.diff(true_index(codes=='HNSC'))].dropna() < 0, 
                  p53_all>0))

fisher_exact_test(del_3p_all.ix[del_3p_all.index.diff(true_index(codes=='HNSC'))].dropna() < 0, 
                  p53_all>0)

venn_pandas(del_3p.ix[del_3p.index.diff(true_index(codes=='HNSC'))].dropna() < 0, 
            p53_mut>0)

len(combine(del_3p.ix[combo.index.diff(true_index(codes=='HNSC'))] < 0, p53_mut>0))

fisher_exact_test(del_3p.ix[combo.index.diff(true_index(codes=='HNSC'))] < 0, p53_mut>0)

combo_all = combine(p53_all>0, del_3p_all<0)
combo_all.name = 'TP53 / 3p14'
ct = pd.crosstab(combo_all, stage[stage != 'missing']=='stage iv').T
(ct.ix[True] / (1.*ct.sum())).order().plot(kind='bar');
plt.ylabel('% of Patients in Stage IV')
plt.xlabel('');

c = clinical['LGG'].clinical

combo_all = combine(p53_all>0, del_3p_all<0)
fisher_exact_test(combo_all=='both', stage[stage != 'M']=='stage iv')

get_cox_ph(surv_5y, combo=='both', print_desc=True);

get_cox_ph_ms(surv_5y, combo=='both', [codes, age, stage], interactions=None)

pts = true_index(codes != 'HNSC')

get_cox_ph(surv_5y, combo.ix[pts].dropna()=='both', print_desc=True);

f = (combo[combo.isin(['del_3p','both'])]=='both').ix[pts].dropna()
cox(f, surv_5y)

cox((combo[combo.isin(['TP53','both'])].dropna()=='both').ix[pts].dropna(),
    surv_5y).ix['LR'].ix['p']

pts_y = pts.intersection(true_index(age < 75))
log_rank((combo[combo.isin(['TP53','both'])]=='both').ix[pts_y].dropna(), surv_5y).ix['p']

cox((combo[combo.isin(['TP53','both'])]=='both').ix[pts_y].dropna(), surv_5y)

get_cox_ph_ms(surv_5y, (combo=='both').ix[true_index(codes!='HNSC')].dropna(), 
              [codes, stage, age, old], interactions=False)

two_hit = combo == 'both'
two_hit.name = 'two_hit'
p53_mut.name = 'TP53'
pct_del.name = 'pct_del'
old = age >= 75
old.name = 'old'

pts = keepers.diff(ti(codes=='HNSC'))
fmla0 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + old + age + codes + stage'
fmla1 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + TP53:pct_del + old + age + codes + stage'
fmla2 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + two_hit + TP53:pct_del + old + age + codes + stage'

m0 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, old, 
                                     age, pct_del, codes, stage], formula=fmla0)
m1 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, old, 
                                     age, pct_del, codes, stage], formula=fmla1)
m2 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, old, 
                                     age, pct_del, codes, stage], formula=fmla2)
LR_test(m2,m0), LR_test(m2, m1)

pts = keepers.diff(ti(codes.isin(['HNSC','LUSC'])))
#pts = keepers.diff(ti(codes=='HNSC'))
fmla0 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + codes + stage + age'
fmla1 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + TP53:pct_del + codes + stage + age'
fmla2 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + two_hit + TP53:pct_del + codes + stage + age'

m0 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, age, pct_del, codes, stage], formula=fmla0)
m1 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, age, pct_del, codes, stage], formula=fmla1)
m2 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, age, pct_del, codes, stage], formula=fmla2)
LR_test(m2,m0), LR_test(m2, m1)

pts = keepers.diff(ti(codes.isin(['HNSC','LUSC'])))
fmla0 = 'Surv(days, event) ~ TP53 + del_3p'
fmla1 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + TP53:pct_del'
fmla2 = 'Surv(days, event) ~ TP53 + del_3p + pct_del + two_hit + TP53:pct_del'

m0 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, old, age, pct_del], formula=fmla0)
m1 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, old, age, pct_del], formula=fmla1)
m2 = get_cox_ph(surv_3y, covariates=[p53_mut.ix[pts], del_3p < 0, two_hit, old, age, pct_del], formula=fmla2)
LR_test(m2,m0), LR_test(m2, m1)

cc = array(colors_th)

fig, ax = subplots(figsize=(5,3))
draw_survival_curve(combo, surv_5y, ax=ax,
                    colors=cc[[2,0,3,1]], ms=30, alpha=.7)
ax.get_legend().set_visible(False)
ax.set_ybound(0,1)
ax.set_xbound(0,5)
prettify_ax(ax)

combo.ix[true_index(codes!='HNSC')].dropna().ix[surv_5y.unstack().index].dropna().value_counts()

k2 = combo.index.intersection(true_index(codes!='HNSC'))

cc = array(colors_th)

fig, ax = subplots(figsize=(4,3))
draw_survival_curve(combo.ix[k2], surv_5y, ax=ax,
                    colors={'both': cc[0], 'TP53': cc[1], 'del_3p': cc[2], 'neither':cc[3]}, 
                    ms=30, alpha=.7)
ax.get_legend().set_visible(False)
ax.set_ybound(0,1)
ax.set_xbound(0,5)
prettify_ax(ax)
fig.tight_layout()
fig.savefig(FIGDIR + 'fig3e.pdf', transparent=True)

df = pd.concat([p53_mut, del_3p, combo.ix[k2], surv[:,'days'], surv[:,'event']],
                keys=['TP53 Mutation', '3p Deletion', 'TP53 / 3p', 'Days to Death/Censoring', 'Death Indicator'],
                axis=1).dropna()
df.to_csv(FIGDIR + 'fig3c.csv')

get_surv_fit(surv, combo.ix[ti(codes != 'HNSCC')], time_cutoff=3)

fmla = 'Surv(days, event) ~ feature + age + old + stage + strata(codes)'

m = {'TP53':'c', 'both':'d', 'neither':'a', 'del_3p':'b'}
cm = combo.map(m).ix[true_index(codes != 'HNSC')]

f = get_cox_ph(surv_3y, cm, interactions=False)

ci = convert_robj(robjects.r.summary(f)[7])
ci.index = ['3p','TP53','both']
n = ci.ix[0]*0 +1
n.name = 'neither'
ci  = ci.append(n)
ci = ci.ix[['neither','3p', 'TP53','both']]
ci_uni = ci

f = get_cox_ph(surv_3y, cm, [codes, age, old, stage], 
              formula=fmla, interactions=False)
ci = convert_robj(robjects.r.summary(f)[7])
ci = ci.ix[:3]
ci.index = ['3p','TP53','both']
n = ci.ix[0]*0 +1
n.name = 'neither'
ci  = ci.append(n)
ci = ci.ix[['neither','3p', 'TP53','both']]
ci_full = ci

fig, ax = subplots(1,1, figsize=(3,2.5))
ci = ci_full
haz = ci['exp(coef)']
b = haz.plot(kind='bar', ax=ax, color=cc[[3,2,1,0]],
     yerr=[haz - ci['lower .95'], ci['upper .95'] - haz], ecolor='black')
prettify_ax(ax)
ax.set_ylabel('hazard')
ax.set_ylim(0)
ax.axhline(y=1, ls='--', lw=3, color='black', alpha=.5)
fig.tight_layout()

fig, ax = subplots(1,1, figsize=(2.3,3))

ci = ci_full.ix[::-1]
haz = ci['exp(coef)']
color_list = colors_th

for j,h in enumerate(haz):
    ax.scatter(h, j, marker='s', s=100, color=color_list[j], 
               edgecolors=['black'], zorder=10, )
    ax.plot(*zip(*((ci.iloc[j]['lower .95'],j), (ci.iloc[j]['upper .95'],j))), 
            lw=3, ls='-', marker='o', dash_joinstyle='bevel', color=color_list[j])
ax.axvline(1, ls='--', color='black')
ax.set_xscale('log')
ax.set_xbound(.7,2.)
ax.set_ybound(-.5,len(ci.index) - .5)
ax.set_xticks([.75, 1, 1.25, 1.5, 2])
ax.set_xticklabels([.75, 1, 1.25, 1.5, 2])
ax.set_yticks(range(len(ci.index)))
ax.set_yticklabels('')
ax.set_ybound(-.5, 3.5)
prettify_ax(ax)
ax.set_xlabel('3 Year Hazard Ratio')
fig.tight_layout()
fig.savefig(FIGDIR + 'fig3f.pdf', transparent=True)

fig, axs = subplots(1,2, figsize=(5,2.5), sharey=True)
for i,ci in enumerate([ci_uni, ci_full]):
    ax = axs[i]
    haz = ci['exp(coef)']
    b = haz.plot(kind='bar', ax=ax, color=cc[[3,2,1,0]],
         yerr=[haz - ci['lower .95'], ci['upper .95'] - haz], ecolor='black')
    prettify_ax(ax)
    ax.set_ylabel('hazard')
    ax.set_ylim(0)
    ax.axhline(y=1, ls='--', lw=3, color='black', alpha=.5)
fig.tight_layout()

from itertools import combinations

c = list(combinations(cm.unique(),2))[0]

cm = combo.ix[true_index(codes!='HNSC')].dropna()
sig = pd.Series({c: get_cox_ph_ms(surv_5y, cm[cm.isin(c)], interactions=False)['LR']
                    for c in combinations(cm.unique(),2)})
sig

cm = combo.ix[true_index(codes!='HNSC')].dropna()
sig = pd.Series({c: get_cox_ph_ms(surv_3y, cm[cm.isin(c)], interactions=False)['LR']
                    for c in combinations(cm.unique(),2)})
sig

cm = combo.ix[true_index(codes!='HNSC')].dropna()
sig = pd.Series({c: get_cox_ph_ms(surv_3y, cm[cm.isin(c)], [codes, age], interactions=False)['LR']
                    for c in combinations(cm.unique(),2)})
sig

cm = combo.ix[true_index(codes!='HNSC')].dropna()
sig = pd.Series({c: get_cox_ph_ms(surv_5y, cm[cm.isin(c)], [codes, age, stage], interactions=False)['LR']
                    for c in combinations(cm.unique(),2)})
sig

pts = keepers.diff(ti(codes=='HNSC'))
cm = combo.ix[pts].dropna()
sig = pd.Series({c: get_cox_ph_ms(surv_3y, cm[cm.isin(c)], [codes, age, old, stage], 
                                  interactions=False)['LR']
                    for c in combinations(cm.unique(),2)})
sig

df = pd.concat([p53_mut, del_3p, codes, age, old, stage.ix[keepers].fillna('missing'), combo.ix[k2], surv[:,'days'], surv[:,'event']],
                keys=['TP53 Mutation', '3p Deletion',
                      'Tissue Type', 'Age', 'Age > 75', 'Stage'
                      'TP53 / 3p', 'Days to Death/Censoring', 'Death Indicator'],
                axis=1).ix[keepers]

df = pd.concat([p53_mut, del_3p, codes, age, old, stage.ix[keepers].fillna('M'), combo.ix[k2], surv[:,'days'], surv[:,'event']],
                keys=['TP53 Mutation', '3p Deletion',
                      'Tissue Type', 'Age', 'Age > 75', 'Stage',
                      'TP53 / 3p', 'Days to Death/Censoring', 'Death Indicator'],
                axis=1).dropna()
df.to_csv(FIGDIR + 'fig3c.csv')

f = {}
for cancer in codes.ix[combo.index].unique():
    try:
        c = combo[combo.isin(['both','TP53'])] == 'both'
        c = c.ix[ti(codes==cancer)].dropna()
        f[cancer] = get_cox_ph_ms(surv_3y, c, [age, old], interactions=False)
    except:
        print 'fail'

pd.DataFrame(f).T.sort('LR')

cc = ti(codes.ix[combo_all.index].value_counts() > 50)
ct = pd.DataFrame({c: fisher_exact_test(p53_all.ix[ti(codes==c)].dropna()>0, del_3p_all<0)
              for c in cc}).T
ct.dropna().sort('odds_ratio', ascending=False)

fig, axs = subplots(1, 7, figsize=(10,3), sharey=True)
cc = array(colors_th)
for i,c in enumerate(['HNSC','BRCA','LUAD', 'OV','KIRC','LGG','LUSC']):
    draw_survival_curve(combo.ix[true_index(codes==c)].dropna(), 
                        surv_3y, ax=axs[i], 
                        colors={'both': cc[0], 'TP53': cc[1], 'del_3p': cc[2], 'neither':cc[3]})
    axs[i].get_legend().set_visible(False)
    axs[i].set_ylim(.4,1.05)
    axs[i].set_xticks([0,1,2,3])
    prettify_ax(axs[i])
    axs[i].annotate(c, (2, .42))
fig.tight_layout()

fig, axs = subplots(1, 7, figsize=(10,3), sharey=True)
cc = array(colors_th)
for i,c in enumerate(['HNSC','BRCA','LUAD', 'OV','KIRC','LGG','LUSC']):
    f = combo[combo.isin(['both','TP53'])].ix[true_index(codes==c)].dropna()
    #f = f.ix[ti(age < 75)].dropna()
    draw_survival_curve(f, surv_5y, ax=axs[i], 
                        colors={'both': cc[0], 'TP53': cc[1], 'del_3p': cc[2], 'neither':cc[3]})
    axs[i].get_legend().set_visible(False)
    axs[i].set_ylim(.4,1.05)
    axs[i].set_xticks([0,1,2,3,4,5])
    prettify_ax(axs[i])
    axs[i].annotate(c, (2, .42))
fig.tight_layout()