# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:aparent]
#     language: python
#     name: conda-env-aparent-py
# ---

# +
import pandas as pd
import scipy
import numpy as np
import scipy.sparse as sp
import scipy.io as spio

import pickle
import os


# +
data = pd.read_csv('unprocessed_data/Alt_5SS_Tag_to_Seq_Map.csv',sep=',',index_col=0)
c = spio.loadmat('unprocessed_data/Alt_5SS_Usage_All_Cells.mat')

c_MCF7 = sp.csc_matrix(c['MCF7'])
c_CHO = sp.csc_matrix(c['CHO'])
c_HELA = sp.csc_matrix(c['HELA'])
c_HEK = sp.csc_matrix(c['HEK'])

# +
#Sort data on counts

total_c_MCF7 = np.ravel(c_MCF7.sum(axis=-1))
total_c_CHO = np.ravel(c_CHO.sum(axis=-1))
total_c_HELA = np.ravel(c_HELA.sum(axis=-1))
total_c_HEK = np.ravel(c_HEK.sum(axis=-1))

avg_c = (total_c_HEK + total_c_HELA + total_c_CHO + total_c_MCF7) / 4.0

sort_index = np.argsort(avg_c)

data = data.iloc[sort_index].copy().reset_index(drop=True)
c_MCF7 = c_MCF7[sort_index, :]
c_CHO = c_CHO[sort_index, :]
c_HELA = c_HELA[sort_index, :]
c_HEK = c_HEK[sort_index, :]


# +
#Constant background sequence context
up_background = 'gggcatcgacttcaaggaggacggcaacatcctggggcacaagctggagtacaactacaacagccacaacgtctatatcatggccgacaagcagaagaacggcatcaaagtgaacttcaagatccgccacaacatcgagg'.upper()
dn_background = 'acagagtttccttatttgtctctgttgccggcttatatggacaagcatatcacagccatttatcggagcgcctccgtacacgctattatcggacgcctcgcgagatcaatacgtatacca'.upper()

print('len(up_background) = ' + str(len(up_background)))
print('len(dn_background) = ' + str(len(dn_background)))


# +
#Extend sequences and count matrices

data['Padded_Seq'] = up_background + data['Seq'].str.slice(0,101) + dn_background

padded_c_MCF7, padded_c_CHO, padded_c_HELA, padded_c_HEK = [
    sp.csr_matrix(
        sp.hstack([
            sp.csc_matrix((c_mat.shape[0], len(up_background))),
            c_mat[:, :101],
            sp.csc_matrix((c_mat.shape[0], len(dn_background))),
            sp.csc_matrix(np.array(c_mat[:, 303].todense()).reshape(-1, 1))
        ])
    )
    for c_mat in [c_MCF7, c_CHO, c_HELA, c_HEK]
]

print('padded_c_MCF7.shape = ' + str(padded_c_MCF7.shape))
print('padded_c_CHO.shape = ' + str(padded_c_CHO.shape))
print('padded_c_HELA.shape = ' + str(padded_c_HELA.shape))
print('padded_c_HEK.shape = ' + str(padded_c_HEK.shape))


# +
#Filter each dataset on > 0 count

hek_keep_index = np.nonzero(np.ravel(padded_c_HEK.sum(axis=-1)) > 0)[0]
hela_keep_index = np.nonzero(np.ravel(padded_c_HELA.sum(axis=-1)) > 0)[0]
mcf7_keep_index = np.nonzero(np.ravel(padded_c_MCF7.sum(axis=-1)) > 0)[0]
cho_keep_index = np.nonzero(np.ravel(padded_c_CHO.sum(axis=-1)) > 0)[0]

#HEK data
data_hek_filtered = data.iloc[hek_keep_index].copy().reset_index(drop=True)
c_hek_filtered = padded_c_HEK[hek_keep_index, :]

#HELA data
data_hela_filtered = data.iloc[hela_keep_index].copy().reset_index(drop=True)
c_hela_filtered = padded_c_HELA[hela_keep_index, :]

