"""
This module include all logic for Grammar Guide Genetic Programming including Derivation Trees, Wihgham and Depth
Control Crossovers and mutation
"""

from functools import reduce
import numpy as np
from gplpy.gggp.grammar import Terminal, Variable, NonTerminal, ProbabilisticModel  

class Derivation:
    """Derivation tree"""
    grammar = None
    probabilistic_model=ProbabilisticModel.uniform

    def __init__(self, max_recursions=0, tree=None, crossover_node=None, subtree=None):
        """Derivation tree constructor

           Keyword arguments:
           grammar -- derivation tree grammar
           max_recursions -- maximum recursions applied on the derivation tree (default math.inf)
           tree -- main subtree for offspring derivation tree creation (default None)
           crossover_node -- non terminal symbol where crossover takes place and subtree must be replaced (default None)
           subtree -- subtree to replace in crossover node (default None)
           mutation_rate -- mutation rate to be applied on crossover creation (default 0.02)
           probabilistic model -- probabilistic model type for new derivation trees initialization ( default Probabilistic_Modell.uniform )
        """

        self.max_recursions = max_recursions
        self._depth = None
        self._recursions = None
        self._word = None
        self._str_word = None

        # Create a new derivation tree
        if tree is None:
            remaining_recursions = np.random.random_integers(0, max_recursions) if Derivation.probabilistic_model is ProbabilisticModel.uniform else max_recursions
            self.tree = Tree(derivation=self,
                             remaining_recursions=remaining_recursions,
                             node=Derivation.grammar.axiom)
        # Copy a derivation tree
        else:
            # TODO check root variable
            self.tree = Tree(derivation=self,
                             node=Derivation.grammar.axiom,
                             tree=tree,
                             crossover_node=crossover_node,
                             subtree=subtree)

    @property
    def depth(self):
        """Returns derivation tree depth"""
        if self._depth is None:
            self._depth = self.tree.depth
        return self._depth

    @property
    def recursions(self):
        """Returns number of recursive derivations applied"""
        if self._recursions is None:
            self._recursions = self.tree.recursions
        return self._recursions

    @property
    def word(self):
        """Returns the word generated by the derivation tree"""
        if self._word is None:
            self._word = self.tree.word
        return self._word

    def __str__(self):
        """Returns the word generated by the derivation tree in string format"""
        if self._str_word is None:
            self._str_word = ' '.join(map(str, self.word))
        return self._str_word

    def __iter__(self):
        """Iterator of the elements of the word generated by the derivation tree"""
        return iter(self.word)

    def __len__(self):
        """Length of the word generated by the derivation tree"""
        return len(self.word)


