# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # 14_promoter_promoter
# 4/6/2021
#
# corresponds to new fig 2F
# A. promoter promoter: Basically for a promoter-promoter loop, we want to know if both genes are being expressed in the loop or if one is acting purely as an enhancer and it's own gene is not being expressed
#
# new fig 2G
# correlate # loops with target gene expression
#

# +
import pandas as pd
import os, glob
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

save_dir = '../data/processed/fig2'

# +
normal_tissues = ['Airway','Astrocytes','Bladder','Colon','Esophageal',
'GDSD0',
'GDSD3',
'GDSD6',
'GM12878',
'HMEC',
'Melanocytes',
'Ovarian',
'Pancreas',
'Prostate',
'Renal',
'Thyroid',
'Uterine']


# -

loop_files = glob.glob('../data/interim/merged/loops/*csv')
loop_files

promoter_anchor_files = glob.glob('../data/interim/annon/promoter_anchors/*bed')
promoter_anchor_files

# +
rna_df = pd.read_csv('../data/interim/rna/tissue_tpm_sym.csv',index_col=0)
print(rna_df.columns)
rna_df[:5]
def get_exp(gene, tissue):
    try:
        return rna_df.loc[gene, tissue]
    except:
        return 0
    
# colnames(rna_df)


# -

file = promoter_anchor_files[0]
tissue = os.path.basename(file).split('_')[1]
print(tissue)
p_anchor_df = pd.read_table(file,header=None)
p_anchor_df.columns = ['chr_g','start_g','stop_g',
                                              'gene','chr_a','start_a','stop_a',
                                               'anchor_name','dist']
p_anchor_df['tpm'] = p_anchor_df.gene.apply(lambda g: get_exp(g, tissue))
p_anchor_df = p_anchor_df[['anchor_name','gene','tpm']]
loop_file = os.path.join('../data/interim/merged/loops/',tissue+'.loops.csv')
loop_df = pd.read_csv(loop_file, index_col=0)

# promoter loop = loop with a gene in both anchor sides
loop_df = (loop_df
    .merge(p_anchor_df,how='left', left_on='source',right_on='anchor_name' )
    .merge(p_anchor_df, how='left', left_on='target',right_on='anchor_name' )
    .dropna()
)


# +
def get_type(row, thres=1):
    """
    types:
    1: both promoters are expression
    2: only one promoter is expressed
    3: no promoters are expressed
    """
    if row['tpm_x']>thres:
        if row['tpm_y']>thres:
            return 1
        else:
            return 2
    else:
        if row['tpm_y']>thres:
            return 2
        else:
            return 3

loop_df['P_P_type'] = loop_df.apply(get_type,axis=1)
# -

tissue_P_P_counts = pd.DataFrame(loop_df.P_P_type.value_counts().sort_index()).T
tissue_P_P_counts.index = [tissue]
tissue_P_P_counts

sns.clustermap(loop_df[['gene_x','gene_y','P_P_type']][loop_df.P_P_type==1]
               .drop_duplicates()
               .pivot(index='gene_x',columns='gene_y',values='P_P_type')
               .fillna(0))


# +
def get_first_gene(row):
    if row['gene_x']<row['gene_y']:
        return row['gene_x']
    else:
        return row['gene_y']
    
def get_second_gene(row):
    if row['gene_x']==row['gene_1']:
        return row['gene_y']
    else:
        return row['gene_x']


# +
# %%time
tissue_P_P_counts_all = pd.DataFrame()
tissue_P_P_gene_df_all = pd.DataFrame()