#MCF7 data
data_mcf7_filtered = data.iloc[mcf7_keep_index].copy().reset_index(drop=True)
c_mcf7_filtered = padded_c_MCF7[mcf7_keep_index, :]

#CHO data
data_cho_filtered = data.iloc[cho_keep_index].copy().reset_index(drop=True)
c_cho_filtered = padded_c_CHO[cho_keep_index, :]

print('len(data_hek_filtered) = ' + str(len(data_hek_filtered)))
print('c_hek_filtered.shape = ' + str(c_hek_filtered.shape))

print('len(data_hela_filtered) = ' + str(len(data_hela_filtered)))
print('c_hela_filtered.shape = ' + str(c_hela_filtered.shape))

print('len(data_mcf7_filtered) = ' + str(len(data_mcf7_filtered)))
print('c_mcf7_filtered.shape = ' + str(c_mcf7_filtered.shape))

print('len(data_cho_filtered) = ' + str(len(data_cho_filtered)))
print('c_cho_filtered.shape = ' + str(c_cho_filtered.shape))



# +
#Get joined min dataset

min_keep_index = (np.ravel(padded_c_HEK.sum(axis=-1)) > 0)
min_keep_index = min_keep_index & (np.ravel(padded_c_HELA.sum(axis=-1)) > 0)
min_keep_index = min_keep_index & (np.ravel(padded_c_MCF7.sum(axis=-1)) > 0)
min_keep_index = min_keep_index & (np.ravel(padded_c_CHO.sum(axis=-1)) > 0)

#MIN data
data_min_filtered = data.iloc[min_keep_index].copy().reset_index(drop=True)

c_hek_min_filtered = padded_c_HEK[min_keep_index, :]
c_hela_min_filtered = padded_c_HELA[min_keep_index, :]
c_mcf7_min_filtered = padded_c_MCF7[min_keep_index, :]
c_cho_min_filtered = padded_c_CHO[min_keep_index, :]

print('len(data_min_filtered) = ' + str(len(data_min_filtered)))

print('c_hek_min_filtered.shape = ' + str(c_hek_min_filtered.shape))
print('c_hela_min_filtered.shape = ' + str(c_hela_min_filtered.shape))
print('c_mcf7_min_filtered.shape = ' + str(c_mcf7_min_filtered.shape))
print('c_cho_min_filtered.shape = ' + str(c_cho_min_filtered.shape))


# +
#Pickle final datasets

data_min_filtered = data_min_filtered.rename(columns={'Padded_Seq' : 'padded_seq'})
data_hek_filtered = data_hek_filtered.rename(columns={'Padded_Seq' : 'padded_seq'})
data_hela_filtered = data_hela_filtered.rename(columns={'Padded_Seq' : 'padded_seq'})
data_mcf7_filtered = data_mcf7_filtered.rename(columns={'Padded_Seq' : 'padded_seq'})
data_cho_filtered = data_cho_filtered.rename(columns={'Padded_Seq' : 'padded_seq'})

data_min_filtered = data_min_filtered[['padded_seq']]
data_hek_filtered = data_hek_filtered[['padded_seq']]
data_hela_filtered = data_hela_filtered[['padded_seq']]
data_mcf7_filtered = data_mcf7_filtered[['padded_seq']]
data_cho_filtered = data_cho_filtered[['padded_seq']]

splicing_5ss_dict = {
    'min_df' : data_min_filtered,
    'hek_df' : data_hek_filtered,
    'hela_df' : data_hela_filtered,
    'mcf7_df' : data_mcf7_filtered,
    'cho_df' : data_cho_filtered,
    
    'hek_count' : c_hek_filtered,
    'hela_count' : c_hela_filtered,
    'mcf7_count' : c_mcf7_filtered,
    'cho_count' : c_cho_filtered,
    
    'min_hek_count' : c_hek_min_filtered,
    'min_hela_count' : c_hela_min_filtered,
    'min_mcf7_count' : c_mcf7_min_filtered,
    'min_cho_count' : c_cho_min_filtered,
}

