import scanpy as sc
import numpy as np
import pandas as pd

from celescope.tools import utils
from celescope.__init__ import HELP_DICT
from celescope.rna.mkref import Mkref_rna
from celescope.tools.step import Step, s_common

# markers adjust p_value
PVAL_CUTOFF = 0.05
# scanpy mitochonrial variable name
MITO_VAR = 'mito'
NORMALIZED_LAYER = 'normalised'
RESOLUTION = 1.2
N_PCS = 25
MITO_GENE_PERCENT_LIST = [5, 10, 15, 20, 50]
# output marker top n in html
MARKER_TOP_N = 100


def read_tsne(tsne_file):
    df = pd.read_csv(tsne_file, sep='\t', index_col=0)
    df.index.names = ['barcode']
    return df

def format_df_marker(df_marker):

    avg_logfc_col = "avg_log2FC"  # seurat 4
    if "avg_logFC" in df_marker.columns:  # seurat 2.3.4
        avg_logfc_col = "avg_logFC"
    df_marker = df_marker.loc[:,
                                ["cluster", "gene", avg_logfc_col, "pct.1", "pct.2", "p_val_adj"]
                                ]
    df_marker["cluster"] = df_marker["cluster"].apply(lambda x: f"cluster {x}")
    df_marker = df_marker[df_marker["p_val_adj"] < PVAL_CUTOFF]

    return df_marker

def get_opts_analysis(parser, sub_program):
    
    parser.add_argument('--genomeDir', help=HELP_DICT['genomeDir'], required=True)
    if sub_program:
        parser.add_argument(
            '--matrix_file',
            help='Required. Matrix directory path.',
            required=True,
        )
        parser = s_common(parser)


