from typing import List
from ruleskit import RuleSet, RegressionRule, ClassificationRule
import logging
from collections import Counter

logger = logging.getLogger(__name__)

# TODO (pcotte) : make it a watcher for data separation


class Aggregation:
    """Abstract class for aggregating node models."""

    # noinspection PyUnresolvedReferences
    def __init__(self, central_server: "CentralServer", **kwargs):
        """
        Parameters
        ----------
        central_server: CentralServer
            `ifra.central_server.CentralServer` object
        kwargs:
            Any additionnal keyword argument that the overleading class accepts. Those arguments will become attributes.
        """
        self.central_server = central_server
        for arg in kwargs:
            setattr(self, arg, kwargs[arg])

    def aggregate(self, rulesets: List[RuleSet]) -> str:
        """Aggregates the node models into central model, then updates the central model's rules predictions

        Parameters
        ----------
        rulesets: List[RuleSet]
            New rulesets provided by the nodes.

        Returns
        -------
        str\n
          * "updated" if central model was updates\n
          * "pass" if no new rules were found but learning should continue\n
          * "stop" otherwise. Will stop learning.\n"""
        result = self._aggregate(rulesets)
        self.aggregate_predictions()
        return result

    def _aggregate(self, rulesets: List[RuleSet]) -> str:
        """To be implemented in daughter class.

         Must modify inplace `ifra.central_server.CentralServer` *ruleset*.

        Parameters
        ----------
        rulesets: List[RuleSet]
            New rulesets provided by the nodes.

        Returns
        -------
        str\n
          * "updated" if central model was updates\n
          * "pass" if no new rules were found but learning should continue\n
          * "stop" otherwise. Will stop learning.\n
        """
        pass

    def aggregate_predictions(self) -> None:
        """To be implemented in daughter class.

         Must modify inplace the rules predictions in `ifra.central_server.CentralServer` *ruleset*
        """
        pass


class AdaBoostAggregation(Aggregation):
    """Among all rules generated by the nodes in the current iteration, will keep only the most recurent rule(s).

    Can be used by giving *adaboost_aggregation* as *aggregation* argument when creating a
    `ifra.central_server.CentralServer`

    Those rules will be added to  the given ruleset, returned by the function,
    and the points in each node's data activated by it will be ignored in order to find other relevant rules.
    In order not to remove too many points, a maximum of coverage is imposed.

    If among the rules extracted from the nodes, none are new compared to the current state of the model, nodes are
    not updated. If no new rules are found but each nodes had new data, learning is finished.
    """

    def _aggregate(self, rulesets: List[RuleSet]) -> str:
        logger.info("Aggregating fit results using AdaBoost method...")
        all_rules = []
        n_nodes = len(self.central_server.nodes)

        for rs in rulesets:
            all_rules += rs.rules

        if len(all_rules) == 0:
            logger.info("... no rules found")
            if len(rulesets) == n_nodes:
                logger.info("No new rules were found, despite all nodes having provided a new model : learning is over")
                return "stop"
            return "pass"

        if self.central_server.ruleset is None:
            occurences = {r: all_rules.count(r)for r in set(all_rules)}
        else:
            occurences = {
                r: all_rules.count(r)
                for r in set(all_rules)
                if r not in self.central_server.ruleset
            }

        if len(occurences) == 0:
            logger.warning("No new rules matched coverage criterion")
            if len(rulesets) == n_nodes:
                logger.info("No new rules were found, despite all nodes having provided a new model : learning is over")
                return "stop"
            return "pass"
        max_occurences = max(list(occurences.values()))

        if self.central_server.ruleset is None:
            new_rules = RuleSet(remember_activation=False)
        else:
            new_rules = RuleSet(rules_list=self.central_server.ruleset.rules, remember_activation=False)

        found_new = 0
        for r in occurences:
            if r not in new_rules and occurences[r] == max_occurences:
                found_new += 1
                new_rules.append(r, update_activation=False)

        if found_new == 0:
            logger.warning("No new rules found")
            if len(rulesets) == n_nodes:
                logger.info(
                    "No new rules were found, despite all nodes having provided a new model" " : learning is over"
                )
                return "stop"
            return "pass"
        else:
            logger.info(f"Aggregated {found_new} new rules")

        # new_rules.check_duplicated_rules(
        #     new_rules.rules, name_or_index="name" if len(new_rules.features_names) > 0 else "index"
        # )

        del self.central_server.ruleset
        self.central_server.ruleset = new_rules
        return "updated"

    def aggregate_predictions(self):
        if self.central_server.ruleset.rule_type == ClassificationRule:
            aggregate_classif_preds(self.central_server.ruleset)
        elif self.central_server.ruleset.rule_type == RegressionRule:
            aggregate_regr_preds(self.central_server.ruleset)
        else:
            raise TypeError(f"Unexpected ruleset's rule type {self.central_server.ruleset.rule_type}")


def aggregate_classif_preds(ruleset: RuleSet):
    """For now, the ruleset can contain rules with identical conditions but different predictions.
    This method aims at solving that : for duplicated conditions, the corresponding rules will have their
    predictions set to be the one of the most recurring prediction. If several predictions are equally frequent, the
    rules corresponding to the duplicated conditions are discarder.

    Parameters
    ----------
    ruleset: RuleSet
        Modified in place
    """
    duplicated_conditions = {}
    to_remove = []
    conditions = [r.condition for r in ruleset]

    if list(set(conditions)) == conditions:  # No duplicated conds -> no contradictory predictions -> do nothing
        return

    for i in range(len(ruleset)):
        if conditions.count(conditions[i]) > 1:
            if conditions[i] not in duplicated_conditions:
                duplicated_conditions[conditions[i]] = [i]
            else:
                duplicated_conditions[conditions[i]].append(i)

    for condition in duplicated_conditions:
        preds = []
        for i in duplicated_conditions[condition]:
            preds.append(ruleset[i].prediction)
        if len(set(preds)) == 1:  # Duplicated conditions predict the same : nothing to do
            continue
        counter = Counter(preds).most_common(2)
        """Is a list of 2 tuples containing the 2 most frequent elements in *preds* with the number of times they
         are present"""
        if counter[0][1] == counter[1][1]:  # 2 or more different predictions are equally frequent : ignore rules
            to_remove += duplicated_conditions[condition]
            continue
        good_pred = counter[0][0]
        for i in duplicated_conditions[condition]:
            # noinspection PyProtectedMember
            ruleset[i]._prediction = good_pred

    ruleset._rules = [ruleset[i] for i in range(len(ruleset)) if i not in to_remove]
    ruleset.remember_activation = False
    ruleset.stack_activation = False
    ruleset._activation = None
    ruleset.stacked_activations = None
    ruleset._coverage = None


def aggregate_regr_preds(self):
    # TODO (pcotte)
    pass