pickle.dump(splicing_5ss_dict, open('alt_5ss_data.pickle', 'wb'))


# +
#Align and consolidate a5ss data

plasmid_dict = pickle.load(open('alt_5ss_data.pickle', 'rb'))
plasmid_df = plasmid_dict['min_df']
hek_cuts = np.array(plasmid_dict['min_hek_count'].todense())
hela_cuts = np.array(plasmid_dict['min_hela_count'].todense())
mcf7_cuts = np.array(plasmid_dict['min_mcf7_count'].todense())
cho_cuts = np.array(plasmid_dict['min_cho_count'].todense())

total_cuts = hek_cuts + hela_cuts + mcf7_cuts + cho_cuts
total_cuts = total_cuts[:, :-1]

# +
fixed_poses = [140, 140 + 44, 140 + 79]

sd_window = 130#120
sd1_pos = 140

negative_sampling_ratio = 2

fixed_pos_mask = np.ones(total_cuts.shape[1])
for j in range(len(fixed_poses)) :
    fixed_pos_mask[fixed_poses[j]] = 0

cut_pos = np.arange(total_cuts.shape[1])

aligned_seqs = []
aligned_libs = []
aligned_mode = []

max_data_len = 3000000

aligned_hek_cuts = sp.lil_matrix((max_data_len, 2 * sd_window + 1))
aligned_hela_cuts = sp.lil_matrix((max_data_len, 2 * sd_window + 1))
aligned_mcf7_cuts = sp.lil_matrix((max_data_len, 2 * sd_window + 1))
aligned_cho_cuts = sp.lil_matrix((max_data_len, 2 * sd_window + 1))

splice_mats = [
    [hek_cuts, aligned_hek_cuts],
    [hela_cuts, aligned_hela_cuts],
    [mcf7_cuts, aligned_mcf7_cuts],
    [cho_cuts, aligned_cho_cuts]
]

old_i = 0
new_i = 0
for _, row in plasmid_df.iterrows() :
    if old_i % 10000 == 0 :
        print("Processing sequence " + str(old_i) + "...")
    
    seq = row['padded_seq']
    
    nonzero_cuts = np.nonzero( ((total_cuts[old_i, :] > 0) & (fixed_pos_mask == 1)) & ((cut_pos >= sd_window) & (cut_pos < total_cuts.shape[1] - sd_window)) )[0].tolist()
    zero_cuts = np.nonzero( ((total_cuts[old_i, :] == 0) & (fixed_pos_mask == 1)) & ((cut_pos >= sd_window + 1) & (cut_pos < total_cuts.shape[1] - sd_window - 1)) )[0].tolist()
    
    #Emit fixed splice positions
    for fixed_pos in fixed_poses :
        aligned_seqs.append(seq[fixed_pos - sd_window: fixed_pos + sd_window])
        aligned_libs.append(fixed_pos - sd1_pos)
        aligned_mode.append("fixed_" + str(fixed_pos - sd1_pos))
        
        for [cuts, aligned_cuts] in splice_mats :
            aligned_cuts[new_i, :2 * sd_window] = cuts[old_i, fixed_pos - sd_window: fixed_pos + sd_window]
            aligned_cuts[new_i, 2 * sd_window] = cuts[old_i, -1]
        
        new_i += 1
    
    #Emit denovo splice positions
    for denovo_pos in nonzero_cuts :
        aligned_seqs.append(seq[denovo_pos - sd_window: denovo_pos + sd_window])
        aligned_libs.append(denovo_pos - sd1_pos)
        aligned_mode.append("denovo_pos_" + str(denovo_pos - sd1_pos))
        
        for [cuts, aligned_cuts] in splice_mats :
            aligned_cuts[new_i, :2 * sd_window] = cuts[old_i, denovo_pos - sd_window: denovo_pos + sd_window]
            aligned_cuts[new_i, 2 * sd_window] = cuts[old_i, -1]
        
        new_i += 1
    
    if negative_sampling_ratio > 0.0 :
        n_neg = int(negative_sampling_ratio * (3 + len(nonzero_cuts)))
        sampled_zero_cuts = np.random.choice(zero_cuts, size=n_neg, replace=False)

        #Emit negative denovo splice positions
        for denovo_pos in sampled_zero_cuts :
            aligned_seqs.append(seq[denovo_pos - sd_window: denovo_pos + sd_window])
            aligned_libs.append(denovo_pos - sd1_pos)
            aligned_mode.append("denovo_neg_" + str(denovo_pos - sd1_pos))

            for [cuts, aligned_cuts] in splice_mats :
                aligned_cuts[new_i, :2 * sd_window] = cuts[old_i, denovo_pos - sd_window: denovo_pos + sd_window]
                aligned_cuts[new_i, 2 * sd_window] = cuts[old_i, -1]

            new_i += 1
    
    old_i += 1