class Scanpy_wrapper(Step):

    def __init__(self, args, display_title=None):

        super().__init__(args, display_title=display_title)

        # data
        self.adata = sc.read_10x_mtx(
            args.matrix_file,  
            var_names='gene_symbols',
        )
        self.adata.layers['raw'] = self.adata.X.copy()
        self.mt_gene_list = Mkref_rna.get_config(args.genomeDir)['files']['mt_gene_list']

        # out
        self.df_marker_file = f'{self.outdir}/markers.tsv'
        self.df_marker_raw_file = f'{self.outdir}/markers_raw.tsv'
        self.df_tsne_file = f'{self.outdir}/tsne_coord.tsv'
        self.h5ad_file = f'{self.outdir}/rna.h5ad'
        self.outs = [self.df_marker_file, self.df_tsne_file, self.h5ad_file]

   
    @utils.add_log
    def calculate_qc_metrics(self):

        if self.mt_gene_list:
            mito_genes, _ = utils.read_one_col(self.mt_gene_list)
            self.adata.var[MITO_VAR] = self.adata.var_names.map(lambda x:True if x in mito_genes else False)
            # if not astype(bool), it will be type object and raise an error
            # https://github.com/theislab/anndata/issues/504
            self.adata.var[MITO_VAR] = self.adata.var[MITO_VAR].astype(bool)
        else:
            self.adata.var[MITO_VAR] = self.adata.var_names.str.upper().str.startswith('MT-')

        sc.pp.calculate_qc_metrics(
            self.adata,
            qc_vars=[MITO_VAR], 
            percent_top=None,
            use_raw=False,
            log1p=False, 
            inplace=True
        )

    @utils.add_log
    def write_mito_stats(self):

        mt_pct_var = f'pct_counts_{MITO_VAR}'
        total_cell_number = self.adata.n_obs

        for mito_gene_percent in MITO_GENE_PERCENT_LIST:
            cell_number = sum(self.adata.obs[mt_pct_var] > mito_gene_percent)
            fraction = round(cell_number / total_cell_number * 100, 2)
            self.add_metric(
                name=f'Fraction of cells have mito gene percent>{mito_gene_percent}%',
                value=f'{fraction}%',
            )


    @utils.add_log
    def normalize(self):
        """
        sc.pp.normalize_per_cell() and sc.pp.log1p()
        """

        sc.pp.normalize_total(
            self.adata,
            target_sum=1e4,
            inplace=True,
        )
        sc.pp.log1p(
            self.adata,
        )
        self.adata.layers[NORMALIZED_LAYER] = self.adata.X

    @utils.add_log
    def hvg(self):
        """
        Wrapper function for sc.highly_variable_genes()
        """
        sc.pp.highly_variable_genes(
            self.adata,
            layer=None,
            n_top_genes=None,
            min_disp=0.5,
            max_disp=np.inf,
            min_mean=0.0125,
            max_mean=3,
            span=0.3,
            n_bins=20,
            flavor='seurat',
            subset=False,
            inplace=True,
            batch_key=None,
            check_values=True
        )        

    @utils.add_log
    def scale(self):
        """
        Wrapper function for sc.pp.scale
        """
        sc.pp.scale(
            self.adata,
            zero_center=True,
            max_value=10,
            copy=False,
            layer=None,
            obsm=None
        )
        
    @utils.add_log
    def pca(self):
        """
        Wrapper function for sc.pp.pca
        """
        sc.pp.pca(
            self.adata,
            n_comps=50,
            zero_center=True,
            svd_solver='auto',
            random_state=0,
            return_info=False,
            use_highly_variable=True,
            dtype='float32',
            copy=False,
            chunked=False,
            chunk_size=None
        )       

    @utils.add_log
    def neighbors(self, ):
        """
        Wrapper function for sc.pp.neighbors(), for supporting multiple n_neighbors
        """
        sc.pp.neighbors(
            self.adata,
            n_neighbors=15,
            n_pcs=N_PCS,
            use_rep=None,
            knn=True,
            random_state=0,
            method='umap',
            metric='euclidean',
            key_added=None,
            copy=False
        )

    @utils.add_log
    def tsne(self):
        """
        Wrapper function for sc.tl.tsne, for supporting named slot of tsne embeddings
        """
        sc.tl.tsne(
            self.adata,
            n_pcs=N_PCS,
            n_jobs=self.thread,
            copy=False,
        )

    @utils.add_log
    def umap(self, ):
        """
        Wrapper function for sc.tl.umap, for supporting named slot of umap embeddings
        """
        sc.tl.umap(
            self.adata,
            min_dist=0.5,
            spread=1.0,
            n_components=2,
            maxiter=None,
            alpha=1.0,
            gamma=1.0,
            negative_sample_rate=5,
            init_pos='spectral',
            random_state=0,
            a=None,
            b=None,
            copy=False,
            method='umap',
            neighbors_key=None
        )

        
    @utils.add_log
    def leiden(self):
        """
        Wrapper function for sc.tl.leiden
        """
        sc.tl.leiden(
            self.adata,
            resolution=RESOLUTION,
            restrict_to=None,
            random_state=0,
            key_added='cluster',
            adjacency=None,
            directed=True,
            use_weights=True,
            n_iterations=-1,
            partition_type=None,
            neighbors_key=None,
            obsp=None,
            copy=False
        )

    @utils.add_log
    def find_marker_genes(self):
        """
        Wrapper function for sc.tl.rank_genes_groups
        """
        sc.tl.rank_genes_groups(
            self.adata,
            "cluster",
            reference='rest',
            pts=True,
            method="wilcoxon",
            use_raw=False,
            layer=NORMALIZED_LAYER
        )

    @utils.add_log
    def write_markers(self):
        '''
        write only p_val_adj < PVAL_CUTOFF to avoid too many markers
        '''
        df_markers = sc.get.rank_genes_groups_df(self.adata, group=None, pval_cutoff=PVAL_CUTOFF)
        df_markers = df_markers[df_markers['logfoldchanges'].notna()]
        markers_name_dict = {
            'group':'cluster',
            'names':'gene',
            'logfoldchanges':'avg_log2FC',
            'pvals':'p_val',
            'pvals_adj':'p_val_adj',
            'pct_nz_group':'pct.1',
            'pct_nz_reference':'pct.2'
         }
        df_markers = df_markers.rename(markers_name_dict,axis='columns')
        df_markers['cluster'] = df_markers['cluster'].map(lambda x : int(x)+1)
        df_markers = df_markers.loc[df_markers['p_val_adj'] < PVAL_CUTOFF, ]
        df_markers.to_csv(self.df_marker_raw_file, index=None, sep='\t')

        df_markers_filter = df_markers.loc[df_markers['avg_log2FC'] > 0].sort_values('p_val_adj').groupby('cluster').head(100)
        df_markers_filter = df_markers_filter.round({
            'avg_log2FC':3,
            'pct.1':3,
            'pct.2':3,
        })
        df_markers_filter.to_csv(self.df_marker_file, index=None, sep='\t')


    @utils.add_log
    def write_tsne(self):
        df_tsne = self.adata.obsm.to_df()[['X_tsne1','X_tsne2']]
        df_tsne['cluster']=self.adata.obs.cluster
        df_tsne['Gene_Counts']=self.adata.obs.n_genes_by_counts
        tsne_name_dict={'X_tsne1':'tSNE_1', 'X_tsne2':'tSNE_2'}
        df_tsne = df_tsne.rename(tsne_name_dict,axis='columns')
        df_tsne['cluster'] = df_tsne['cluster'].map(lambda x:int(x)+1)
        df_tsne.to_csv(self.df_tsne_file, sep='\t')

    @utils.add_log
    def write_h5ad(self):
        self.adata.write(self.h5ad_file)
    

    @utils.add_log
    def run(self):

        self.calculate_qc_metrics()
        self.write_mito_stats()
        self.normalize()
        self.hvg()
        self.scale()
        self.pca()
        self.neighbors()
        self.tsne()
        self.umap()
        self.leiden()
        self.find_marker_genes()

        self.write_markers()
        self.write_tsne()
        self.write_h5ad()


    def get_df(self):
        """
        return df_tsne, df_marker
        """
        df_tsne = read_tsne(self.df_tsne_file)
        df_marker = pd.read_csv(self.df_marker_file, sep="\t")
        df_marker = format_df_marker(df_marker)
        return df_tsne, df_marker