class Tree:
    """Tree structure for derivation tree"""
    def __init__(self, derivation, remaining_recursions=0, max_recursions=0,
                 root=None, node=None, tree=None,
                 mutation_rate=0.0, crossover_node=None, subtree=None):
        """Tree constructor"""
        # Sore derivation information
        self.derivation = derivation
        # Store node information
        self.root = root
        self.leaves = None

        # Information generated during properties method calls
        self._word = None
        self.str_word = None
        self._non_terminal_nodes = None
        self._non_terminal_nodes_by_depth = None
        self._depth = None
        self._recursions = None
        self._height = None

        # If the node is a Terminal or a Critical Terminal there is nothing else to do
        # Variable have their own value, a copy of Critical Terminal is created
        if isinstance(node, Variable):
            self.node = node
            self._word = [self.node]
            return

        # Terminal
        if isinstance(node, Terminal):
            self.node = node
            self._word = [self.node]
            return

        # Axiom & Non terminal
        self.node = node

        # Create a derivation tree using parent's derivation trees, if it doesn't mutate
        if tree:
            self.leaves = []

            if crossover_node:
                # If the root node of the tree is the crossover node we continue copying the subtree
                if tree is crossover_node:
                    tree = subtree
                    self.production = tree.production
                    for leave in tree.leaves:
                        self.leaves.append(Tree(derivation=derivation,
                                                remaining_recursions=remaining_recursions-1 if self.production.left in self.production.right else remaining_recursions,
                                                root=self,
                                                node=leave.node,
                                                tree=leave,
                                                mutation_rate=mutation_rate))
                # Continue copying the main tree
                else:
                    self.production = tree.production
                    for leave in tree.leaves:
                        self.leaves.append(Tree(derivation=derivation,
                                                remaining_recursions=remaining_recursions-1 if self.production.left in self.production.right else remaining_recursions,
                                                root=self,
                                                node=leave.node,
                                                tree=leave,
                                                crossover_node=crossover_node,
                                                subtree=subtree,
                                                mutation_rate=mutation_rate))
            # Continue copying the subtree
            else:
                self.production = tree.production
                for leave in tree.leaves:
                    self.leaves.append(Tree(derivation=derivation,
                                            remaining_recursions=remaining_recursions-1 if self.production.left in self.production.right else remaining_recursions,
                                            root=self,
                                            node=leave.node,
                                            tree=leave,
                                            mutation_rate=mutation_rate))

        # Create a new derivation tree or mutate a tree
        else:
            self.production = self.derivation.grammar.get_production_with_probabilistic_model(symbol=self.node.symbol, 
                                                                                              model=self.derivation.probabilistic_model,
                                                                                              remaining_recursions=remaining_recursions)
            if self.production.left in self.production.right:
                remaining_recursions -= 1

            self.leaves = []
            # Split remaining recursions between recursive nonterminal symbols
            # TODO WARNING It can reduce or increase by 1 the remaining recursions
            remaining_recursions_per_recursive_nt = np.round(np.random.dirichlet(np.ones(self.production.recursion_arity)) * remaining_recursions).tolist()

            for leave in self.production.right:
                self.leaves.append(Tree(derivation=derivation,
                                        remaining_recursions=remaining_recursions_per_recursive_nt.pop() if leave.recursive else 0,                                        root=self,
                                        node=leave))

    @property
    def depth(self):
        if self._depth is None:
            if self.leaves is None:
                self._depth = 0
            else:
                self._depth = reduce(lambda x, y: x if (x.depth > y.depth) else y, self.leaves).depth + 1
        return self._depth

    @property
    def recursions(self):
        if self._recursions is None:
            if self.leaves is None:
                self._recursions = 0
            else:
                self._recursions = reduce(lambda x, y: x + y.recursions, self.leaves, 1 if self.production.left in self.production.right else 0)
        return self._recursions

    @property
    def height(self):
        if self._height is None:
            if self.root is None:
                self._height = 0
            else:
                self._height = self.root.height + 1
        return self._height

    @property
    def word(self):
        if self._word is None:
            self._word = []
            for leave in self.leaves:
                self._word += leave.word
        return self._word

    def __str__(self):
        if self.str_word is None:
            self.str_word = str()
            for letter in self.word:
                self.str_word += ' ' + letter.__str__()
        return self.str_word

    @property
    def non_terminal_subtrees_by_node(self):
        if self._non_terminal_nodes is None:
            self._non_terminal_nodes = {}
            if isinstance(self.node, NonTerminal):
                # Axiom is included as a crossover node
                #if not self.node.axiom:
                self._non_terminal_nodes[self.node] = [self]
                if self.leaves is not None:
                    for leave in self.leaves:
                        for nt in leave.non_terminal_subtrees_by_node:
                            if nt in self._non_terminal_nodes:
                                self._non_terminal_nodes[nt] += leave.non_terminal_subtrees_by_node[nt]
                            else:
                                self._non_terminal_nodes[nt] = leave.non_terminal_subtrees_by_node[nt]

        return self._non_terminal_nodes

    def reset_non_terminal_nodes(self):
        self._non_terminal_nodes is None
        if self.leaves is not None:
            for leave in self.leaves:
                leave.reset_non_terminal_nodes()

    @property
    def non_terminal_nodes_by_depth(self):
        if self._non_terminal_nodes_by_depth is None:
            self._non_terminal_nodes_by_depth = {}
            if isinstance(self.node, NonTerminal):
                if not self.node.axiom:
                    self._non_terminal_nodes_by_depth[self.node][self.node.depth] = [self]
                if self.leaves is not None:
                    for leaf in self.leaves:
                        for nt in leaf.non_terminal_nodes_by_depth:
                            if nt in self._non_terminal_nodes_by_depth:
                                for d in self._non_terminal_nodes_by_depth[nt]:
                                    if d in self._non_terminal_nodes_by_depth[nt]:
                                        self._non_terminal_nodes_by_depth[nt][d] += \
                                        leaf.non_terminal_nodes_by_depth[nt][d]
                                    else:
                                        self._non_terminal_nodes_by_depth[nt][d] = \
                                        leaf.non_terminal_nodes_by_depth[nt][d]
                            else:
                                self._non_terminal_nodes_by_depth[nt] = leaf.non_terminal_nodes_by_depth[nt]

        return self._non_terminal_nodes_by_depth
        