for file in sorted(promoter_anchor_files):
    tissue = os.path.basename(file).split('_')[1]
    if tissue not in normal_tissues:
        continue
    print(tissue)
    # get anchor to gene annotation and annotate with correct tissue tpm
    p_anchor_df = pd.read_table(file,header=None)
    p_anchor_df.columns = ['chr_g','start_g','stop_g',
                                                  'gene','chr_a','start_a','stop_a',
                                                   'anchor_name','dist']
    p_anchor_df['tpm'] = p_anchor_df.gene.apply(lambda g: get_exp(g, tissue))
    p_anchor_df = p_anchor_df[['anchor_name','gene','tpm']]
    
    # add to looping information
    loop_file = os.path.join('../data/interim/merged/loops/',tissue+'.loops.csv')
    loop_df = pd.read_csv(loop_file, index_col=0)
    # promoter loop = loop with a gene in both anchor sides
    loop_df = (loop_df
        .merge(p_anchor_df,how='left', left_on='source',right_on='anchor_name' )
        .merge(p_anchor_df, how='left', left_on='target',right_on='anchor_name' )
        .dropna()
        )
    # get P-P type info 1: both exp, 2: one gene exp, 3: no gene exp
    loop_df['P_P_type'] = loop_df.apply(get_type,axis=1)
    loop_df['gene_1'] = loop_df.apply(get_first_gene,axis=1)
    loop_df['gene_2'] = loop_df.apply(get_second_gene,axis=1)
    loop_df_sel = loop_df[['gene_1', 'gene_2', 'P_P_type']]
    loop_df_sel['tissue'] = tissue
    
    # add info to global table 
    loop_1= loop_df[['count','gene_x','tpm_x']]
    loop_1.columns = ['count','gene','tpm']
    loop_2= loop_df[['count','gene_y','tpm_y']]
    loop_2.columns = ['count','gene','tpm']
    loop_bi = (pd.concat([loop_1, loop_2])
                .groupby('gene').agg({'tpm':min, 'count':['count','sum']})
                .reset_index())
    loop_bi.columns = ['gene','tpm','num_loops','loop_count']
    loop_bi['log_tpm'] = np.log10(loop_bi.tpm+1)
    loop_bi['log_loop_count'] = np.log10(loop_bi.loop_count+1)
    
    
    # add count info
    tissue_P_P_counts = pd.DataFrame(loop_df.P_P_type.value_counts().sort_index()).T
    tissue_P_P_counts.index = [tissue]
    tissue_P_P_counts_all = pd.concat([tissue_P_P_counts_all, tissue_P_P_counts])
    
    # add gene pairs info
    tissue_P_P_gene_df_all = pd.concat([tissue_P_P_gene_df_all, loop_df_sel])

# +
# loop_1 = loop_df[['anchor_name_x','P_P_type', 'tpm_x']]
# loop_1.columns = ['anchor_name','P_P_type', 'tpm']
# loop_2= loop_df[['anchor_name_y','P_P_type', 'tpm_y']]
# loop_2.columns = ['anchor_name','P_P_type', 'tpm']
# loop_bi_anchor = (pd.concat([loop_1, loop_2])
#             .groupby('anchor_name').agg({'P_P_type':['count','min','median','mean','std','max'],
#                                         'tpm':['min','median','mean','std','max']})
#                 .reset_index()
#                  .fillna(0))
# loop_bi_anchor.columns = ['anchor_name',
#                           'count_P_P','min_type','median_type','mean_type','std_type','max_type',
#                          'min_tpm','median_tpm','mean_tpm','std_tpm','max_tpm']
# loop_bi_anchor.plot.scatter(x='max_type',y='mean_type')
# loop_bi_anchor.plot.scatter(x='min_tpm',y='max_type')
# -

tissue_P_P_counts_all['tot_P_P'] = tissue_P_P_counts_all.sum(axis=1)

print(tissue_P_P_counts_all.sum(axis=0))
tissue_P_P_counts_all.sum(axis=0).div(tissue_P_P_counts_all.tot_P_P.sum())

tissue_P_P_counts_all.columns = ['both_exp','one_exp','none_exp','total_P_P']
tissue_P_P_counts_all.to_csv(os.path.join(save_dir,'tissue_P_P_counts_all.csv' ))

tissue_P_P_counts_all

tissue_P_P_gene_df_all.to_csv(os.path.join(save_dir,'tissue_P_P_gene_df_all.csv' ))

# +
# tissue_P_P_gene_df_all.pivot(index=['gene_1','gene_2'], values='P_P_type', columns='tissue')
# -

# # does # loops correlate with expression
# no...

