# ---
# 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
# ---

# Design MERFISH probes using the example inputs from <NAME>. The original MATLAB design pipeline can be found at https://github.com/ZhuangLab/MERFISH_analysis.

# # Download the input data
# # This is for the UNIX-like operating systems. If you are using Windows, just download the files accordingly.
# # !mkdir temporary_data
# # !wget http://zhuang.harvard.edu/merfish/MERFISHData/MERFISH_Examples2.zip -O temporary_data/MERFISH_Examples2.zip
# # !unzip -o temporary_data/MERFISH_Examples2.zip -d temporary_data 
# # Make a path for output
# # !mkdir temporary_data/MERFISH_Examples2/outputs

# +
# Define all the input files you need in this script

import os

base_folder = r'\\10.245.74.212\Chromatin_NAS_2\Chromatin_Libraries'
rpkm_filename = os.path.join(base_folder, r'CTP-10_Aire\MERFISH_designer\Aire_smFISH\design\E-MTAB-6081.processed.3\mouse_rpkm.txt')
pool_folder = os.path.join(base_folder, 'CTP-10_Aire')


ref_folder = os.path.join(base_folder, r'CTP-10_Aire\MERFISH_designer\Aire_smFISH')
output_folder = os.path.join(ref_folder, 'outputs')
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

codebook_file = os.path.join(ref_folder, 'design', r'codebook.csv')


genome_folder = os.path.join(base_folder, r'Genomes\mouse\GRCm39_ensembl')
transcripts_fasta_file = os.path.join(genome_folder, r'Mus_musculus.GRCm39.cdna.all.fa')
ncRNA_file = os.path.join(genome_folder, r'Mus_musculus.GRCm39.ncrna.fa')
fpkm_tracking_file = os.path.join(ref_folder, r'aire_wt_mtec_genes_ensembl.fpkm_tracking')

readout_fasta_file = os.path.join(ref_folder, r'readouts.fasta')

forward_primer_file = os.path.join(base_folder, r'Primers\forward_primers.fasta')
reverse_primer_file = os.path.join(base_folder, r'Primers\reverse_primers.fasta')

# Define the output files
ottable_transcriptome_file = os.path.join(output_folder, r'ottable_transcriptome.pkl')
selected_primers_file = os.path.join(output_folder, r'selected_primers.csv')
probe_output_file = os.path.join(output_folder, r'designed_probes.csv')
transcript_level_report_file = os.path.join(output_folder, r'transcript_level_report.csv')
# -

print(transcripts_fasta_file)

# # Initialize data structures

# +
# Import the modules
import os
import sys
import pandas as pd

from IPython.display import display
import MERFISH_probe_design

import MERFISH_probe_design.IO.file_io as fio
import MERFISH_probe_design.probe_design.probe_dict as p_d
import MERFISH_probe_design.probe_design.OTTable_dict as ot
import MERFISH_probe_design.probe_design.readout_sequences as rs
import MERFISH_probe_design.probe_design.probe_selection as ps
import MERFISH_probe_design.probe_design.quality_check as qc
from MERFISH_probe_design.probe_design import filters
from MERFISH_probe_design.probe_design import plot
from MERFISH_probe_design.probe_design import primer_design
# -

# ## load transcriptome

from importlib import reload
reload(fio)
reload(qc)

# Load the transcriptome as a pandas data frame
transcriptome = fio.load_transcriptome(transcripts_fasta_file, None)


def find_gene_id_in_description(description_string):
    import re
    return re.search(r'gene:([A-Za-z0-9\.]+)', description_string)[0].split('gene:')[1]
def find_gene_name_in_description(description_string):
    import re
    gene_matches = re.search(r'gene_symbol:([A-Za-z0-9\.]+)', description_string)
    if gene_matches is None:
        print(description_string)
        return float('nan')
    else:
        return gene_matches.group(0).split('gene_symbol:')[1]
        