def get_opts_analysis_match(parser, sub_program):
    """
    Do not perform analysis. Only read data from scRNA-seq match_dir.
    """
    if sub_program:
        parser.add_argument("--match_dir", help=HELP_DICT['match_dir'])
        parser.add_argument("--tsne_file", help=HELP_DICT['tsne_file'])
        parser.add_argument("--df_marker_file", help=HELP_DICT['df_marker_file'])

        parser = s_common(parser)

class Report_runner(Step):

    def __init__(self, args, display_title=None):
    
        super().__init__(args, display_title=display_title)

    def add_marker_help(self):
        self.add_help_content(
            name='Marker Genes by Cluster',
            content='differential expression analysis based on the non-parameteric Wilcoxon rank sum test'
        )
        self.add_help_content(
            name='avg_log2FC',
            content='log fold-change of the average expression between the cluster and the rest of the sample'
        )
        self.add_help_content(
            name='pct.1',
            content='The percentage of cells where the gene is detected in the cluster'
        )
        self.add_help_content(
            name='pct.2',
            content='The percentage of cells where the gene is detected in the rest of the sample'
        )
        self.add_help_content(
            name='p_val_adj',
            content='Adjusted p-value, based on bonferroni correction using all genes in the dataset'
        )

    @staticmethod
    def get_df_file(match_dir):
        """
        return df_tsne_file, df_marker_file
        """
        match_dict = utils.parse_match_dir(match_dir)
        df_tsne_file = match_dict['tsne_coord']
        df_marker_file = match_dict.get('markers', None)
        return df_tsne_file, df_marker_file

    def get_df(self):
        """
        return df_tsne, df_marker
        """
        if utils.check_arg_not_none(self.args, 'match_dir'):
            df_tsne_file, df_marker_file = self.get_df_file(self.args.match_dir)
        elif utils.check_arg_not_none(self.args, 'tsne_file'):
            df_tsne_file = self.args.tsne_file
            df_marker_file = self.args.df_marker_file
        else:
            raise ValueError('match_dir or tsne_file must be specified')
        df_tsne = read_tsne(df_tsne_file)
        if df_marker_file:
            df_marker = pd.read_csv(df_marker_file, sep="\t")
            df_marker = format_df_marker(df_marker)
        else:
            df_marker = None
        return df_tsne, df_marker

    def run(self):
        pass
