""" Cluster analysis module """
import numpy as np
import threading
import concurrent.futures

from typing import Union
from operator import attrgetter
from functools import partial


from AnyQt.QtCore import (
    Qt, QAbstractListModel, QVariant, Slot, QThread
)
from Orange.widgets.utils.concurrent import ThreadExecutor, FutureWatcher, methodinvoke
from Orange.widgets.gui import ProgressBar

from orangecontrib.bioinformatics.geneset import GeneSet
from orangecontrib.bioinformatics.widgets.utils.gui import gene_scoring_method
from orangecontrib.bioinformatics.utils.statistics import FDR, ALT_GREATER
from orangecontrib.bioinformatics.ncbi.gene import Gene

DISPLAY_GENE_COUNT = 20
DISPLAY_GENE_SETS_COUNT = 5


class ClusterGene(Gene):
    __slots__ = ['score', 'p_val', 'fdr']

    def __init__(self, gene_name, gene_id):
        super().__init__()
        self.input_name = gene_name
        self.ncbi_id = gene_id

        # default values
        self.fdr = 1
        self.p_val = 1
        self.score = 0


class ClusterGeneSet(GeneSet):
    __slots__ = ['p_val', 'fdr', 'count']

    def __init__(self):
        super().__init__()


class Cluster:

    CLUSTER_VS_REST = False
    CLUSTER_VS_CLUSTER = True

    def __init__(self, name, index):
        # type: (str, int) -> None
        self.name = name
        self.index = index

        # gene enrichment
        self.genes = []
        self.filtered_genes = []

        # set enrichment
        self.gene_sets = []
        self.filtered_gene_sets = []

        # Statistical method used when performing analysis
        self.method_used = None

    def set_genes(self, gene_names, gene_ids):
        self.genes = []

        for name, entrez in zip(gene_names, gene_ids):
            self.genes.append(ClusterGene(name, entrez))

    @staticmethod
    def apply_filter(obj, p_val, fdr):
        filter_status = []

        if p_val is not None:
            filter_status.append(obj.p_val < p_val)

        if fdr is not None:
            filter_status.append(obj.fdr < fdr)

        return all(filter_status)

    def filter_enriched_genes(self, p_val, fdr, max_gene_count=None):
        all_args_none = all(arg is None for arg in [p_val, fdr])
        filter_function = partial(self.apply_filter, p_val=p_val, fdr=fdr)
        sorted_list = sorted(self.genes, key=attrgetter('p_val' if all_args_none else 'fdr'))
        filtered_list = list(filter(filter_function, sorted_list))

        if max_gene_count is not None:
            self.filtered_genes = filtered_list[:max_gene_count]
        else:
            self.filtered_genes = filtered_list

    def filter_gene_sets(self, p_val, fdr, max_set_count=None):
        all_args_none = all(arg is None for arg in [p_val, fdr])
        filter_function = partial(self.apply_filter, p_val=p_val, fdr=fdr)
        sorted_list = sorted(self.gene_sets, key=attrgetter('p_val' if all_args_none else 'fdr'))
        filtered_list = list(filter(filter_function, sorted_list))

        if max_set_count is not None:
            self.filtered_gene_sets = filtered_list[:max_set_count]
        else:
            self.filtered_gene_sets = filtered_list

    def gene_set_enrichment(self, gene_sets, selected_sets, genes, ref_genes):
        self.gene_sets = []

        if not genes:
            return

        # calculate gene set enrichment
        for gene_set in gene_sets:

            if gene_set.hierarchy not in selected_sets:
                continue

            gs = ClusterGeneSet()
            enrichment_result = gene_set.set_enrichment(ref_genes, genes.intersection(genes))
            gs.count = len(enrichment_result.query)
            gs.p_val = enrichment_result.p_value
            gs.name = gene_set.name
            gs.gs_id = gene_set.gs_id
            self.gene_sets.append(gs)

        # calculate FDR
        p_vals = [gs.p_val for gs in self.gene_sets]
        fdr_values = FDR(p_vals)

        for gs, fdr in zip(self.gene_sets, fdr_values):
            gs.fdr = fdr

    def __update_gene_objects(self, scores, p_vals, fdr_vals):
        # type: (Union[np.ndarray, list], Union[np.ndarray, list], Union[np.ndarray, list]) ->  None
        """ update gene objects with computed results

        :param p_vals:   Computed p-values
        :param fdr_vals: Computed fdr-values

        """
        for score, p_val, fdr, gene in zip(scores, p_vals, fdr_vals, self.genes):
            gene.score = score
            gene.p_val = p_val
            gene.fdr = fdr

    def cluster_scores(self, table_x, rows_by_cluster, method, design, **kwargs):
        # type: (np.ndarray, np.ndarray, gene_scoring_method, str) -> None
        """
        General scoring of genes in the cluster.
        Ways to score genes are determined by design.
        If a batch variable index is defined, it is accounted for in gene scoring.

        :param table_x:
        :param rows_by_cluster:
        :param method:
        :param design:
        :param kwargs:
        :param aggregation:
        :param rows_by_batch:
        :return:
        """
        aggregation = kwargs.get('aggregation', 'max')
        alternative = kwargs.get('alternative', ALT_GREATER)
        rows_by_batch = kwargs.get('rows_by_batch', None)
        if not isinstance(rows_by_batch, np.ndarray):
            rows_by_batch = np.zeros((len(table_x),))
        uniq_batches = set(rows_by_batch)

        # Determine clusters
        self.method_used = method.name
        uniq_clusters = set(rows_by_cluster) - {self.index}
        this_cluster = self.index
        if design == self.CLUSTER_VS_REST:
            rows_by_cluster = rows_by_cluster == self.index
            uniq_clusters = {False}
            this_cluster = True

        calculated_p_values = np.ones((table_x.shape[1],      # genes
                                      len(uniq_clusters),     # other clusters
                                      len(uniq_batches)))     # batches

        calculated_scores = np.ones((table_x.shape[1],       # genes
                                     len(uniq_clusters),     # other clusters
                                     len(uniq_batches)))     # batches

        for bi, b in enumerate(uniq_batches):
            for ci, c in enumerate(uniq_clusters):
                cluster = table_x[np.logical_and(rows_by_cluster == this_cluster, rows_by_batch == b)]
                rest = table_x[np.logical_and(rows_by_cluster == c, rows_by_batch == b)]
                if cluster.any() and rest.any():
                    scores, p_values = method.score_function(cluster, rest, alternative=alternative)
                    scores[np.isnan(p_values)] = 0
                    calculated_scores[:, ci, bi] = scores
                    p_values[np.isnan(p_values)] = 1
                    calculated_p_values[:, ci, bi] = p_values

        if aggregation == 'max':
            max_p_values = np.max(calculated_p_values, axis=(1, 2))
            max_p_indexes = np.where(np.max(calculated_p_values, axis=(1, 2), keepdims=True) == calculated_p_values)
            # this holds true only if max_p_indexes.ndim == 3
            scores = calculated_scores[max_p_indexes[0], max_p_indexes[1], max_p_indexes[2]]
            fdr_values = FDR(max_p_values.tolist())
            self.__update_gene_objects(scores, max_p_values, fdr_values)
        else:
            raise NotImplementedError("Aggregation %s is not implemented" % aggregation)
        return

    def to_html(self):
        gene_sets = '(no enriched gene sets)'
        if self.filtered_gene_sets:
            sets_to_display = self.filtered_gene_sets[:DISPLAY_GENE_SETS_COUNT]
            gene_sets = '<br>'.join(['<b>{}</b> (FDR={:0.2e}, n={})'.format(g_set.name, g_set.fdr, g_set.count)
                                    for g_set in sets_to_display])

            if len(self.filtered_gene_sets) > len(sets_to_display):
                gene_sets += \
                    '<br> ... ({} more gene sets)'.format(len(self.filtered_gene_sets) - DISPLAY_GENE_SETS_COUNT)

        genes = '(all genes are filtered out)'
        if self.filtered_genes:
            genes_to_display = [gene.input_name for gene in self.filtered_genes[:DISPLAY_GENE_COUNT]]

            genes = ', '.join(genes_to_display)
            if len(self.filtered_genes) > len(genes_to_display):
                genes += ', ... ({} more genes)'.format(len(self.filtered_genes) - DISPLAY_GENE_COUNT)

        html_string = """
        <html>
        <body>
        <table >
            <tr>
            </tr>
            <tr>
                <td width=\"20%\" heigth=\"100%\" align=left valign=top><p>{}</p></td>
                <td width=\"80%\" heigth=\"100%\" align=left valign=top>
                    {}
                </td>
            </tr>
            <tr> </tr>
            <tr>
                <td width=\"20%\" ></td>
                <td width=\"80%\" align=left valign=top>{}</td>
            </tr>
            <tr>
            </tr>
        </table>
        </body>
        </html>
        """.format(self.name, gene_sets, genes)
        return html_string


