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

# # 0. Map reads onto genome sequence using Galaxy
# ### Unmapped BAM files are available from the Gene Expression Omnibus under accession number GSE163258. You can map them on usegalaxy.org
# ### Mapped BAM files are available here https://doi.org/10.6084/m9.figshare.14214131

# # 1. Count reads per gene
# ### This uses the R script 

# # 2. Calculate relative expression
#
# ### This uses a MATLAB script 

import numpy as np
import pandas as pd
import os
import json
import csv

# # 3. Format organism metadata for PEA

project = 'Smutans'
os.chdir('data')

df = pd.read_csv(os.path.join('external','file_list_external_data.csv'))
FILE_GFF = os.path.join('external',df.at[0,'filename'])
print(FILE_GFF)
FILE_JSON = os.path.join('external',df.at[1,'filename'])
print(FILE_JSON)

# ## 3.1 Create Gene Name to Locus Map from GFF

FILE_OUT = os.path.join('interim',project+'_Genename_to_locusID.csv')
label = 'RefSeq\tCDS'
#label = 'Genbank\tgene'
#### choose whether RefSeq or Genbank gff file source

# +
entries = [] #list of dictionaries 
with open(FILE_GFF, 'r') as gff:
    for line in gff:
        line_split = line.strip().split(';')
        if label not in line_split[0]:
            continue    
        row = {}
        for item in line_split[1:]:
            key,value = item.split('=')
            row[key] = value
        entries.append(row)
        
## call up gene and locus tags
## write out gene and locus map to csv for record keeping
with open(FILE_OUT, "w") as fid:
    fid.write('name,locus_tag\n')
    for row in entries:
        fid.write(f'{row["Name"]},{row["locus_tag"]}\n')           

# mark genome size
GENOME_SIZE = len(entries)
print(GENOME_SIZE)
# -

# convert GFF dictionary to pandas dataframe
df_locus_IDs = pd.DataFrame.from_dict(entries)
df_locus_IDs.head()


# +
#df_locus_IDs.drop(columns=['gbkey','gene_biotype','gene'], inplace=True)

df_locus_IDs.drop(columns=['Parent','Dbxref','Note','gbkey','transl_table'], inplace=True)
df_locus_IDs.rename(columns={'locus_tag':'Locus'},inplace=True)

# -

# ## 3.2 Create KEGG Pathway to gene locus map from JSON

FILE_OUT_PATHWAYMAP = os.path.join('interim',project+'_KEGG_Pathway_to_locusID.txt')

# +
with open(FILE_JSON, 'r') as fid:
    kdf = json.load(fid)
    
with open(FILE_OUT_PATHWAYMAP, 'w') as fid:
    organism = kdf['name']
    fid.write('\t'.join(('organism','system','category','Pathway_db','Locus',
                        'gene_label')) + '\n')
    for system_dict in kdf['children']:
        system = system_dict['name']
        if 'children' not in system_dict:
            continue
        for category_dict in system_dict['children']:
            category = category_dict['name']
            if 'children' not in category_dict:
                continue
            for pathway_dict in category_dict['children']:
                pathway = pathway_dict['name']
                if 'children' not in pathway_dict:
                    continue
                for gene_dict in pathway_dict['children']:
                    gene = gene_dict['name'].split('\t')[0]
                    gene_id, gene_label = gene.split(' ', 1)
                    fid.write('\t'.join((organism, system, category, pathway, gene_id, gene_label)) + '\n')
# -

#

# Create a pandas table to genes and lowest level pathways
kegg_df = pd.read_csv(FILE_OUT_PATHWAYMAP, sep='\t')
kegg_df.drop(columns=['organism','system','category','gene_label'], inplace = True)

kegg_df.head()

# # 4. Pathway enrichment analysis

# ## 4.0 Merge MATLAB outputs & add gene names & locus_IDs

# +
FILE_OUTPUT = os.path.join('processed', 'DifferentialExpression_Complete_List.csv')
FILE_LIST_INPUT = os.path.join('interim','file_list_relative_expression.csv')

with open(FILE_LIST_INPUT, newline='') as fid:
    reader = csv.reader(fid)
    FILE_LIST = list(reader)
    
print(FILE_LIST)
# -

# merge MATLAB outputs and rename columns
all_df = None
for (header, fname) in FILE_LIST:
    df = pd.read_csv(os.path.join('interim',fname))
    df.rename(columns={'Gene':'Order'}, inplace=True)
    #df.set_index('Order', inplace=True)
    df.rename(columns={'p_fdr': f'{header}_pfdr'}, inplace=True)
    df.rename(columns={'log2_fold_change': f'{header}_l2fc'}, inplace=True)
    df.rename(columns={'pvals': f'{header}_pvals'}, inplace=True)
    if all_df is None:
        # First file
        all_df = df
    else:
        # Need to append
        all_df = all_df.merge(df, how='outer', on='Order')        

all_df

all_df.at[1958,'Order'] = "NP_720543.2" #editing ComS
all_df.rename(columns={'Order':'Name'},inplace=True)
all_df.set_index('Name',inplace=True)

all_df

df_locus_IDs.set_index('Name',inplace=True)

# Add in information for comS
df_locus_IDs.at['NP_720543.2','gene'] = 'comS'
df_locus_IDs.at['NP_720543.2','Locus'] = 'SMU_62'
df_locus_IDs.at['NP_720543.2','product'] = 'comS'
df_locus_IDs.at['NP_720543.2','protein_id'] = 'NP_720543.2'

df_locus_IDs

# add gene names to dataframe
data = df_locus_IDs.join(all_df, how='outer', on='Name')
# output the differential expression with gene names & pathway maps for record keeping
data.to_csv(FILE_OUTPUT)

