import time
from collections import defaultdict
from operator import itemgetter

import numba
import numpy as np
import pandas as pd
import scipy
from cusingler import cusingler
from joblib import (
    Parallel,
    delayed,
    cpu_count
)
from tqdm import tqdm

from stereo.core.stereo_exp_data import StereoExpData
from stereo.log_manager import logger
from .utils import (
    corr_spearman,
    apply_along_axis
)
from ..algorithm_base import AlgorithmBase


class _TestData(object):

    def __init__(self, exp_matrix, cell_names, gene_names):
        self.exp_matrix = exp_matrix
        self.cell_names = cell_names
        self.gene_names = gene_names

    def to_df(self):
        return pd.DataFrame(self.exp_matrix.toarray(), index=self.cell_names, columns=self.gene_names)


@numba.jit(forceobj=True, nogil=True, parallel=True)
def _spearman_parallel(label, ranked_mat_ref, y, quantile):
    ranked_mat_qry = apply_along_axis(y)
    sim = corr_spearman(ranked_mat_ref, ranked_mat_qry)
    return label, np.percentile(sim, quantile, axis=1)


def _get_label_genes_set(top_labels, trained_data, de_n_len=0):
    de_n = np.round(500 * (2 / 3) ** np.log2(de_n_len if de_n_len else len(top_labels)))
    genes_filtered = []
    for i in top_labels:
        for j in top_labels:
            if i == j:
                continue
            s = trained_data[i][j]
            s = s.loc[s[0] > 0]
            genes_filtered.append(s.index[:int(de_n)])
    return genes_filtered


