"""Vizualise switching genes and the differences in isoforms expression between different cell types"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_gene_report_iso_data.ipynb.

# %% auto 0
__all__ = ['download_test_data', 'create_iso_adata', 'plot_switch_gen_bar', 'plot_isoforms_frequencies', 'plot_genes_cell_type',
           'plot_isoforms_summary', 'filter_switch_genes', 'AnnDataIso']

# %% ../nbs/02_gene_report_iso_data.ipynb 3
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 requests
import sys
import patchworklib as pw
import urllib 
import os
import gzip
import shutil

# %% ../nbs/02_gene_report_iso_data.ipynb 4
def download_test_data() -> str: #The absolute path of the extracted file 'sample_isomatrix.txt' if the download is successful.
    """
    This function downloads a test data file from a specified URL, saves it locally, and extracts it.
    """

    # URL of the file to be downloaded
    url = "https://ftp.ncbi.nlm.nih.gov/geo/samples/GSM3748nnn/GSM3748087/suppl/GSM3748087%5F190c.isoforms.matrix.txt.gz"

    # Download the file from `url` and save it locally under `file.txt.gz`:
    urllib.request.urlretrieve(url, 'file.txt.gz')

    # Check if the file is downloaded correctly
    if os.path.exists('file.txt.gz'):
        print("File downloaded successfully")
        # Now we need to extract the file
        with gzip.open('file.txt.gz', 'rb') as f_in:
            with open('sample_isomatrix.txt', 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        print("File extracted successfully")
        return os.path.abspath('sample_isomatrix.txt')
    else:
        print("Failed to download the file")
        return None 

# %% ../nbs/02_gene_report_iso_data.ipynb 7
def create_iso_adata(
    path: str  # Path to a tab-separated file, typically the Sicelore output with transcript counts.
):
    """
    Creates an AnnData object containing isoform counts from a Sicelore output file.

    Parameters:
    - path (str): Path to the input file, which is expected to be tab-separated and contain:
        - `transcriptId`: Transcript identifiers.
        - `geneId`: Gene identifiers.
        - Transcript counts for individual barcodes (cells) as additional columns.

    Returns:
    - adata_iso (AnnData): An AnnData object with isoform counts.
    """

    # Open the input file and read it into a Pandas DataFrame
    with open(path) as isoforms:
        df = pd.read_table(isoforms, sep='\t')

    # Extract the count matrix
    # - `df.iloc[0:, 2:]`: assuming these are counts.
    counts = df.iloc[0:, 2:]

    # Create an AnnData object with the count matrix
    
    adata_iso = ad.AnnData(counts).transpose()

    # Assign transcript IDs as variable names (columns in the original file)
    adata_iso.var_names = df['transcriptId'].to_list()

    # Add additional variable-level metadata
    # - Include both `transcriptId` and `geneId` as annotations for the variables.
    adata_iso.var = df[['transcriptId', 'geneId']]

    # Create observation-level metadata for barcodes (cells)
    # - Extract column names (barcodes) starting from the 3rd column (index 2).
    barcodes = {'barcodes': df.columns.values[2:]}
    barcodes = pd.DataFrame(data=barcodes)

    # Assign the barcodes DataFrame to the AnnData object as observation metadata
    adata_iso.obs = barcodes

    # Set observation names (row indices) in the AnnData object to match the barcodes
    adata_iso.obs_names = barcodes['barcodes'].tolist()

    # Return the AnnData object containing isoform counts
    return adata_iso

# %% ../nbs/02_gene_report_iso_data.ipynb 12
def _plot_switch_gen_bar(adata_iso, _ax):
    """
    Creates a bar chart to visualize the proportion of genes with multiple isoforms versus single isoforms,
    and the total number of transcripts.

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.
        - `adata_iso.var`: Should include columns for 'transcriptId' and 'geneId'.
    - _ax (matplotlib.axes.Axes or None): The axes on which to render the plot.
        - If `None`, a new figure and axes are created.
    """

    if _ax is None:
        # Step 1: Calculate the number of isoforms per gene
        # - Group by 'geneId' and count the number of 'transcriptId' entries for each gene.
        iso_per_gene = adata_iso.var.reset_index().groupby(by='geneId').count()

        # Step 2: Define categories for the bar plot
        x = ['genes']  # Category for genes
        x1 = ['transcripts']  # Category for transcripts

        # Step 3: Compute the number of genes with multiple isoforms and single isoforms
        multiple_iso = sum(iso_per_gene['transcriptId'] > 1)  # Genes with more than one isoform
        mono_iso = sum((iso_per_gene['transcriptId'] > 1) == False)  # Genes with a single isoform

        # Step 4: Generate percentage labels for the plot
        total_genes = multiple_iso + mono_iso
        labels = [
            str(round(1000 * multiple_iso / total_genes) / 10) + '%',  # Percentage of multiple-isoform genes
            str(round(1000 * mono_iso / total_genes) / 10) + '%'       # Percentage of single-isoform genes
        ]

        # Step 5: Create a new figure and axes for the plot
        fig, ax = plt.subplots()

        # Step 6: Plot the bar for multiple-isoform genes
        mult = ax.bar(x, multiple_iso, color='deepskyblue', label=labels[1])

        # Step 7: Plot the bar for single-isoform genes, stacked on top of the multiple-isoform bar
        mono = ax.bar(x, mono_iso, bottom=multiple_iso, color='sandybrown', label=labels[0])

        # Step 8: Plot the total number of transcripts in a separate category
        tran = ax.bar(x1, len(adata_iso.var['transcriptId']))

        # Step 9: Add percentage labels to the bars
        ax.text(
            ax.patches[0].get_x() + ax.patches[0].get_width() / 2,  # Center of the first bar
            ax.patches[0].get_height() / 2,  # Middle of the bar height
            labels[0], ha="center", va="center"  # Align text at the center
        )
        ax.text(
            ax.patches[1].get_x() + ax.patches[1].get_width() / 2,  # Center of the second bar
            ax.patches[1].get_height() / 2 + ax.patches[0].get_height(),  # Adjust for stacked bars
            labels[1], ha="center", va="center"
        )

        # Step 10: Add a legend and display the plot
        plt.legend(['Multiple isoforms', 'Single isoform'])
        plt.show()

    else:
        # If an axes object (`_ax`) is provided, repeat the same steps on the given axes

        iso_per_gene = adata_iso.var.reset_index().groupby(by='geneId').count()
        x = ['genes']
        x1 = ['transcripts']
        multiple_iso = sum(iso_per_gene['transcriptId'] > 1)
        mono_iso = sum((iso_per_gene['transcriptId'] > 1) == False)
        total_genes = multiple_iso + mono_iso
        labels = [
            str(round(1000 * multiple_iso / total_genes) / 10) + '%',
            str(round(1000 * mono_iso / total_genes) / 10) + '%'
        ]

        # Plot bars on the given axes
        mult = _ax.bar(x, multiple_iso, color='deepskyblue', label=labels[1])
        mono = _ax.bar(x, mono_iso, bottom=multiple_iso, color='sandybrown', label=labels[0])
        tran = _ax.bar(x1, len(adata_iso.var['transcriptId']))

        # Add text labels to the bars
        _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"
        )

        # Add a legend
        _ax.legend(['Multiple isoforms', 'Single isoform'])
        plt.show()

def plot_switch_gen_bar(adata_iso):
    """
    Public wrapper function to generate the bar plot.
    Creates a new figure and axes for the plot.

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.
    """
    _plot_switch_gen_bar(adata_iso, None)


# %% ../nbs/02_gene_report_iso_data.ipynb 15
# This code visualizes the frequency distribution of genes based on the number of isoforms they have.
# It creates a bar plot where the x-axis represents the number of isoforms per gene, 
# and the y-axis represents the number of genes with that specific number of isoforms.

def _plot_isoforms_frequencies(adata_iso, _ax):
    """
    Creates a bar chart to visualize the frequency distribution of isoforms per gene.

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.
        - `adata_iso.var`: Should include columns for 'transcriptId' and 'geneId'.
    - _ax (matplotlib.axes.Axes or None): The axes on which to render the plot.
        - If `None`, a new figure and axes are created.
        - If an `Axes` object is provided, the plot is rendered on the specified axes.
    """

    # Step 1: Calculate the number of isoforms per gene
    # - Group by 'geneId' and count the number of 'transcriptId' entries for each gene.
    iso_per_gene = adata_iso.var.reset_index().groupby(by='geneId').count()

    if _ax is None:
        # Create a new figure and axes if no axes object is provided
        fig, ax = plt.subplots()

        # Plot the frequency distribution as a bar chart
        # - `value_counts()` calculates the frequency of each unique isoform count.
        iso_per_gene['transcriptId'].value_counts().plot(
            ax=ax,  # Plot on the new axes
            kind='bar',  # Create a bar chart
            xlabel='number of isoforms per gene',  # Label for the x-axis
            ylabel='quantity of genes'  # Label for the y-axis
        )
    else:
        # If an axes object is provided, plot on the given axes
        iso_per_gene['transcriptId'].value_counts().plot(
            ax=_ax,  # Plot on the provided axes
            kind='bar',  # Create a bar chart
            xlabel='number of isoforms per gene',  # Label for the x-axis
            ylabel='quantity of genes'  # Label for the y-axis
        )

def plot_isoforms_frequencies(adata_iso):
    """
    Public wrapper function to create the bar chart showing isoform frequencies.
    It always creates a new figure and axes for the plot.

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.
    """
    # Call the private function with no axes provided, ensuring a new plot is created
    _plot_isoforms_frequencies(adata_iso, None)


# %% ../nbs/02_gene_report_iso_data.ipynb 17
# Boxplot with the number of genes expressed per cell type.
# This function visualizes the distribution of the number of genes expressed across different cell types.

def _plot_genes_cell_type(adata_iso, _ax):
    """
    Creates a boxplot and overlayed strip plot to show the number of genes expressed per cell type.

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.
        - `adata_iso.X`: The data matrix where rows are cells and columns are features (e.g., genes or isoforms).
        - `adata_iso.obs['cell_type']`: Cell type annotations for each cell.
    - _ax (matplotlib.axes.Axes or None): The axes on which to render the plot.
        - If `None`, a new figure and axes are created.
        - If an `Axes` object is provided, the plot is rendered on the specified axes.
    """

    if _ax is None:
        # Create a DataFrame with gene expression data grouped by cell type
        df = pd.DataFrame(np.transpose(adata_iso.X), columns=adata_iso.obs['cell_type'])

        # Sum the number of expressed genes for each cell type
        # - Sum expression values across all genes for each cell type.
        # - Convert the resulting Series into a DataFrame.
        df = df.sum(axis=0).to_frame().reset_index()

        # Rename the columns for clarity
        df.columns = ['cell_type', 'n_of_genes']  # 'n_of_genes' represents the total number of genes expressed per cell type.

        # Create a boxplot to visualize the distribution of the number of expressed genes per cell type
        ax = sns.boxplot(x='cell_type', y='n_of_genes', data=df)

        # Overlay a strip plot to show individual data points (gene counts per cell type)
        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:
        # Repeat the same process but render the plot on the provided axes (`_ax`)
        df = pd.DataFrame(np.transpose(adata_iso.X), columns=adata_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(adata_iso):
    """
    Public wrapper function to create a boxplot with the number of genes expressed per cell type.
    It always creates a new figure and axes for the plot.

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.
    """
    # Call the private function with no axes provided, ensuring a new plot is created
    _plot_genes_cell_type(adata_iso, None)


# %% ../nbs/02_gene_report_iso_data.ipynb 19
def plot_isoforms_summary(adata_iso):
    """
    Generates a summary visualization of isoform-related metrics for a dataset.

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.

    Workflow:
    1. Creates a bar plot showing the percentage of genes with multiple isoforms.
    2. Creates a bar plot showing the frequency distribution of isoforms per gene.
    3. Creates a boxplot showing the number of genes expressed per cell type.
    4. Combines these plots into a single figure and saves it.

    Returns:
    - Saves the combined figure to a file.
    """

    # - Use a pre-defined function `_plot_switch_gen_bar` to create the plot.
    ax1 = pw.Brick(figsize=(4, 4))
    _plot_switch_gen_bar(adata_iso, ax1)  # Plot genes with multiple isoforms
    ax1.set_title("Multiple isoforms genes %")  # Set the title of the first plot
    # - Use the `_plot_isoforms_frequencies` function to create the plot.
    ax2 = pw.Brick(figsize=(4, 4))
    _plot_isoforms_frequencies(adata_iso, ax2)  # Plot isoform frequencies
    ax2.set_title("Frequency of isoforms per gene")  # Set the title of the second plot
    # Create the third plot (Boxplot of the number of genes expressed per cell type)
    ax3 = pw.Brick(figsize=(3, 2))
    _plot_genes_cell_type(adata_iso, ax3)  # Plot the number of genes expressed per cell type
    ax3.set_title("Nb of genes per cell type")  # Set the title of the third plot

    # Combine the three plots into a single figure
    return (ax1 | ax2 | ax3).savefig()



# %% ../nbs/02_gene_report_iso_data.ipynb 23
# This function filters an AnnData object (`adata_iso`) to retain only genes with multiple isoforms.

def filter_switch_genes(adata_iso):
    """
    Subsets an AnnData object to include only genes with multiple isoforms (switching genes).

    Parameters:
    - adata_iso (AnnData): An AnnData object containing isoform-level data.
        - `adata_iso.var['geneId']`: Gene IDs associated with each transcript.

    Returns:
    - filtr_adata_iso (AnnData): A subset of the input AnnData object, containing only isoforms of genes
      that have more than one isoform (switching genes).
    """

    # Count the number of isoforms per gene
    # - `np.unique(adata_iso.var['geneId'], return_counts=True)`:
    #   - Identifies unique gene IDs in `adata_iso.var['geneId']`.
    #   - Returns two arrays:
    #     - `genes`: Unique gene IDs.
    #     - `frequencies`: The number of isoforms associated with each gene ID.
    genes, frequencies = np.unique(adata_iso.var['geneId'], return_counts=True)

    # Step 2: Create a DataFrame to store gene IDs and their isoform counts
    # - This helps in filtering genes with more than one isoform.
    df = pd.DataFrame({"geneId": genes, "frequencies": frequencies})

    # Step 3: Identify genes with more than one isoform
    # - Filter the DataFrame to include only rows where the 'frequencies' column is greater than 1.
    # - Convert the filtered 'geneId' column to a list of gene IDs with multiple isoforms.
    multi_iso_genes = df[df["frequencies"] > 1]['geneId'].tolist()

    # Step 4: Subset the AnnData object to retain only isoforms of genes with multiple isoforms
    # - Use `.var['geneId'].isin(multi_iso_genes)` to create a boolean mask for selecting isoforms.
    filtr_adata_iso = adata_iso[:, adata_iso.var['geneId'].isin(multi_iso_genes)]

    # Step 5: Return the filtered AnnData object
    return filtr_adata_iso


# %% ../nbs/02_gene_report_iso_data.ipynb 25
# A class extending AnnData to add specific functionality for isoform-level data analysis.
class AnnDataIso(ad.AnnData):

    # Private method to filter dataset for genes with multiple isoforms.
    def __filter_isodata(self):
        """
        Filters the data to include only genes with multiple isoforms.

        Logic:
        - Counts the number of isoforms for each gene using `np.unique`.
        - Creates a DataFrame with gene IDs and their isoform counts.
        - Identifies genes with more than one isoform.
        - Subsets the AnnData object to include only these genes.

        Returns:
        - Filtered AnnData object with only genes having multiple isoforms.
        """
        genes, frequencies = np.unique(self.var['geneId'], return_counts=True)
        df = pd.DataFrame({"geneId": genes, "frequencies": frequencies})
        multi_iso_genes = df[df["frequencies"] > 1]['geneId'].tolist()
        filtr_adata_iso = self[:, self.var['geneId'].isin(multi_iso_genes)]
        return filtr_adata_iso

    # Method to compute isoform expression percentages relative to total gene expression.
    def iso_percent(self, df, barcodes_regex="^[ACGT]+$"):
        """
        Calculates isoform expression percentages for each barcode.

        Parameters:
        - df (DataFrame): Input DataFrame containing isoform expression data.
        - barcodes_regex (str): Regex to identify columns representing barcodes.

        Returns:
        - DataFrame with normalized isoform percentages for each barcode.
        """
        iso_perc_df = df.__deepcopy__()
        if len(iso_perc_df.filter(regex=(barcodes_regex)).columns.to_list()) < 1:
            raise ValueError("No barcode was identified. Please check the names of the columns.")
        # Normalize isoform counts by gene-level expression.
        iso_perc_df[iso_perc_df.filter(regex=(barcodes_regex)).columns.to_list()] = (
            iso_perc_df.filter(regex=(barcodes_regex)) /
            iso_perc_df.groupby(['geneId']).transform('sum').filter(regex=(barcodes_regex))
        )
        # Replace NaN values with 0.
        iso_perc_df = iso_perc_df.replace(np.nan, 0.0)
        return iso_perc_df

    # Class initializer to process the AnnData object and compute relevant annotations.
    def __init__(self, anndata: ad.AnnData, cell_types: pd.DataFrame):
        """
        Initializes the AnnDataIso object with isoform and gene-level annotations.

        Parameters:
        - anndata (AnnData): Input AnnData object containing isoform-level data.
        - cell_types (DataFrame): DataFrame with cell type annotations.
        """
        self._init_as_actual(anndata.copy())
        # Group genes and count the number of isoforms.
        self.gene_counts = self.var.reset_index().groupby(by='geneId').count()
        # Filter dataset for genes with multiple isoforms.
        self.__filtered_anndata = self.__filter_isodata()
        # Annotate cell types in the observation data.
        self.obs['cell_type'] = cell_types
        # Compute isoform percentages.
        df = self.__filtered_anndata.to_df().set_index(self.__filtered_anndata.obs['barcodes'])
        df = df.transpose()
        df[['transcriptId', 'geneId']] = self.__filtered_anndata.var
        df_m_iso = self.iso_percent(df)
        df_m_iso = df_m_iso.iloc[0:, :-2].transpose()
        self.__filtered_anndata.obsm['Iso_prct'] = df_m_iso

    # Plotting a summary of isoform-related metrics.
    def plot_isoforms_summary(self):
        """
        Generates a combined plot summarizing:
        1. Percentage of genes with multiple isoforms.
        2. Frequency distribution of isoforms per gene.
        3. Number of genes expressed per cell type.

        Returns:
        - Combined plot saved as an image.
        """
        ax1 = pw.Brick(figsize=(4, 4))
        self._plot_switch_gen_bar(ax1)
        ax1.set_title("Multiple isoforms genes %")
        ax2 = pw.Brick(figsize=(4, 4))
        self._plot_isoforms_frequencies(ax2)
        ax2.set_title("Frequency of isoforms per gene")
        ax3 = pw.Brick(figsize=(3, 2))
        self._plot_genes_cell_type(ax3)
        ax3.set_title("Nb of genes per cell type")
        return (ax1 | ax2 | ax3).savefig()

    # Private method to create a boxplot of genes expressed per cell type.
    def _plot_genes_cell_type(self, _ax):
        """
        Plots the number of genes expressed per cell type as a boxplot.

        Parameters:
        - _ax (Axes): Matplotlib axes to render the plot.
        """
        if _ax is None:
            df = pd.DataFrame(np.transpose(self.X), columns=self.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(self.X), columns=self.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(self):
        """
        Public method to plot the number of genes expressed per cell type.
        """
        self._plot_genes_cell_type(None)

    # Method to visualize isoform frequency distribution.
    def _plot_isoforms_frequencies(self, _ax):
        """
        Plots the frequency distribution of isoforms per gene.

        Parameters:
        - _ax (Axes): Matplotlib axes to render the plot.
        """
        if _ax is None:
            fig, ax = plt.subplots()
            self.gene_counts['transcriptId'].value_counts().plot(
                ax=ax, kind='bar', xlabel='number of isoforms per gene', ylabel='quantity of genes'
            )
        else:
            self.gene_counts['transcriptId'].value_counts().plot(
                ax=_ax, kind='bar', xlabel='number of isoforms per gene', ylabel='quantity of genes'
            )

    def plot_isoforms_frequencies(self):
        """
        Public method to plot the frequency distribution of isoforms per gene.
        """
        self._plot_isoforms_frequencies(None)

    def _plot_transcripts_per_cell_type(self, gene_name, trs_to_show, _ax):
        if trs_to_show == []:
            transcripts_id = self.__get_transcripts_from_gene(gene_name)
        else:
            transcripts_id = trs_to_show
        if _ax is None:
            fig, ax = plt.subplots()
            grouped = self.__filtered_anndata.obsm['Iso_prct']
            grouped['cell_type'] = self.obs['cell_type']
            res = grouped.groupby('cell_type').mean().transpose()
            res = res.assign(transcriptId=self.__filtered_anndata.var['transcriptId'].to_list())
            res = res.assign(geneId=self.__filtered_anndata.var['geneId'].to_list())
            res = res[res['geneId'] == gene_name].drop(['geneId'], axis=1)
            res = res[res['transcriptId'].isin(transcripts_id)]
            plot_data = res.set_index('transcriptId').transpose()
            plot_data.plot(kind='barh', ax=ax, stacked=True, color=self.colors).legend(loc='center left',bbox_to_anchor=(1.0, 1.0))
            plt.legend(self.get_transcripts_common_names(transcripts_id), loc="upper left", bbox_to_anchor=(1, 1))
            #plt.legend(self.get_transcripts_common_names(trs_to_show))
            plt.ylabel('Cell type')
        else:
            grouped = self.__filtered_anndata.obsm['Iso_prct']
            grouped['cell_type'] = self.obs['cell_type']
            res = grouped.groupby('cell_type').mean().transpose()
            res = res.assign(transcriptId=self.__filtered_anndata.var['transcriptId'].to_list())
            res = res.assign(geneId=self.__filtered_anndata.var['geneId'].to_list())
            res = res[res['geneId'] == gene_name].drop(['geneId'], axis=1)
            res = res[res['transcriptId'].isin(transcripts_id)]
            plot_data = res.set_index('transcriptId').transpose()
            plot_data.plot(kind='barh', stacked=True, ax=_ax).legend(loc='center left',bbox_to_anchor=(1.0, 1.0))
            
    def plot_transcripts_per_cell_type(self, gene_name, trs_to_show = []):
        self._plot_transcripts_per_cell_type(gene_name, trs_to_show, None)
    
    def _trsct_counts_cell_type (self, gene_name, trs_to_show, _ax):
        if _ax is None:
        # create df with filtered isoform counts and labeled cell types:
            df = self.__filtered_anndata.to_df().set_index(self.__filtered_anndata.obs['cell_type'])
            df = df.transpose() 
            df[['transcriptId', 'geneId']] = self.__filtered_anndata.var[['transcriptId', 'geneId']]
            gene_iso_count = df[df['geneId'] == gene_name]
            gene_iso_count = gene_iso_count.drop('geneId', axis=1).set_index('transcriptId').transpose()
            gene_iso_count_long = gene_iso_count.reset_index().melt(id_vars='cell_type', var_name='transcriptId', value_name='count')
            if trs_to_show != []:
                gene_iso_count_long = gene_iso_count_long[gene_iso_count_long['transcriptId'].isin(trs_to_show)]
            g = sns.catplot(x="cell_type", y="count", col="transcriptId", aspect=1, dodge=False, kind="violin", data=gene_iso_count_long)
            # Set custom facet titles
            g.set_titles(col_template="{col_name}", size = 8)
            # Remove x ticks
            g.set_xticklabels(rotation=90)
            g.fig.suptitle(gene_name)
            plt.show()
        else:
            df = self.__filtered_anndata.to_df ().set_index(self.__filtered_anndata.obs['cell_type'])
            df = df.transpose()
            df[['transcriptId', 'geneId']] = self.__filtered_anndata.var[['transcriptId', 'geneId']]
            gene_iso_count = df[df['geneId']== gene_name]
            gene_iso_count = gene_iso_count.drop('geneId', axis=1).set_index('transcriptId').transpose()
            gene_iso_count_long = gene_iso_count.reset_index().melt(id_vars='cell_type', var_name='transcriptId', value_name='count')
            if trs_to_show != []:
                gene_iso_count_long = gene_iso_count_long[gene_iso_count_long['transcriptId'].isin(trs_to_show)]
            g = sns.catplot(x="cell_type", y="count", col="transcriptId", aspect=1, dodge=False, kind="violin", data=gene_iso_count_long)
            # Set custom facet titles
            g.set_titles(col_template="{col_name}", size = 8)
            # Remove x ticks
            g.set_xticklabels(rotation=90, labels=self.obs['cell_type'].unique())
            g.fig.suptitle(gene_name)
            return g

    
    def trsct_counts_cell_type (self, gene_name, trs_to_show = []):
        # create df with filtered isoform counts and labeled cell types:
        self._trsct_counts_cell_type(gene_name, trs_to_show, None)

    def __get_coord_from_tscrpt_id(self, transcript_id):
        if '.' in transcript_id:
            transcript_id = transcript_id.split('.')[0]
        server = "https://rest.ensembl.org"
        ext = "/lookup/id/" + transcript_id + "?expand=1"
    
        r = requests.get(server+ext, headers={ "Content-Type" : "application/json"})
    
        if not r.ok:
            r.raise_for_status()
            sys.exit()
    
        decoded = r.json()
        exon_list = list(decoded['Exon'])
        exon_coord = []
        for i, e in enumerate(exon_list):
            coord = [e.get('end'), e.get('start')]
            exon_coord.append(coord)
        strand = decoded['strand']
        return(exon_coord, strand)

    def __draw_exons(self, exons, direction, color, transcript_name, offset=0, start_override=None, end_override=None, no_render=False):
        if not no_render:
            plt.axes()
            plt.xlim((-0.1, 1))
            plt.ylim((-0.3, 0.3))
            plt.margins(0.2)
            plt.axis('off')
            fig = plt.gcf()
            fig.set_size_inches(20, 2)
        height = 0.2
        plt.plot([offset + 0.1, offset + 0.1], linestyle='solid', linewidth=0.5, c='grey')
        j = 0
        k = 1
        if direction == 1:
            pos_start = exons[0][1]
            pos_end = exons[-1][0]
        else: #direction == -1
            pos_start = exons[-1][1]
            pos_end = exons[0][0]
            j = 1
            k = 0
        real_start = pos_start
        real_end = pos_end
        if start_override is not None and end_override is not None:
            pos_start = start_override
            pos_end = end_override
        total_length = pos_end - pos_start
        total_length_with_margin = 1.05 * total_length
        pos_start_with_margin = pos_start - 0.025*total_length
        for i, exon in enumerate(exons):
            rectangle = plt.Rectangle(((exon[j] - pos_start_with_margin)/total_length_with_margin,offset), (exon[k] - exon[j])/total_length_with_margin, height, fc=color,ec="black")
            plt.gca().add_patch(rectangle)
        if i > 0:
            arrow = None
            if direction < 0:
                arrow = plt.arrow(1, offset - height/4, -1, 0, width=0.0015, head_length=0.01, head_width=0.1, length_includes_head=True, overhang=1)
            else:
                arrow = plt.arrow(0, offset - height/4, 1, 0, width=0.0015, head_length=0.01, head_width=0.1, length_includes_head=True, overhang=1)
            plt.gca().add_patch(arrow)
        plt.plot(np.array([0.025 + (real_start - pos_start) / (total_length)/1.05, 0.025 + (real_start - pos_start) / (total_length)/1.05]), np.array([offset - height/4 - 0.03, offset - height/4 + 0.03]), color='black')
        plt.plot(np.array([1 - 0.025 - (pos_end - real_end) / (total_length)/1.05, 1 - 0.025 - (pos_end - real_end) / (total_length)/1.05]), np.array([offset - height/4 - 0.03, offset - height/4 + 0.03]), color='black')
        plt.text(0.025 + (real_start - pos_start) / (total_length)/1.05, offset - height/4 - 0.075, real_start, horizontalalignment='center', verticalalignment='center', fontsize=9)
        plt.text(1 - 0.025 - (pos_end - real_end) / (total_length)/1.05, offset - height/4 - 0.075, real_end, horizontalalignment='center', verticalalignment='center', fontsize=9)
        plt.text(1, offset - height, transcript_name, horizontalalignment='right', verticalalignment='top', fontsize=12)
        if not no_render:
            plt.show()
    
    def __get_transcripts_from_gene(self, gene_name):
        elems = self.__filtered_anndata.var
        return elems[elems['geneId'] == gene_name]['transcriptId'].to_list()

    def __draw_transcripts_list(self, gene_name, trs_to_show, _ax, colors=None):
        if trs_to_show == []:
            transcripts_id = self.__get_transcripts_from_gene(gene_name)
        else:
            transcripts_id = trs_to_show
        exons = []
        directions = []
        for tr in transcripts_id:
            t, d = self.__get_coord_from_tscrpt_id(tr)
            exons += [t]
            directions += [d]
        if colors is None:
            colors = []
            for i in range(len(exons)):
                colors.append(['#898D90','#8D93A1','#9F99B5','#AFACC9','#D7CADE','#DAEDF3','#F7EABD'][i % 7])
        def get_limits(ex, dir):
            start = sys.maxsize
            end = -sys.maxsize
            for (e, d) in zip(ex, dir):
                if d == 1:
                    start = min(start,e[0][1])
                    end = max(end, e[-1][0])
                else:
                    start = min(start,e[-1][1])
                    end = max(end, e[0][0])
            return (start, end)
        plt.axes()
        plt.xlim((-0.1, 1.1))
        plt.ylim((0.1 - 0.5 *  len(exons), 0.3))
        plt.margins(0.2)
        plt.axis('off')
        fig = plt.gcf()
        fig.set_size_inches(20, len(exons) * 2)
        i = 0
        (start, end) = get_limits(exons, directions)
        for (ex, di, co, name) in zip(exons, directions, colors, transcripts_id):
            self.__draw_exons(ex, di, co, name, offset= -0.5 * i, start_override=start, end_override=end, no_render=True)
            i+=1
        if _ax is None:
            plt.show()
        else: 
            return plt

    def draw_transcripts_list(self, gene_name, colors=None):
        self.__draw_transcripts_list(gene_name, [], None, colors)
    
    def draw_gene_summary(self, gene_name, trs_to_show = []):
        ax1 = pw.Brick(figsize=(12,4))
        pw.overwrite_axisgrid()
        fg = self._trsct_counts_cell_type(gene_name, trs_to_show, ax1)
        ax4 = pw.load_seaborngrid(fg)
        ax1.set_title("Transcripts count per cell type")
        ax2 = pw.Brick(figsize=(12,4))
        self._plot_transcripts_per_cell_type(gene_name, trs_to_show, ax2)
        ax2.set_title("Transcripts count per cell type")
        ax3 = pw.Brick(figsize=(12,4))
        self.__draw_transcripts_list(gene_name, trs_to_show, ax3)
        ax3.set_title("Transcripts list")
        return (ax4/(ax2/ax3)).savefig()
    
    def _trsct_counts_cell_type (self, gene_name, trs_to_show, _ax):
        if _ax is None:
        # create df with filtered isoform counts and labeled cell types:
            df = self.__filtered_anndata.to_df().set_index(self.__filtered_anndata.obs['cell_type'])
            df = df.transpose() 
            df[['transcriptId', 'geneId']] = self.__filtered_anndata.var[['transcriptId', 'geneId']]
            gene_iso_count = df[df['geneId'] == gene_name]
            gene_iso_count = gene_iso_count.drop('geneId', axis=1).set_index('transcriptId').transpose()
            gene_iso_count_long = gene_iso_count.reset_index().melt(id_vars='cell_type', var_name='transcriptId', value_name='count')
            if trs_to_show != []:
                gene_iso_count_long = gene_iso_count_long[gene_iso_count_long['transcriptId'].isin(trs_to_show)]
            g = sns.catplot(x="cell_type", y="count", col="transcriptId", aspect=1, dodge=False, kind="violin", data=gene_iso_count_long)
            # Set custom facet titles
            g.set_titles(col_template="{col_name}", size = 8)
            # Remove x ticks
            g.set_xticklabels(rotation=90)
            g.fig.suptitle(gene_name)
            plt.show()
        else:
            df = self.__filtered_anndata.to_df ().set_index(self.__filtered_anndata.obs['cell_type'])
            df = df.transpose()
            df[['transcriptId', 'geneId']] = self.__filtered_anndata.var[['transcriptId', 'geneId']]
            gene_iso_count = df[df['geneId']== gene_name]
            gene_iso_count = gene_iso_count.drop('geneId', axis=1).set_index('transcriptId').transpose()
            gene_iso_count_long = gene_iso_count.reset_index().melt(id_vars='cell_type', var_name='transcriptId', value_name='count')
            if trs_to_show != []:
                gene_iso_count_long = gene_iso_count_long[gene_iso_count_long['transcriptId'].isin(trs_to_show)]
            g = sns.catplot(x="cell_type", y="count", col="transcriptId", aspect=1, dodge=False, kind="violin", data=gene_iso_count_long)
            # Set custom facet titles
            g.set_titles(col_template="{col_name}", size = 8)
            # Remove x ticks
            g.set_xticklabels(rotation=90, labels=self.obs['cell_type'].unique())
            g.fig.suptitle(gene_name)
            return g

    def __draw_transcripts_list(self, gene_name, trs_to_show, _ax, colors=None):
        if trs_to_show == []:
            transcripts_id = self.__get_transcripts_from_gene(gene_name)
        else:
            transcripts_id = trs_to_show
        exons = []
        directions = []
        for tr in transcripts_id:
            t, d = self.__get_coord_from_tscrpt_id(tr)
            exons += [t]
            directions += [d]
        if colors is None:
            colors = []
            for i in range(len(exons)):
                colors.append(['#898D90','#8D93A1','#9F99B5','#AFACC9','#D7CADE','#DAEDF3','#F7EABD'][i % 7])
        def get_limits(ex, dir):
            start = sys.maxsize
            end = -sys.maxsize
            for (e, d) in zip(ex, dir):
                if d == 1:
                    start = min(start,e[0][1])
                    end = max(end, e[-1][0])
                else:
                    start = min(start,e[-1][1])
                    end = max(end, e[0][0])
            return (start, end)
        plt.axes()
        plt.xlim((-0.1, 1.1))
        plt.ylim((0.1 - 0.5 *  len(exons), 0.3))
        plt.margins(0.2)
        plt.axis('off')
        fig = plt.gcf()
        fig.set_size_inches(20, len(exons) * 2)
        i = 0
        (start, end) = get_limits(exons, directions)
        for (ex, di, co, name) in zip(exons, directions, colors, transcripts_id):
            self.__draw_exons(ex, di, co, name, offset= -0.5 * i, start_override=start, end_override=end, no_render=True)
            i+=1
        if _ax is None:
            plt.show()
        else: 
            return plt

    def draw_transcripts_list(self, gene_name, colors=None):
        self.__draw_transcripts_list(gene_name, [], None, colors)

    # More detailed functions omitted for brevity.


# %% ../nbs/02_gene_report_iso_data.ipynb 38
#trsct_counts_cell_type(gene)