transcriptome['gene_id'] = list(map(find_gene_id_in_description, transcriptome['description']))
transcriptome['gene_short_name'] = list(map(find_gene_name_in_description, transcriptome['description']))

# Make sure that the transcriptome data frame has the standard column names.
# The standard columns are: transcript_id, sequence, gene_id, gene_short_name and FPKM.
# Also remove the non-standard columns for clarity.
transcriptome = qc.check_and_standardize_transcriptome(transcriptome, remove_non_standard_columns=True)
transcriptome # Let's have a look at what's inside the transcriptome

# check wether this name exist
print(transcriptome['gene_id'][transcriptome['gene_short_name']=='Krt5'])
print(transcriptome['gene_id'][transcriptome['gene_short_name']=='Aire'])

# +
# Generate a codebook table for the target transcript
# Choose one of the full length protein coding transcript as the target
bit_names = ['RS0015', 'RS0083', 'RS0095', 'RS0109']

barcode_table = pd.DataFrame({'name':['Aire', 'Krt5'], 
                              #'id':['ENSMUSG00000061527.8', 'ENSMUSG00000000731.16'], 
                              'id':['ENSMUST00000023709.7', 'ENSMUST00000145975.8'], 
                              'barcode_str':['1100','0011']})
barcode_table
# -

# Initialize the probe dictionary which is the carrier of the probes throught the design process.
probe_dict = p_d.init_probe_dict(['Aire', 'Krt5'], transcriptome, 
                                 'gene_short_name', K=30)
p_d.print_probe_dict(probe_dict)

# Select the transcripts that we want to target
# The target transcripts are already defined in the codebook
probe_dict = p_d.select_transcripts_by_ids(probe_dict, 
                                           transcript_ids=['ENSMUST00000023709.7', 'ENSMUST00000145975.8'])
p_d.print_probe_dict(probe_dict) # We excluded all the transcripts that are not our direct targets

# Initialize the off-target counting tables
# OTTable for rRNA/tRNAs
ncRNAs =  fio.load_fasta_into_df(ncRNA_file)
ottable_rtRNAs = ot.get_OTTable_for_rtRNAs(ncRNAs, 15)

# OTTables for the genes we target
gene_ottable_dict = ot.get_gene_OTTables(transcriptome, ['Aire', 'Krt5'], 'gene_short_name', 17)

# %%time
# OTTable for the transcriptome.
ottable_transcriptome = ot.get_OTTable_for_transcriptome(transcriptome, 17)

# # Select target regions

# Calculate and plot the GC contents of the target regions
filters.calc_gc_for_probe_dict(probe_dict, column_key_seq='target_sequence', column_key_write='target_GC')
plot.plot_hist(probe_dict, column_key='target_GC')

# Filter GC cotent and plot the GC content after filtering
filters.filter_probe_dict_by_metric(probe_dict, 'target_GC', lower_bound=43, upper_bound=73)
plot.plot_hist(probe_dict, column_key='target_GC')

# Calculate and plot the melting-temperatures (Tm)
filters.calc_tm_for_probe_dict(probe_dict, Na_conc=390, fmd_percentile=30, probe_conc=5,
                              column_key_seq='target_sequence', column_key_write='target_Tm')
plot.plot_hist(probe_dict, column_key='target_Tm')

# Filter by Tm
# NOTE: here we used a higher upper bound for GC content and Tm than JM's original
# cutoffs. It was shown that higher Tm gives better signal-to-noise ratios in SM-FISH
filters.filter_probe_dict_by_metric(probe_dict, 'target_Tm', lower_bound=52)
plot.plot_hist(probe_dict, column_key='target_Tm')

# Calculate and plot the off-targets to rRNA/tRNAs
ot.calc_OTs(probe_dict, ottable_rtRNAs, 'target_sequence', 'target_OT_rtRNA', 15)
plot.plot_hist(probe_dict, 'target_OT_rtRNA', y_max=1000)
# Filter out probes that have any rRNA/tRNA off-targets
filters.filter_probe_dict_by_metric(probe_dict, 'target_OT_rtRNA', upper_bound=0.5)
plot.plot_hist(probe_dict, 'target_OT_rtRNA')