class SingleR(AlgorithmBase):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.n_jobs = None
        self.quantile = None
        self.ref_exp_data = None
        self.fine_tune_times = None
        self.group_data_frame = None
        self.fine_tune_threshold = None

    def main(
            self,
            ref_exp_data: StereoExpData,
            ref_use_col: str = "celltype",
            cluster_res_key: str = None,
            quantile: int = 80,
            fine_tune_threshold: float = 0.05,
            fine_tune_times: int = 0,
            n_jobs: int = int(cpu_count() / 2),
            res_key: str = 'annotation',
            method: str = 'default',
            gpuid: int = 0
    ):
        """
        Single-cell recognition is a tool to automatically annotate a test sample by a reference sample.

        :param ref_exp_data: a `StereoExpData` as reference data.
        :param ref_use_col: `ref_use_col` mean cluster-like or annotation result key in ref's `StereoExpData.tl.result`.
        :param cluster_res_key: test's cluster-like key in `StereoExpData.tl.result`, If cluster_res_key is set, CPU will be used by default regardless of whether gpu acceleration is set or not.
        :param quantile: quantile will influence scoring and fine_tune result.
        :param fine_tune_threshold: while in fine_tuning, if result greater than `max(result) - fine_tune_threshold`, will be filtered.
        :param fine_tune_times: default to 0, meaning that it will fine_tune until results decreasing to only 1.
                                If it is set to num(eg: 5), it will only loop only 5 times, and choose the first one.
        :param n_jobs: `joblib` parameter, will create `n_jobs` num of threads to work.
        :param res_key: default to `annotation`, means the result will be stored as key `annotation` in the `tl.result`.
        :param methods：whether to use GPU acceleration, if methods is `rapids`, it means using, It is not used by default.
        :param gpuid：slots used by gpu, default to 0.

        :return: `pandas.DataFrame`
        """  # noqa
        interact_genes = pd.Index(self.stereo_exp_data.gene_names) & pd.Index(ref_exp_data.gene_names)
        assert not interact_genes.empty, "no gene of `test_exp_data.gene_names` in `ref_exp_data.gene_names`"

        total_start_time = time.time()
        test_exp_data = self.stereo_exp_data.sub_by_name(gene_name=interact_genes)
        ref_exp_data = ref_exp_data.sub_by_name(gene_name=interact_genes)
        if not cluster_res_key and method == 'rapids':
            self.check_gpu_env()

            qry_data = test_exp_data.exp_matrix.data.tolist()
            qry_indices = test_exp_data.exp_matrix.indices.tolist()
            qry_indptr = test_exp_data.exp_matrix.indptr.tolist()
            ref_geneidx = ref_exp_data.gene_names.tolist()
            temp_ref_exp_data = ref_exp_data.array2sparse()
            ref_data = temp_ref_exp_data.data.tolist()
            ref_indices = temp_ref_exp_data.indices.tolist()
            ref_indptr = temp_ref_exp_data.indptr.tolist()
            ref_height, ref_width = temp_ref_exp_data.shape
            qry_height, qry_width = test_exp_data.shape
            cellnames = test_exp_data.cell_names.tolist()
            qry_geneidx = test_exp_data.gene_names.tolist()
            codes, celltypes = self.load_cell_types(ref_exp_data, ref_use_col)
            res = np.array(
                cusingler.run(
                    n_jobs, gpuid, ref_height, ref_width, ref_data, ref_indices, ref_indptr,
                    qry_height, qry_width, qry_data, qry_indices, qry_indptr,
                    codes, celltypes, cellnames, ref_geneidx, qry_geneidx,
                    quantile=quantile, finetune_thre=fine_tune_threshold, finetune_times=fine_tune_times
                )
            )
            bins = res[:, 0].tolist()
            labels_array = res[:, 1].tolist()
            ret_labels = res[:, 2].tolist()
        else:
            assert ref_use_col in ref_exp_data.tl.result
            self.group_data_frame = ref_exp_data.tl.result[ref_use_col]
            self.ref_gene_names = ref_exp_data.to_df().columns
            self.test_gene_names = {name: i for i, name in enumerate(test_exp_data.gene_names.tolist())}
            self.group_data_frame = self.group_data_frame.reset_index()
            self.group_data_frame['bins'] = self.group_data_frame.index
            self._group_data_frame_checker()

            self.ref_exp_data = _TestData(
                ref_exp_data.exp_matrix if scipy.sparse.issparse(ref_exp_data.exp_matrix) else scipy.sparse.csr_matrix(
                    ref_exp_data.exp_matrix),
                np.array(range(len(ref_exp_data.cells.cell_name))),
                np.array(range(len(ref_exp_data.genes.gene_name)))
            )

            self.n_jobs = n_jobs
            self.quantile = quantile
            self.fine_tune_times = fine_tune_times
            self.fine_tune_threshold = fine_tune_threshold

            logger.info(f'start single-r with n_jobs={n_jobs} fine_tune_times={fine_tune_times}')

            logger.info('start training ref...')
            start_time = time.time()
            trained_data, common_gene = self._train_ref()
            logger.info(f'training ref finished, cost {time.time() - start_time} seconds')

            if cluster_res_key:
                tmp_exp_matrix = pd.DataFrame(
                    test_exp_data.exp_matrix.todense() if scipy.sparse.issparse(
                        test_exp_data.exp_matrix) else test_exp_data.exp_matrix,
                    index=test_exp_data.cell_names,
                    columns=test_exp_data.gene_names
                )
                test_cluster_result = test_exp_data.tl.result[cluster_res_key]
                test_group_data_frame = test_cluster_result['group']
                if 'bins' in test_cluster_result:
                    test_group_data_frame.index = test_cluster_result['bins'].values
                tmp_exp_matrix = tmp_exp_matrix.groupby(test_group_data_frame).sum()
                test_data = _TestData(
                    scipy.sparse.csr_matrix(tmp_exp_matrix.values),
                    tmp_exp_matrix.index,
                    np.array(range(len(test_exp_data.gene_names)))
                )
            else:
                test_data = _TestData(
                    test_exp_data.exp_matrix,
                    np.array(range(len(test_exp_data.cell_names))),
                    np.array(range(len(test_exp_data.gene_names)))
                )

            logger.info('start scoring test_data...')
            start_time = time.time()
            output, labels_array = self._score_test_data(test_data, common_gene)
            logger.info(f'scoring test_data finished, cost {time.time() - start_time} seconds')

            logger.info('start fine-tuning...')
            start_time = time.time()
            ret_labels = self._fine_tune(test_data, output, trained_data)
            logger.info(f'fine-tuning finished, cost {time.time() - start_time} seconds')
            bins = test_data.cell_names if cluster_res_key else test_exp_data.cell_names

        res = pd.DataFrame(columns=['bins', 'group', 'first_labels'])
        res['bins'] = bins
        res['group'] = ret_labels
        res['first_labels'] = labels_array
        res.index = res['bins'].values
        logger.info(f'single-r finished, cost {time.time() - total_start_time} seconds')
        if cluster_res_key:
            res_single_r = pd.DataFrame(res.loc[test_cluster_result['group']])
            bins = test_cluster_result['bins'] if 'bins' in test_cluster_result else test_cluster_result.index
            res_single_r['bins'].astype(bins.dtype)
            res_single_r['bins'] = bins.values
            res_single_r.index = res_single_r['bins']
            self.pipeline_res[res_key] = res_single_r
            return res_single_r
        else:
            self.pipeline_res[res_key] = res
            return res

    @staticmethod
    def check_gpu_env():
        try:
            import cudf  # noqa
            import cugraph  # noqa
        except ImportError:
            raise ImportError(
                "Your env don't have GPU related RAPIDS packages, if you want to run this option, follow the "
                "guide at https://stereopy.readthedocs.io/en/latest/Tutorials/clustering_by_gpu.html")

    @staticmethod
    def load_cell_types(ref_exp_data, ref_use_col):
        if ref_use_col not in ref_exp_data._ann_data.obs:
            raise IndexError("No celltype found.")

        group = ref_exp_data._ann_data.obs[ref_use_col]
        codes = group.cat.categories.to_list()
        celltypes = [int(celltype) for celltype in group.cat.codes]
        return codes, celltypes

    def _train_ref(self):
        dict_of_median_exp = dict()
        for label, y in self.group_data_frame.groupby('group'):
            cells_bool_list = np.isin(self.ref_exp_data.cell_names, y['bins'].values)
            dict_of_median_exp[label] = pd.DataFrame(
                np.median(self.ref_exp_data.exp_matrix[cells_bool_list].toarray(), axis=0))
        median_exp = pd.concat(dict_of_median_exp, axis=1)
        median_exp.index = self.ref_gene_names

        ret_all = defaultdict(dict)
        ret_gene = defaultdict(dict)

        ct = self.group_data_frame['group'].astype('category').cat.categories
        de_n = np.round(500 * (2 / 3) ** np.log2(len(ct)))
        for i in ct:
            for j in ct:
                if i == j:
                    continue

                ret_all[i][j] = median_exp[i] - median_exp[j]
                ret_tmp = pd.DataFrame(ret_all[i][j]).round(6)
                ret_tmp['arg'] = [x for x in range(ret_tmp.shape[0])][::-1]
                ret_all[i][j] = ret_tmp.sort_values([0, 'arg'], ascending=False)
                ret_gene[i][j] = ret_all[i][j].index[:int(de_n)]

        return ret_all, np.unique([y for x in ret_gene.values() for y in x.values()])

    def _score_test_data(self, test_data, common_gene):
        original_exp = {}
        tmp_gene_names = self.ref_gene_names.tolist()
        test_data_length = test_data.exp_matrix.shape[1]
        common_gene_index = []
        for com in common_gene:
            index = tmp_gene_names.index(com)
            if index < test_data_length:
                common_gene_index.append(index)

        for x, y in self.group_data_frame.groupby('group'):
            y_cell_bool_list = np.isin(self.ref_exp_data.cell_names, y['bins'].values)
            original_exp[x] = self.ref_exp_data.exp_matrix[y_cell_bool_list][:, common_gene_index].toarray()
        test_mat = test_data.exp_matrix[:, common_gene_index].toarray()
        ranked_mat_ref = apply_along_axis(test_mat.T)

        res_dict = dict(Parallel(n_jobs=self.n_jobs, backend="threading")(
            delayed(_spearman_parallel)(label, ranked_mat_ref, y.T, self.quantile)
            for label, y in tqdm(original_exp.items())
        ))

        ret = pd.DataFrame(res_dict, index=test_data.cell_names)
        return ret, ret.columns[np.argmax(ret.values, axis=1)]

    def _fine_tune(self, test_data, output, trained_data):
        tmp = output.values.copy()
        tmp[tmp < np.array(np.max(output.values, axis=1) - self.fine_tune_threshold).reshape([-1, 1])] = 0
        tmp[tmp > 0] = 1

        ref = pd.DataFrame(
            self.ref_exp_data.exp_matrix.toarray(),
            index=self.ref_exp_data.cell_names,
            columns=self.ref_exp_data.gene_names
        )

        logger.info(f'fine-tuning with test_data(shape={test_data.exp_matrix.shape})')

        ret_labels = Parallel(n_jobs=self.n_jobs, backend="threading")(
            delayed(self._fine_tune_parallel)(
                ref,
                output.columns[tmp[x].astype(bool)].values,
                y[1].to_frame().T,
                trained_data
            )
            for x, y in tqdm(enumerate(test_data.to_df().iterrows()))
        )
        return ret_labels

    def _fine_tune_parallel(self, ref, labels, y, trained_data):
        if self.fine_tune_times:
            try_num = 0
            while try_num < self.fine_tune_times and len(labels) > 1:
                labels = self._fine_tune_one_time(labels, ref, y, trained_data, try_num)
                try_num += 1
        else:
            while len(labels) > 1:
                labels = self._fine_tune_one_time(labels, ref, y, trained_data)
        return labels[0]

    @numba.jit(forceobj=True, nogil=True, parallel=True)
    def _fine_tune_one_time(self, top_labels, ref, test, trained_data, try_num=None):
        genes_filtered = _get_label_genes_set(top_labels, trained_data)
        genes_filtered = np.unique([y for x in genes_filtered for y in x])
        if len(genes_filtered) < 20:
            return [top_labels[0]]

        test_data_length = test.shape[1]
        ref_data_length = ref.shape[1]
        test_genes_filtered_index = []
        ref_genes_filtered_index = []
        for gen in genes_filtered:
            index = self.test_gene_names.get(gen, None)
            if index is None:
                continue
            if test_data_length > index:
                test_genes_filtered_index.append(index)
            if ref_data_length > index:
                ref_genes_filtered_index.append(index)
        test_filtered = test.loc[:, test_genes_filtered_index].values
        if np.std(test_filtered) <= 0:
            return [top_labels[0]]

        bool_list = np.isin(self.group_data_frame['group'].values, top_labels)
        ref_filtered = ref.loc[bool_list, ref_genes_filtered_index]
        group_data_frame_filtered = self.group_data_frame.loc[bool_list].reset_index()

        ranked_mat_ref = apply_along_axis(test_filtered.T)
        res_labels = {}
        for p, q in group_data_frame_filtered.groupby('group'):
            ref_filtered_by_group = ref_filtered.iloc[q.index].values.T
            ranked_mat_qry = apply_along_axis(ref_filtered_by_group)
            sim = corr_spearman(ranked_mat_ref, ranked_mat_qry)
            if not (sim.shape and sim.shape[0] and sim.shape[1]):
                continue
            res_labels[p] = np.percentile(sim, self.quantile, axis=1)[0]

        if not res_labels:
            return [top_labels[0]]

        # remove the labels lower than median
        mid_value = (min(res_labels.values()) + np.median(list(res_labels.values()))) / 2
        res_labels = {p: value for p, value in res_labels.items() if value > mid_value}
        if not res_labels:
            return [top_labels[0]]

        # remove the label lower than fine_tune_threshold
        max_value = max(res_labels.values()) - self.fine_tune_threshold
        res_labels = {p: value for p, value in res_labels.items() if value >= max_value}
        if not res_labels:
            return [top_labels[0]]

        if try_num and try_num == self.fine_tune_times - 1:
            return [sorted(res_labels.items(), key=itemgetter(1))[-1][0]]

        return list(res_labels.keys())

    @staticmethod
    def test_rank():
        mat = np.random.sample((5, 5)).astype(np.float32)
        res_self = apply_along_axis(mat)

        from scipy.stats import rankdata
        res_sci = rankdata(mat, axis=0).astype(np.float32)

        assert np.all(res_self == res_sci), 'self define rank is different with the SciPy\'s versions'
        logger.info('test_rank success')

    def _group_data_frame_checker(self):
        assert 'group' in self.group_data_frame.columns
        if 'bins' not in self.group_data_frame.columns:
            self.group_data_frame = self.group_data_frame.reset_index().rename(
                columns={'index': 'bins', 'CellID': 'bins'}
            )
        assert 'bins' in self.group_data_frame.columns

    @staticmethod
    def test_group_data_frame():
        s1 = SingleR(None)
        # cluster result data like
        s1.group_data_frame = pd.DataFrame(index=['cell1', 'cell2'], columns=['bins', 'group'])
        s1._group_data_frame_checker()

        s2 = SingleR(None)
        # h5ad with cell_type data like
        s2.group_data_frame = pd.DataFrame(index=['cell1', 'cell2'], columns=['CellID', 'group'])
        s2._group_data_frame_checker()

        s3 = SingleR(None)
        # h5ad with cell_type data like
        s3.group_data_frame = pd.DataFrame(index=['cell1', 'cell2'], columns=['group'])
        s3._group_data_frame_checker()
        logger.info('test_group_data_frame success')
