import pandas as pd
import scanpy as sc
import numpy as np
import scanpy.external as sce
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns



class remove_mt_ribo:

    def __init__(self, anndata):
        # anndata = the unprocessed anndata with spatial coordinates generated by MakeAnndata
        self.anndata = anndata

        # metrics for mito genes
        self.mito_obs, self.mito_data = self._before_mt_ribo()

        # Filtered anndata without mt and ribo genes
        self.data = self._remove_mt_ribo()

    def _before_mt_ribo(self):
        # Function to get metrics before filtering
        anndata = self.anndata
        # mitochondrial genes
        anndata.var['mt'] = anndata.var_names.str.startswith('MT-')
        # ribosomal genes
        anndata.var['ribo'] = anndata.var_names.str.startswith(("RPS", "RPL"))
        # hemoglobin genes.
        anndata.var['hb'] = anndata.var_names.str.contains(("^HB[^(P)]"))

        sc.pp.calculate_qc_metrics(anndata, qc_vars=['mt', 'ribo', 'hb'], percent_top=None, log1p=False, inplace=True)

        anndata = self.anndata
        mito_genes = anndata.var_names.str.startswith('MT-')
        # for each cell compute fraction of counts in mito genes vs. all genes
        # the `.A1` is only necessary as X is sparse (to transform to a dense array after summing)
        anndata.obs['percent_mt2'] = np.sum(
            anndata[:, mito_genes].X, axis=1).A1 / np.sum(anndata.X, axis=1).A1
        # add the total counts per cell as observations-annotation to adata
        anndata.obs['n_counts'] = anndata.X.sum(axis=1).A1

        return anndata[:, mito_genes].obs, anndata

    def _remove_mt_ribo(self):
        # Function to filter out noise genes
        anndata = self.anndata
        # remove mt genes, malat1 (non-informative) and hb-genes -
        malat1 = anndata.var_names.str.startswith('MALAT1')
        # need to redefine the mito_genes since they were first calculated on the full object before removing low expressed genes.
        mito_genes = anndata.var_names.str.startswith('MT-')
        hb_genes = anndata.var_names.str.contains('^HB[^(P)]')

        remove = np.add(mito_genes, malat1)
        remove = np.add(remove, hb_genes)
        keep = np.invert(remove)

        anndata = anndata[:, keep]

        return anndata