# +
# Get the FPKMs of the transcripts
transcript_fpkms = dict(zip(list(transcriptome['transcript_id']), list(transcriptome['FPKM'])))

# Calculate the specificities and isoform specificities of the target regions
ot.calc_specificity(probe_dict, ottable_transcriptome, gene_ottable_dict, transcript_fpkms,
                    'target_sequence', 'target_specificity', 'target_isospecificity', 17)

plot.plot_hist(probe_dict, 'target_specificity')
plot.plot_hist(probe_dict, 'target_isospecificity')

# +
# Filter the specificities of the target regions
# Here we allow the probes to target other isoforms
# Don't allow any gene-level off-targets
filters.filter_probe_dict_by_metric(probe_dict, 'target_specificity', lower_bound=0.99999)
#filters.filter_probe_dict_by_metric(probe_dict, 'target_isospecificity', upper_bound=0.1)

plot.plot_hist(probe_dict, 'target_specificity')
plot.plot_hist(probe_dict, 'target_isospecificity')
# -

# # Design readout sequences

# Load the readout sequences into a data frame
readout_seqs = fio.load_fasta_into_df(readout_fasta_file)
rs.append_on_bit_ids_to_readout_sequences(readout_seqs, bit_names)
readout_seqs

# +
import importlib
importlib.reload(rs)

# Add the readout sequences. Here we randomly add 2 readout sequences to each probe.
# Add an "T" between the concatenated sequences.
# Force each probe to have the same readout sequence
rs.add_readout_seqs_to_probes_random(probe_dict, readout_seqs, barcode_table, 2, 
                                     spacer='A', gene_id_key='name', each_probe_1_on_bit=True)
# -

probe_dict['Aire'].keys()

# probe_dict['Krt5']['ENSMUST00000023709.7']

# Filter out probes that have off-targets to rRNA/tRNAs
ot.calc_OTs(probe_dict, ottable_rtRNAs, 'target_readout_sequence', 'target_readout_OT_rtRNA', 15)
plot.plot_hist(probe_dict, 'target_readout_OT_rtRNA', y_max=1000)
filters.filter_probe_dict_by_metric(probe_dict, 'target_readout_OT_rtRNA', upper_bound=0.5)
plot.plot_hist(probe_dict, 'target_readout_OT_rtRNA')

# Calculate how many more off-targets to the transcriptome are introduced due to the readout sequences.
# The off-target counts are weighted down by the FPKMs of the on-target transcripts
ot.calc_OT_diffs(probe_dict, ottable_transcriptome, gene_ottable_dict, transcript_fpkms, 
                 'target_sequence', 'target_readout_sequence', 'readout_OT_increase', 17)
plot.plot_hist(probe_dict, 'readout_OT_increase', y_max=2000)

import numpy as np
np.unique(probe_dict['Aire']['ENSMUST00000145975.8']['probe_barcode'], return_counts=True)

# Filter out the probes with extra off-targets due to the readouts
# Require the new weighted off-targets to be minor compared to the on-target weight.
filters.filter_probe_dict_by_metric(probe_dict, 'readout_OT_increase', upper_bound=0.5)
plot.plot_hist(probe_dict, 'readout_OT_increase')

np.unique(probe_dict['Aire']['ENSMUST00000145975.8']['probe_barcode'], return_counts=True)

# # Select probes

# %%time
# Select probes by a stochastic greedy algorithms that optimizes the on-bit coverage
# and minimizes the overlapping between probes.
ps.select_probes_greedy_stochastic(probe_dict, N_probes_per_transcript=48, 
                                   N_on_bits=2)