class Task:
    future = None
    watcher = None
    cancelled = False

    def cancel(self):
        self.cancelled = True
        self.future.cancel()
        concurrent.futures.wait([self.future])


class ClusterModel(QAbstractListModel):

    def __init__(self, parent=None):
        QAbstractListModel.__init__(self)
        self.__items = []
        self.parent = parent

        self._task = None  # type: Union[Task, None]
        self._executor = ThreadExecutor()

    def add_rows(self, rows):
        self.__items = rows

    def get_rows(self):
        return self.__items

    def rowCount(self, *args, **kwargs):
        return len(self.__items)

    def data(self, model_index, role=None):
        # check if data is set
        if not self.__items:
            return QVariant()

        # return empty QVariant if model index is unknown
        if not model_index.isValid() or not (0 <= model_index.row() < len(self.__items)):
            return QVariant()

        row_obj = self.__items[model_index.row()]

        if role == Qt.DisplayRole:
            return row_obj

    @Slot(concurrent.futures.Future)
    def _end_task(self, f):
        assert self.thread() is QThread.currentThread()
        assert threading.current_thread() == threading.main_thread()
        assert self._task is not None
        assert self._task.future is f
        assert f.done()

        self._task = None
        self.parent.progressBarFinished()
        self.parent.filter_genes()

        try:
            f.result()
        except Exception as ex:
            raise ex

    def _score_genes(self, callback, **kwargs):
        for item in self.get_rows():
            item.cluster_scores(**kwargs)
            callback()

    @Slot(bool)
    def progress_advance(self, finish):
        # GUI should be updated in main thread. That's why wex are calling advance method here
        if self.parent.progress_bar:
            if finish:
                self.parent.progressBarFinished()
            else:
                self.parent.progress_bar.advance()

    def cancel(self):
        """
        Cancel the current task (if any).
        """
        if self._task is not None:
            self._task.cancel()
            assert self._task.future.done()
            # disconnect the `_task_finished` slot
            self._task.watcher.done.disconnect(self._end_task)
            self._task = None

    def score_genes(self, **kwargs):
        """ Run gene enrichment.

        :param design:
        :param data_x:
        :param rows_by_cluster:
        :param method:


        Note:
            We do not apply filter nor notify view that data is changed. This is done after filters
        """
        if self._task is not None:
            # First make sure any pending tasks are cancelled.
            self.cancel()
        assert self._task is None

        progress_advance = methodinvoke(self, "progress_advance", (bool,))

        def callback():
            if self._task.cancelled:
                raise KeyboardInterrupt()
            progress_advance(self._task.cancelled)

        self.parent.progress_bar = ProgressBar(self.parent, iterations=len(self.get_rows()))
        f = partial(self._score_genes, callback=callback, **kwargs)
        self._task = Task()
        self._task.future = self._executor.submit(f)

        self._task.watcher = FutureWatcher(self._task.future)
        self._task.watcher.done.connect(self._end_task)

    def gene_sets_enrichment(self, gs_object, gene_sets, reference_genes):
        """ Run gene sets enrichment.

        :param gs_object:
        :param gene_sets:
        :param reference_genes:

        Note:
            We do not apply filter nor notify view that data is changed. This is done after filters

        """

        for item in self.get_rows():
            genes = [gene.ncbi_id for gene in item.filtered_genes]
            item.gene_set_enrichment(gs_object,
                                     gene_sets,
                                     set(genes),
                                     reference_genes)

    def apply_gene_filters(self, p_val=None, fdr=None, count=None):
        [item.filter_enriched_genes(p_val, fdr, max_gene_count=count) for item in self.get_rows()]
        self.dataChanged.emit(self.createIndex(0, 0), self.createIndex(self.rowCount(0), 0))

    def apply_gene_sets_filters(self, p_val=None, fdr=None, count=None):
        [item.filter_gene_sets(p_val, fdr, max_set_count=count) for item in self.get_rows()]
        self.dataChanged.emit(self.createIndex(0, 0), self.createIndex(self.rowCount(0), 0))




