import matplotlib.pyplot as plt

from .MakeAnndata import MakeAnndata
from .remove_mt_ribo import remove_mt_ribo
from matplotlib_venn import venn3
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import scanpy as sc
import scanpy.external as sce
class analysation:
    def __init__(self,h5ad_path, barcodes_path, cell_cycle_gene_path, resolution):
        # h5ad_path = path to h5ad file generated by CURIO pipeline
        # barcodes_path = path to bead barcode whitelist for specific slide
        self.barcodes_path = barcodes_path
        self.h5ad_path = h5ad_path
        # cell cycle genes
        self.cell_cycle_gene_path = cell_cycle_gene_path

        self.s_genes,self.g2m_genes = self.cell_cycle()

        self.resolution = resolution

        # The mitocondrial statisitcs:
        self.mito_obs, self.mito_data = self.mito_stat()

        self.normalised_data = self.preprocessing_anndata()

        self.final_anndata, self.ranked_genes = self.clustering()

    def mito_stats(self):
        # Function which gives statistics of mitocondrial and ribosomal genes

        tmp = MakeAnndata(self.h5ad_path, self.barcodes_path)

        data = tmp.data

        tmp = remove_mt_ribo(data)

        mito_obs = tmp.mito_obs
        mito_data = tmp.mito_data

        sc.settings.set_figure_params(dpi_save=400, fontsize=30, figsize=(10, 10), facecolor='white', frameon=False,
                                      transparent=True, format="pdf")

        sc.pl.violin(mito_data, ['n_genes_by_counts', 'total_counts', 'pct_counts_mt', 'pct_counts_ribo', 'pct_counts_hb'],
                     jitter=0.4, rotation=45, title= "Violin Plots of Quality Indicators")
        sc.pl.violin(mito_data, ['pct_counts_mt'],
                     jitter=0.4, rotation=45, title = "Violin Plot of Mito Data")
        sc.pl.scatter(mito_data, x='total_counts', y='pct_counts_mt', title= "Scatter Plot of pct_counts_mt")
        sc.pl.scatter(mito_data, x='total_counts', y='n_genes_by_counts', title="Scatter Plot of n_genes_by_counts")

        sc.settings.set_figure_params(dpi_save=400, fontsize=6, figsize=(10, 10), facecolor='white', frameon=False,
                                      transparent=True, format="pdf")


        sc.pl.spatial(
        mito_data,
        color=["total_counts", "n_genes_by_counts", "pct_counts_mt"],
        spot_size=30, vmin=0, vmax='p99.2', save="heart_spatial_before_filtering")

        sc.pl.spatial(
        mito_data,
        color=["pct_counts_mt"],
        spot_size=30, vmin=0, vmax='p99.2', save="heart_spatial_before_filtering")

        data.obs.pct_counts_mt.hist(bins=100)
        plt.title("Histogram of Percentage Counts Mito")

        data.obs.pct_counts_ribo.hist(bins=100)
        plt.title("Histogram of Percentage Counts Ribo")

        return mito_obs, mito_data

    def cell_cycle(self):
        cell_cycle_genes = [x.strip() for x in open(self.cell_cycle_gene_path)]
        # Split into 2 lists
        s_genes = cell_cycle_genes[:43]
        g2m_genes = cell_cycle_genes[43:]

        cell_cycle_genes = [x for x in cell_cycle_genes if x in data.var_names]
        print(len(cell_cycle_genes))

        return s_genes, g2m_genes

    def preprocessing_anndata(self):

        tmp = MakeAnndata(self.h5ad_path, self.barcodes_path)
        data = tmp.data
        tmp = remove_mt_ribo(data)

        data = tmp.data

        sc.pl.highest_expr_genes(data, n_top=20)

        # Normalise - cpm = target_sum = 1e6
        sc.pp.normalize_total(data, inplace=True, target_sum=1e6)

        # Log - OBS! Bug in write_h5ad - see lower down.
        sc.pp.log1p(data)

        # Note that you do want to exclude junk - high expression probably common genes expressed in all cells
        sc.pp.highly_variable_genes(data, max_mean=10, min_mean=0, min_disp=0.25)
        sc.pl.highly_variable_genes(data)

        sc.tl.score_genes_cell_cycle(data, s_genes=self.s_genes, g2m_genes=self.g2m_genes)

        # Select only high var genes
        data.raw = data
        data = data[:, data.var.highly_variable]

        sc.pl.violin(data, ['S_score', 'G2M_score'],
                     jitter=0.4, rotation=45)

        # Scale before pca with exclusion of given sd
        sc.pp.scale(data, max_value=10, zero_center=False)

        # regress out unnecessary sources of variation
        # sc.pp.regress_out(data, ['total_counts', 'pct_counts_mt'])
        # data.raw = data
        sc.pp.pca(data, random_state=0)

        # captured variance
        sc.pl.pca_variance_ratio(data, n_pcs=50, log=True)
        # Calculate neighbour hood graph using pc
        sc.pp.neighbors(data, n_pcs=5)
        # Embedd neighbourhood graph through umap
        sc.tl.umap(data)
        # confirm scaling

        data.to_df().mean(axis=0).hist(bins=100)

        data.to_df().iloc[:, 1000:2000].var(axis=0).hist()

        return data

    def clustering(self):
        data.uns['log1p'] = {'base': None}  # Bug in scanpy! OBS!
        # Leiden clustering
        sc.tl.leiden(data, key_added="leiden_0.3", resolution = 0.3)
        sc.tl.leiden(data, key_added="leiden_0.5", resolution = 0.5)
        sc.tl.leiden(data, key_added="leiden_0.8", resolution = 0.8)
        sc.tl.leiden(data, key_added="leiden_1.0", resolution = 1.0)

        sc.pl.spatial(data, color=["leiden_0.3","leiden_0.5","leiden_0.8", "leiden_1.0"], spot_size=30, wspace=0.5)

        if self.resolution == None:
            sc.pl.spatial(data, color=["leiden_1.0"], spot_size=30, wspace=0.5)
            sc.tl.dendrogram(data, groupby="leiden_1.0")
            sc.pl.dendrogram(data, groupby="leiden_1.0")
            sc.pl.umap(data, color=["leiden_1.0"], palette="tab20")

            sc.tl.rank_genes_groups(data, 'leiden_1.0', method='wilcoxon', key_added='wilcoxon')
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="wilcoxon")

            sc.tl.rank_genes_groups(data, 'leiden_1.0', method='t-test', key_added='t-test')
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="t-test")

            sc.tl.rank_genes_groups(data, 'leiden_1.0', method='logreg', key_added="logreg")
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="logreg")

            sc.tl.rank_genes_groups(data, 'leiden_1.0', method='t-test_overestim_var', key_added="t-test_ov")
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="t-test_ov")

            ranked_genes = pd.DataFrame(data.uns['wilcoxon']['names'])

            sc.pl.rank_genes_groups_heatmap(data, n_genes=5, key="wilcoxon", groupby="leiden_1.0", show_gene_labels=True)

            sc.pl.rank_genes_groups_dotplot(data, n_genes=5, key="wilcoxon", groupby="leiden_1.0")

            sc.pl.rank_genes_groups_stacked_violin(data, n_genes=5, key="wilcoxon", groupby="leiden_1.0")

            sc.pl.rank_genes_groups_matrixplot(data, n_genes=5, key="wilcoxon", groupby="leiden_1.0")


        else:
            sc.pl.spatial(data,color=[f"leiden_{self.resolution}"], spot_size=30, wspace=0.5)
            sc.tl.dendrogram(data, groupby=f"leiden_{self.resolution}")
            sc.pl.dendrogram(data, groupby=f"leiden_{self.resolution}")
            sc.pl.umap(data, color=[f"leiden_{self.resolution}"], palette="tab20")

            sc.tl.rank_genes_groups(data, f"leiden_{self.resolution}", method='wilcoxon', key_added='wilcoxon')
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="wilcoxon")

            sc.tl.rank_genes_groups(data, f"leiden_{self.resolution}", method='t-test', key_added='t-test')
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="t-test")

            sc.tl.rank_genes_groups(data, f"leiden_{self.resolution}", method='logreg', key_added="logreg")
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="logreg")

            sc.tl.rank_genes_groups(data, f"leiden_{self.resolution}", method='t-test_overestim_var', key_added="t-test_ov")
            sc.pl.rank_genes_groups(data, n_genes=25, sharey=False, key="t-test_ov")

            ranked_genes = pd.DataFrame(data.uns['wilcoxon']['names'])
            sc.pl.rank_genes_groups_heatmap(data, n_genes=5, key="wilcoxon", groupby=f"leiden_{self.resolution}", show_gene_labels = True)

            sc.pl.rank_genes_groups_dotplot(data, n_genes=5, key="wilcoxon", groupby=f"leiden_{self.resolution}")

            sc.pl.rank_genes_groups_stacked_violin(data, n_genes=5, key="wilcoxon", groupby=f"leiden_{self.resolution}")

            sc.pl.rank_genes_groups_matrixplot(data, n_genes=5, key="wilcoxon", groupby=f"leiden_{self.resolution}")

        wc = sc.get.rank_genes_groups_df(data, group='0', key='wilcoxon', pval_cutoff=0.01, log2fc_min=0)['names']
        tt = sc.get.rank_genes_groups_df(data, group='0', key='t-test', pval_cutoff=0.01, log2fc_min=0)['names']
        tt_ov = sc.get.rank_genes_groups_df(data, group='0', key='t-test_ov', pval_cutoff=0.01, log2fc_min=0)['names']


        # compare cluster1 genes, only stores top 100 by default
        venn3([set(wc), set(tt), set(tt_ov)], ('Wilcox', 'T-test', 'T-test_ov'))
        plt.show()

        return data, ranked_genes




