# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_gene_report.ipynb.

# %% auto 0
__all__ = ['create_iso_adata', 'plot_switch_gen_bar', 'plot_isoforms_frequencies', 'plot_genes_cell_type',
           'plot_isoforms_summary', 'filter_mdata']

# %% ../nbs/01_gene_report.ipynb 2
import muon as mu
import numpy as np
import pandas as pd
import scanpy as sc
from scipy.sparse import csr_matrix
import anndata as ad
from matplotlib import pyplot as plt
import seaborn as sns
import patchworklib as pw
print(ad.__version__)

# %% ../nbs/01_gene_report.ipynb 6
# creates isoform level anndata object from df produced by Sicelore:
def create_iso_adata (path):
    with open(path) as isoforms:
        df = pd.read_table(isoforms, sep='\t')

    counts = df.iloc[0:,2:]
    adata_iso = ad.AnnData(counts).transpose()
    adata_iso.var_names = df['transcriptId'].to_list()
    adata_iso.var = df[['transcriptId', 'geneId']]
    barcodes = {'barcodes': df.columns.values[2:]}   
    barcodes = pd.DataFrame(data=barcodes)
    adata_iso.obs = barcodes
    adata_iso.obs_names = barcodes['barcodes'].tolist()
    return(adata_iso)

# %% ../nbs/01_gene_report.ipynb 10
def _plot_switch_gen_bar (mdata, _ax):
    if _ax is None:
        fig, ax = plt.subplots()
        mdata.mod['gen'].var['n_of_iso'] = mdata.mod['iso'].var.groupby(by='geneId').count()
        x = ['genes']
        x1 = ['transcripts']
        multiple_iso = sum(mdata.mod['gen'].var['n_of_iso'] > 1)
        mono_iso = sum((mdata.mod['gen'].var['n_of_iso'] > 1) == False)
        labels = [str(round(1000*multiple_iso/(multiple_iso+mono_iso))/10) + '%', str(round(1000*mono_iso/(multiple_iso+mono_iso))/10) + '%']
        mult = ax.bar(x, multiple_iso, color = 'deepskyblue', label=labels[1]) #multiple_iso/(multiple_iso+mono_iso))
        mono = ax.bar(x, mono_iso, bottom=multiple_iso, color='sandybrown', label=labels[0]) #mona_iso/(multiple_iso+mono_iso))
        tran = ax.bar(x1, len(mdata.mod['iso'].var['transcriptId']))
        ax.text(
            ax.patches[0].get_x() + ax.patches[0].get_width() / 2, ax.patches[0].get_height() / 2, labels[0], ha="center", va="center"
        )
        ax.text(
            ax.patches[1].get_x() + ax.patches[1].get_width() / 2, ax.patches[1].get_height() / 2 + ax.patches[0].get_height(), labels[1], ha="center", va="center"
        )
        plt.legend(['Multiple isoforms', 'Single isoform'])
        plt.show()
    else:
        mdata.mod['gen'].var['n_of_iso'] = mdata.mod['iso'].var.groupby(by='geneId').count()
        x = ['genes']
        x1 = ['transcripts']
        multiple_iso = sum(mdata.mod['gen'].var['n_of_iso'] > 1)
        mono_iso = sum((mdata.mod['gen'].var['n_of_iso'] > 1) == False)
        labels = [str(round(1000*multiple_iso/(multiple_iso+mono_iso))/10) + '%', str(round(1000*mono_iso/(multiple_iso+mono_iso))/10) + '%']
        mult = _ax.bar(x, multiple_iso, color = 'deepskyblue', label=labels[1]) #multiple_iso/(multiple_iso+mono_iso))
        mono = _ax.bar(x, mono_iso, bottom=multiple_iso, color='sandybrown', label=labels[0]) #mona_iso/(multiple_iso+mono_iso))
        tran = _ax.bar(x1, len(mdata.mod['iso'].var['transcriptId']))
        _ax.text(
            _ax.patches[0].get_x() + _ax.patches[0].get_width() / 2, _ax.patches[0].get_height() / 2, labels[0], ha="center", va="center"
        )
        _ax.text(
            _ax.patches[1].get_x() + _ax.patches[1].get_width() / 2, _ax.patches[1].get_height() / 2 + _ax.patches[0].get_height(), labels[1], ha="center", va="center"
        )
        _ax.legend(['Multiple isoforms', 'Single isoform'])