aligned_min_hek_cuts = sp.csr_matrix(aligned_hek_cuts[:len(aligned_seqs), :])
aligned_min_hela_cuts = sp.csr_matrix(aligned_hela_cuts[:len(aligned_seqs), :])
aligned_min_mcf7_cuts = sp.csr_matrix(aligned_mcf7_cuts[:len(aligned_seqs), :])
aligned_min_cho_cuts = sp.csr_matrix(aligned_cho_cuts[:len(aligned_seqs), :])

aligned_min_df = pd.DataFrame({
    'seq'  : aligned_seqs,
    'library'  : aligned_libs,
    'origin' : aligned_mode
})

aligned_min_df = aligned_min_df[['seq', 'library', 'origin']]

print("len(aligned_min_df) = " + str(len(aligned_min_df)))

print("aligned_min_hek_cuts.shape = " + str(aligned_min_hek_cuts.shape))
print("aligned_min_hela_cuts.shape = " + str(aligned_min_hela_cuts.shape))
print("aligned_min_mcf7_cuts.shape = " + str(aligned_min_mcf7_cuts.shape))
print("aligned_min_cho_cuts.shape = " + str(aligned_min_cho_cuts.shape))


# +
#Filter out zeros

keep_index = (np.ravel(aligned_min_hek_cuts.sum(axis=-1)) > 0)
keep_index = keep_index & (np.ravel(aligned_min_hela_cuts.sum(axis=-1)) > 0)
keep_index = keep_index & (np.ravel(aligned_min_mcf7_cuts.sum(axis=-1)) > 0)
keep_index = keep_index & (np.ravel(aligned_min_cho_cuts.sum(axis=-1)) > 0)

aligned_min_df = aligned_min_df.iloc[keep_index].copy().reset_index(drop=True)

aligned_min_hek_cuts = aligned_min_hek_cuts[keep_index, :]
aligned_min_hela_cuts = aligned_min_hela_cuts[keep_index, :]
aligned_min_mcf7_cuts = aligned_min_mcf7_cuts[keep_index, :]
aligned_min_cho_cuts = aligned_min_cho_cuts[keep_index, :]

print("len(aligned_min_df) = " + str(len(aligned_min_df)))

print("aligned_min_hek_cuts.shape = " + str(aligned_min_hek_cuts.shape))
print("aligned_min_hela_cuts.shape = " + str(aligned_min_hela_cuts.shape))
print("aligned_min_mcf7_cuts.shape = " + str(aligned_min_mcf7_cuts.shape))
print("aligned_min_cho_cuts.shape = " + str(aligned_min_cho_cuts.shape))


# +
data_version = '_neg_rate_2'#'_neg_rate_1'#''

aligned_5ss_dict = {
    'min_df' : aligned_min_df,
    
    'min_hek_count' : aligned_min_hek_cuts,
    'min_hela_count' : aligned_min_hela_cuts,
    'min_mcf7_count' : aligned_min_mcf7_cuts,
    'min_cho_count' : aligned_min_cho_cuts,
}

pickle.dump(aligned_5ss_dict, open('alt_5ss_data_aligned' + data_version + '.pickle', 'wb'))
# -



StarcoderdataJupyterScriptsDedupFiltered