data

# ## 4.3 Select significantly up-regulated and down-regulated genes for each sample

FDR = 0.05 #set false discovery rate threshold

data.set_index('Locus', inplace=True)
#data.drop(columns=['Name','Order','protein_id'], inplace=True)
data.head(5)

# +
# Check whether each gene is significantly up or down regulated

all_gene_enrichments = {}
for (header, fname) in FILE_LIST:
    file_data = pd.DataFrame(index=data.index) #data[[f'{header}_pfdr', f'{header}_l2fc']].copy()
    file_data['downregulated'] = (data[f'{header}_l2fc'] < 0) & (data[f'{header}_pfdr'] < FDR)
    file_data['upregulated'] = (data[f'{header}_l2fc'] > 0) & (data[f'{header}_pfdr'] < FDR)
    all_gene_enrichments[header] = file_data
# -

# ## 4.4 Use Kegg-annotated pathways to check for significant enrichment
#

# +
from scipy.stats import hypergeom
from statsmodels.stats.multitest import multipletests

count = 0
all_pathway_enrichments = {}
for header, gene_enrichments in all_gene_enrichments.items():
    pathway_enrichments = {}
    
    for pathway, df in kegg_df.groupby('Pathway_db'):
        gene_list = df['Locus'].unique()
        
        # Calculate number of genes in this pathway
        num_genes_total = len(gene_list)
        if num_genes_total < 5:
            continue
        #freq_background = num_genes_total / GENOME_SIZE
        
        # Downselect gene list to only include ones observed in our data
        gene_list_obs = [g for g in gene_list if g in gene_enrichments.index]
        
        # Slice enrichment dataframe to only include these relevant genes
        relevant_gene_enrichments = gene_enrichments.loc[gene_list_obs]
        
        # Look for significant upregulation/downregulation
        pathway_enrichments[pathway] = {}
        for direction in ['upregulated', 'downregulated']:
            num_genes_sig = sum(relevant_gene_enrichments[direction])
            #freq_gene_sig = num_genes_sig / sum(gene_enrichments[direction])
            
            # Upper tail of hypergeometric distribution
            pval = hypergeom.sf(
                num_genes_sig,
                GENOME_SIZE,
                num_genes_total,
                sum(gene_enrichments[direction]),
            )
            
            pathway_enrichments[pathway][direction] = {
                'pval': pval,
                'input': num_genes_sig,
                'background': num_genes_total,
                'genome_size': GENOME_SIZE,
                'num_sig_genes': sum(gene_enrichments[direction]),
            }
    #print(pathway_enrichments[pathway])
    # Correct for FDR
    for direction in ['upregulated', 'downregulated']:
        pvals = [pathway_enrichments[pathway][direction]['pval'] \
                 for pathway in pathway_enrichments]
        _, pvals_corr, _, _ = multipletests(pvals, method='fdr_bh')
        for pathway, pval_corr in zip(pathway_enrichments, pvals_corr):
            pathway_enrichments[pathway][direction]['pval_corr'] = pval_corr
    count+=1
    all_pathway_enrichments[header] = pathway_enrichments
    #print(all_pathway_enrichments)

# -

all_gene_enrichments.items()

all_pathway_enrichments

FILE_PEA_OUTPUT = os.path.join('processed','PathwayEnrichmentAnalysis_Results.csv')
for header, pathway_enrichments in all_pathway_enrichments.items():
    with open(f'{FILE_PEA_OUTPUT}_{header}.txt', 'w') as fid:
        fid.write('Pathway\t')
        for direction in ['upregulated','downregulated']:
            fid.write('\t'.join(['background', f'{direction}_input', f'{direction}_FDR','FDR<0.05','']))
        fid.write('\n')
        for pathway in pathway_enrichments.keys():
            row = f'{pathway}'
            for direction in ['upregulated','downregulated']:
                row += f'\t{pathway_enrichments[pathway][direction]["background"]}'
                row += f'\t{pathway_enrichments[pathway][direction]["input"]}'
                row += f'\t{pathway_enrichments[pathway][direction]["pval_corr"]}'
                if pathway_enrichments[pathway][direction]["pval_corr"] < FDR:
                    row += '\t*'
                else:
                    row +="\t"
                
            row += '\n'
            fid.write(row)


# ## 4.5 Comparing overall gene expression profiles of samples
#

# calculate the Pearson's correlation coefficient
correlations = data.corr('pearson')

correlations

# +
plot = data.plot.scatter(y='MUC5ACGlycanvsMedia_l2fc', x='MUC5BGlycanvsMedia_l2fc', c='DarkBlue', figsize=(4,4))
plot.set_ylabel('log2(Fold Change MUC5AC Glycans vs CMedia)')
plot.set_xlabel('log2(Fold Change MUC5B Glycans vs CMedia)')


# -

plot = data.plot.scatter(y='MUC5BMucinvsMedia_l2fc', x='MUC5BGlycanvsMedia_l2fc', c='DarkBlue', figsize=(4,4))
plot.set_ylabel('log2(Fold Change MUC5B Mucin vs CMedia)')
plot.set_xlabel('log2(Fold Change MUC5B Glycans vs CMedia)')

plot = data.plot.scatter(y='MUC5BMucinvsMedia_l2fc', x='MUC5ACGlycanvsMedia_l2fc', c='DarkBlue', figsize=(4,4))
plot.set_ylabel('log2(Fold Change MUC5B Mucin vs CMedia)')
plot.set_xlabel('log2(Fold Change MUC5AC Glycans vs CMedia)')



StarcoderdataJupyterScriptsDedupFiltered