class WX:
    @staticmethod
    def crossover(derivations, max_recursions=0):
        common_non_terminals = derivations[0].tree.non_terminal_subtrees_by_node.keys() & derivations[1].tree.non_terminal_subtrees_by_node.keys()
        choosing_pool = []
        if common_non_terminals:
            for nt in common_non_terminals:
                choosing_pool += (
                    [nt] * (len(derivations[0].tree.non_terminal_subtrees_by_node[nt]) + len(derivations[1].tree.non_terminal_subtrees_by_node[nt])))

            while True:
                crossover_non_terminal = np.random.choice(choosing_pool)

                genome_0_crossover_node = np.random.choice(derivations[0].tree.non_terminal_subtrees_by_node[crossover_non_terminal])
                genome_1_crossover_node = np.random.choice(derivations[1].tree.non_terminal_subtrees_by_node[crossover_non_terminal])

                if (genome_0_crossover_node.recursions + derivations[1].recursions - genome_1_crossover_node.recursions) <= max_recursions and \
                   (genome_1_crossover_node.recursions + derivations[0].recursions - genome_0_crossover_node.recursions) <= max_recursions:
                    break

            son_1 = Derivation(max_recursions=max_recursions, tree=derivations[0].tree, crossover_node=genome_0_crossover_node, subtree=genome_1_crossover_node)
            son_2 = Derivation(max_recursions=max_recursions, tree=derivations[1].tree, crossover_node=genome_1_crossover_node, subtree=genome_0_crossover_node)
            return son_1, son_2
        else:
            return None, None


class OnePointMutation:
    @staticmethod
    def mutate(individuals, mutation_rate=0.05):
        for i in individuals:
            d = i.derivation
            if np.random.random() < mutation_rate:
                tree = np.random.choice([tree for nt in d.tree.non_terminal_subtrees_by_node.values() for tree in nt])
                d.tree.reset_non_terminal_nodes()
                node_remaining_recursions = d.max_recursions - d.recursions + tree.recursions
                if node_remaining_recursions<0:
                    print("max" + str(d.max_recursions))
                    print(d.recursions)
                    print(d)
                    print(tree.recursions)
                    print(tree)
                    node_remaining_recursions=0

                if tree.node.recursive:
                    remaining_recursions = np.random.random_integers(0, node_remaining_recursions) if d.probabilistic_model is ProbabilisticModel.uniform else node_remaining_recursions
                else:
                    remaining_recursions = 0
                tree.production = d.grammar.get_production_with_probabilistic_model(symbol=tree.node.symbol,
                                                                                    model=d.probabilistic_model,
                                                                                    remaining_recursions=remaining_recursions)

                # Information generated during properties method calls
                tree._word = None
                tree.str_word = None
                tree._non_terminal_nodes = None
                tree._non_terminal_nodes_by_depth = None
                tree._depth = None
                tree._recursions = None
                tree._height = None
                tree.leaves = []

                if tree.production.left in tree.production.right:
                    remaining_recursions -= 1

                # Split remaining recursions between recursive nonterminal symbols
                remaining_recursions = np.round(np.random.dirichlet(np.ones(tree.production.recursion_arity)) * remaining_recursions).tolist()

                for leave in tree.production.right:
                    tree.leaves.append(Tree(derivation=d,
                                            remaining_recursions=remaining_recursions.pop() if leave.recursive else 0,
                                            root=tree,
                                            node=leave))
        return individuals


if __name__ == "__main__":
    from gplpy.gggp.grammar import CFG
    gr = CFG("../../gr/symbolic_regression_problem.gr")
    Derivation.grammar = gr
    b = Derivation(max_recursions=50)
    print(b.recursions, b)


__author__ = "aturing"
__license__ = "Apache License 2.0"
__version__ = "1.1.0"
__maintainer__ = "Pablo Ramos"
__email__ = "pablo.ramos@aturing.com"
__status__ = "Production"