# Let's plot the probe coverage of an example transcript
seq_len = len(transcriptome[transcriptome['transcript_id'] == 'ENSMUST00000145975.8'].iloc[0]['sequence'])
plot.plot_sequence_coverage(probe_dict['Aire']['ENSMUST00000145975.8'], seq_len)

# Let's plot the probe coverage of an example transcript
seq_len = len(transcriptome[transcriptome['transcript_id'] == 'ENSMUST00000023709.7'].iloc[0]['sequence'])
plot.plot_sequence_coverage(probe_dict['Krt5']['ENSMUST00000023709.7'], seq_len)

probe_dict['Aire']['ENSMUST00000145975.8']

probe_dict['Krt5']['ENSMUST00000023709.7']

# # Quality check

# +
# %%time
# Filter out the probes that self complement or complement with other probes.

# Iterately remove the probes with high numbers of cis/trans-complementarity
# This filtering strategy is a compromise between speed and the number of probes to keep
while True:
    # Make a OTTable from the reverse-complement sequences of the probes.
    ottable_probes_rc = ot.get_OTTable_for_probe_dictionary(probe_dict, 'target_readout_sequence', 15, rc=True)
    
    # The off-targets in this table indicates cis/trans-complementarity
    ot.calc_OTs(probe_dict, ottable_probes_rc, 'target_readout_sequence', 'probe_cis_trans_OT', 15)
    max_ot = max(plot.get_values_from_probe_dict(probe_dict, 'probe_cis_trans_OT'))
    if max_ot == 0:
        break
    
    # Remove probes that have any cis/trans-complementarity
    filters.filter_probe_dict_by_metric(probe_dict, 'probe_cis_trans_OT', upper_bound=max_ot - 0.5)
    
plot.plot_hist(probe_dict, 'probe_cis_trans_OT')
# -



# +
# Get the reverse-complementary sequences of the designed probes
p_d.get_rc_sequences(probe_dict, 'target_readout_sequence', 'target_readout_sequence_rc')

# Get a data frame of all probes
df_all_probes = p_d.probe_dict_to_df(probe_dict).sort_values(['probe_barcode', 'shift'])

# Make a column of probe names
probe_names = np.array(df_all_probes['gene_id']) + '_' + np.array(df_all_probes['transcript_id'])
probe_names = probe_names + '_shift:' + np.array(df_all_probes['shift']).astype(str)
probe_names = probe_names + '_GC:' + np.array(df_all_probes['target_GC']).astype('|S4').astype(str)
probe_names = probe_names + '_Tm:' + np.array(df_all_probes['target_Tm']).astype('|S4').astype(str)
probe_names = probe_names + '_specificity:' + np.array(df_all_probes['target_specificity']).astype('|S1').astype(str)
probe_names = probe_names + '_isospecificity:' + np.array(df_all_probes['target_isospecificity']).astype('|S4').astype(str)
probe_names = probe_names + '_' + np.array(df_all_probes['readout_names'])
df_all_probes['probe_names'] = probe_names

# Write the designed probes
df_all_probes.to_csv(probe_output_file, index=False)
# -

# generate IDT list
well_positions = []
for _col in 'ABCDEFGH':
    for _row in range(12):
        well_positions.append(f"{_col}{_row+1}")

# +
# export IDT format
# Get the reverse-complementary sequences of the designed probes
p_d.get_rc_sequences(probe_dict, 'target_readout_sequence', 'target_readout_sequence_rc')
# Get a data frame of all probes
idt_plate_df = pd.DataFrame({
    'Well Position': well_positions,
    'Name': df_all_probes['probe_names'][:96],
    'Sequence': [_seq for _seq in df_all_probes['target_readout_sequence_rc'][:96]],
})
# Write the designed probes
idt_probe_output_file = os.path.join(output_folder, r'idt_order_new.xlsx')

idt_plate_df.to_excel(idt_probe_output_file, index=False)
# -

print(output_folder)

idt_df



StarcoderdataJupyterScriptsDedupFiltered