def plot_switch_gen_bar (mdata):
    _plot_switch_gen_bar (mdata, None)

    

# %% ../nbs/01_gene_report.ipynb 12
# count total frequencis of genes having 1 , 2 , 3 , ... n isoforms
def _plot_isoforms_frequencies (mdata, _ax):
    if _ax is None:
        fig, ax = plt.subplots()
        mdata.mod['gen'].var['n_of_iso'].value_counts().plot(ax=ax, kind='bar', 
                                                             xlabel='number of isoforms per gene', ylabel='quantity of genes')
    else:
        mdata.mod['gen'].var['n_of_iso'].value_counts().plot(ax=_ax, kind='bar', 
                                                             xlabel='number of isoforms per gene', ylabel='quantity of genes')
def plot_isoforms_frequencies(mdata):
    _plot_isoforms_frequencies(mdata, None)

# %% ../nbs/01_gene_report.ipynb 17
# boxplot with number of genes expressed per cell type:
def _plot_genes_cell_type (mdata, _ax): #mdata with annotated cell types in isoform and gen anndata objects
    if _ax is None:
        df = pd.DataFrame(np.transpose(mdata.mod['iso'].X), columns=mdata.mod['iso'].obs['cell_type'])
        df = df.sum(axis = 0).to_frame().reset_index()
        df.columns = ['cell_type', 'n_of_genes']
        ax = sns.boxplot(x='cell_type', y='n_of_genes', data=df) 
        ax1 = sns.stripplot(x='cell_type', y='n_of_genes', data=df, color = 'black', size = 3)
        ax.set_xticklabels(ax.get_xticklabels(), rotation=90)
        plt.show()
    else:
        df = pd.DataFrame(np.transpose(mdata.mod['iso'].X), columns=mdata.mod['iso'].obs['cell_type'])
        df = df.sum(axis = 0).to_frame().reset_index()
        df.columns = ['cell_type', 'n_of_genes']
        sns.boxplot(x='cell_type', y='n_of_genes', data=df, ax=_ax) 
        sns.stripplot(x='cell_type', y='n_of_genes', data=df, color = 'black', size = 3, ax=_ax)
        _ax.set_xticklabels(_ax.get_xticklabels(), rotation=90)

def plot_genes_cell_type(mdata):
    _plot_genes_cell_type(mdata, None)

# %% ../nbs/01_gene_report.ipynb 19
def plot_isoforms_summary(mdata):
    ax1 = pw.Brick(figsize=(4,4))
    _plot_switch_gen_bar(mdata, ax1)
    ax1.set_title("Multiple isoforms genes %")
    ax2 = pw.Brick(figsize=(4,4))
    _plot_isoforms_frequencies(mdata, ax2)
    ax2.set_title("Frequency of isoforms per gene")
    ax3 = pw.Brick(figsize=(3,2))
    _plot_genes_cell_type(mdata, ax3)
    ax3.set_title("Nb of genes per cell type")
    return (ax1|ax2|ax3).savefig()
plot_isoforms_summary(mdata)

# %% ../nbs/01_gene_report.ipynb 24
# subset mdata object for switching genes
def filter_mdata (mdata):
    gene_m = mdata.mod['gen'][:, mdata.mod['gen'].var['n_of_iso'] > 1]
    gene_iso_m = mdata.mod['iso'][:, mdata.mod['iso'].var['geneId'].isin(gene_m.var.index)]
    mdata_m = mu.MuData({'gen': gene_m, 'iso': gene_iso_m})
    return mdata_m