loop_1= loop_df[['count','gene_x','tpm_x']]
loop_1.columns = ['count','gene','tpm']
loop_2= loop_df[['count','gene_y','tpm_y']]
loop_2.columns = ['count','gene','tpm']
loop_bi = (pd.concat([loop_1, loop_2])
            .groupby('gene').agg({'tpm':min, 'count':['count','sum']})
            .reset_index())
loop_bi.columns = ['gene','tpm','num_loops','loop_count']
loop_bi['log_tpm'] = np.log10(loop_bi.tpm+1)
loop_bi['log_loop_count'] = np.log10(loop_bi.loop_count+1)
display(loop_bi[:5])
# loop_bi = loop_bi[(loop_bi.tpm>1) & (loop_bi.loop_count>0)]
loop_bi['tpm'] = loop_bi.tpm.clip(0,10)
loop_bi['loop_count'] = loop_bi.loop_count.clip(0,100)
loop_bi.plot.scatter(x='tpm',y='num_loops')
loop_bi.plot.scatter(x='log_tpm',y='num_loops')
loop_bi.plot.scatter(x='log_tpm',y='log_loop_count')

# %%time
data_all = pd.read_csv('../data/processed/tissue_crms/all_count_comb_overall.csv')
data_all[:5]

data_all = data_all.iloc[:,[0,1,2,3,4]]
data_all.columns = ['gene','tissue','tpm','num_loop_counts', 'num_loops']
data_all = data_all.fillna(0)
data_all.sort_values('gene',inplace=True)



# +
sns.set_style("whitegrid")
def get_scatter_gene(gene, df=data_all, get_text=False):
    df_sel = df[df.gene==gene].reset_index()
    ax = sns.scatterplot(data=df_sel, x='tpm',y='num_loops')
    if get_text:
        added_x = max(df_sel.tpm)/10
        for line in range(0,df_sel.shape[0]):
             plt.text(df_sel.tpm[line]+added_x, df_sel.num_loops[line], df_sel.tissue[line], horizontalalignment='left', size='medium', color='black', weight='semibold')
    ax.set_xlabel('TPM')
    ax.set_ylabel('# loops')
    ax.set_title('Loop-expression correlation: '+ gene)
    plt.show()
    return df_sel

def get_per_gene_r(gene, df=data_all):
    df_sel = df[df.gene==gene].reset_index()
    r =  np.corrcoef(df_sel.tpm, df_sel.num_loops)[0,1]
    if np.isnan(r):
        return 0
    return r


# -

data_all[(data_all.tpm>20)&(data_all.tissue=='GDSD6')].gene.unique()#.shape

df_sel = get_scatter_gene('ADAMTS18')

df_sel = data_all[data_all.gene=='ADAMTS13'].reset_index()
display(df_sel)
np.corrcoef(df_sel.tpm, df_sel.num_loops)

# %%time
corr_df = pd.DataFrame()
corr_df['gene'] = data_all.gene.unique()
corr_df['r'] = corr_df.gene.apply(get_per_gene_r)

corr_df.to_csv(os.path.join(save_dir, 'corr_df.csv'))

corr_df_filt = corr_df[corr_df.r!=0]
corr_df_filt['r2'] = corr_df_filt.r.apply(lambda x: x**2)
print(corr_df_filt.r2.describe())
ax = sns.distplot(corr_df_filt.r2)
ax.set_title('distribution of r^2 of tpm vs # loops for each gene')
plt.savefig(os.path.join(save_dir, 'corr_df_filt_dist.pdf'))

# df_sel = get_scatter_gene('NFATC1',get_text=True)
# df_sel = get_scatter_gene('FOSB',get_text=True)
# df_sel = get_scatter_gene('ADAMTS13',get_text=True)
# df_sel = get_scatter_gene('ABCA1',get_text=True)
# df_sel = get_scatter_gene('PRMT1',get_text=True)
# df_sel = get_scatter_gene('SOX9',get_text=True)
# df_sel = get_scatter_gene('YY1',get_text=True)


genes = corr_df_filt[(corr_df_filt.r2<.8)&(corr_df_filt.r2>.5)].gene.sort_values().values
print(len(genes))
for x in genes:
    print(x)
#     df_sel = get_scatter_gene(x,get_text=True)



StarcoderdataJupyterScriptsDedupFiltered