import os 
if os.getcwd().endswith('Validation'):
    os.chdir('..')

import NotebookImport
from Setup.Imports import *

"""
Explanation: Process Normalized Cross-Study Dataset
For differential expression analysis, we processed all of the data together on a high memory node.  This was done using the minfi package in R, and all data were normalized using the preprocessQuantile function.  For a full script of the R pipeline see the Methylation_Normalization_MINFI notebook.
Imports and helper functions from Imports notebook.
End of explanation
"""


path = '/cellar/users/agross/TCGA_Code/Methlation/data/Validation/'

df_hiv =  pd.read_csv(path + 'BMIQ_Hannum2.csv', index_col=0)

cell_type = pd.Series(df_hiv.columns, index=df_hiv.columns)
cell_type = cell_type.map(lambda s: s.strip()[-3:])

df_hiv.shape

cell_type.value_counts()

"""
Explanation: Read in Validation Data
End of explanation
"""


hiv_ann = pd.read_excel(path + 'Methylome human NEUvs WB.xlsx', 
                        sheetname='HIV', skiprows=2, index_col=0)
control_ann = pd.read_excel(path + 'Methylome human NEUvs WB.xlsx',
                            sheetname='Negative', skiprows=2, index_col=0)
ann = pd.concat([hiv_ann, control_ann], keys=['HIV+','HIV-'])
ann = ann.reset_index()
ann = ann.rename(columns={'Age (years)': 'age', 'level_0':'HIV'})
ann = ann.set_index('PATID')

hiv_ann = pd.read_excel(path + 'Methylome human NEUvs WB.xlsx', 
                        sheetname='HIV', skiprows=2, index_col=0)
control_ann = pd.read_excel(path + 'Methylome human NEUvs WB.xlsx',
                            sheetname='Negative', skiprows=2, index_col=0)
ann = pd.concat([hiv_ann, control_ann], keys=['HIV+','HIV-'])
ann = ann.reset_index()
ann = ann.rename(columns={'Age (years)': 'age', 'level_0':'HIV'})
ann = ann.set_index('PATID')

m2 = pd.read_excel('/cellar/users/agross/Downloads_Old/Methylome Study Demographic Data.xlsx',
              skiprows=5, index_col=0)
age2 = m2['Age (years)']
ann = ann.ix[ann.index.union(age2.index)]
ann['age'] = age2.combine_first(ann.age)
ann['HIV'] = ann.HIV.fillna('HIV+')

age_n = pd.Series({'{}_Neu'.format(i[3:]): v for i,v in ann.age.iteritems()})
age_t = pd.Series({'{}_CD4'.format(i[3:]): v for i,v in ann.age.iteritems()})
age = pd.concat([age_n, age_t])

hiv_n = pd.Series({'{}_Neu'.format(i[3:]): v for i,v in ann.HIV.iteritems()})
hiv_t = pd.Series({'{}_CD4'.format(i[3:]): v for i,v in ann.HIV.iteritems()})
hiv = pd.concat([hiv_n, hiv_t])
hiv = hiv == 'HIV+'

import Setup.DX_Imports as dx
import Parallel.Age_HIV_Features as fx

neu_corr = pd.read_csv('/cellar/users/agross/Data/Methylation_Controls/neutrophils_age_corr.csv',
                            index_col=0, squeeze=True)
cd4_corr = pd.read_csv('/cellar/users/agross/Data/Methylation_Controls/CD4T_age_corr.csv',
                            index_col=0, squeeze=True)

r2 = ti(fx.rr > 1)
r4 = r2.intersection(ti(cd4_corr.abs() > .2))
#r4 = cd4_corr.abs().ix[r2].dropna().order()[-10000:].index
len(r4)

df = df_hiv.ix[r4, ti(cell_type=='CD4')].dropna(1)
dd = logit_adj(df)
m = dd.mean(1)
s = dd.std(1)
df_norm = dd.subtract(m, axis=0).divide(s, axis=0)

U,S,vH = frame_svd(df_norm)

p = S ** 2 / sum(S ** 2)
p[:5]

fig, ax = subplots(1,1, figsize=(4,3))
rr1 = -1*vH[0]

sns.regplot(*match_series(age, rr1.ix[ti(hiv==0)]),
            ax=ax, label='HIV+', ci=None)
sns.regplot(*match_series(age, rr1.ix[ti(hiv>0)]),
            ax=ax, label='Control', ci=None)
ax.set_ylabel('First PC', size=12)
ax.set_xlabel('Chronological age (years)', size=14)

ax.set_yticks([0])
ax.axhline(0, ls='--', lw=2.5, color='grey', zorder=-1)
ax.set_xbound(23,70)
prettify_ax(ax)
fig.tight_layout()

r2 = ti(fx.rr > 1)
r4 = r2.intersection(ti(neu_corr.abs() > .2))
#r4 = neu_corr.abs().ix[r2].dropna().order()[-10000:].index
len(r4)

df = df_hiv.ix[r4, ti(cell_type=='Neu')].dropna(1)
dd = logit_adj(df)
m = dd.mean(1)
s = dd.std(1)
df_norm = dd.subtract(m, axis=0).divide(s, axis=0)

U,S,vH = frame_svd(df_norm)
p = S ** 2 / sum(S ** 2)
p[:5]

((vH[0] - vH[0].mean()) / vH[0].std()).abs().order().tail()

U,S,vH = frame_svd(df_norm[[p for p in df_norm.columns if p != '383_Neu']])
p = S ** 2 / sum(S ** 2)
p[:5]

fig, ax = subplots(1,1, figsize=(4,3))
rr2 = -1*vH[0]

sns.regplot(*match_series(age, rr2.ix[ti(hiv==0)]),
            ax=ax, label='HIV+', ci=None)
sns.regplot(*match_series(age, rr2.ix[ti(hiv>0)]),
            ax=ax, label='Control', ci=None)
ax.set_ylabel('First PC', size=12)
ax.set_xlabel('Chronological age (years)', size=14)

ax.set_yticks([0])
ax.axhline(0, ls='--', lw=2.5, color='grey', zorder=-1)
ax.set_xbound(23,70)
#ax.set_ybound(-.25,.25)
prettify_ax(ax)
fig.tight_layout()

fig, axs = subplots(1,2, figsize=(8,3))

ax = axs[0]
sns.regplot(*match_series(age, rr1.ix[ti(hiv==0)]),
            ax=ax, label='HIV+', ci=None)
sns.regplot(*match_series(age, rr1.ix[ti(hiv>0)]),
            ax=ax, label='Control', ci=None)
ax.set_ylabel('First PC\n(CD4+)', size=12)

ax = axs[1]
sns.regplot(*match_series(age, rr2.ix[ti(hiv==0)]),
            ax=ax, label='HIV+', ci=None)
sns.regplot(*match_series(age, rr2.ix[ti(hiv>0)]),
            ax=ax, label='Control', ci=None)
ax.set_ylabel('First PC\n(Neutrophil)', size=12)

for ax in axs:
    ax.legend(loc='lower right', frameon=True, fancybox=True)
    ax.set_xlabel('Chronological age (years)', size=14)

    ax.set_yticks([0])
    ax.axhline(0, ls='--', lw=2.5, color='grey', zorder=-1)
    ax.set_xbound(23,70)
    prettify_ax(ax)
    
fig.tight_layout()
fig.savefig(FIGDIR + 'PCA_sorted.png', dpi=200)

"""
Explanation: Read in annotations
End of explanation
"""


