import torch, torch.nn, torch.nn.functional, torch.distributions

"""
This file contains an implementation of a probabilistic decision tree.
If a node is a leaf, it should generate an action when sampled (see ./action.py)
Internal/brach nodes contain a list of child nodes and their relative weights.

To make these tree amenable to reinforcement learning, they act much like a torch.distribution.
They provide a .sample(), which walks the tree from top-to-bottom to generate a list of actions.
.log_prob() starts from an action, and walks it way backward to the root computing the log prob along the way.

For this reason, actions need a pointer to their generating nodes, and nodes need pointers to their parents.
TODO: Make parent a arg of __init__.

Instead of hard-coding weights into the trees, a dictionary of weights is passed to sample(...) and log_prob(...).
This allows the same tree to be re-used between different policies.
If multiple p-trees existed, it would be painful to compute the log prob of an action generated by
one tree under a second tree.
"""
# Define a node in a probabilistic decision tree.
class ProbablisticNode:
    def sample(self, *args): assert 0
    def log_prob(self, *args): assert 0

# Represents an internal node in a probabilistic decision tree.
class ProbabalisticBranch(ProbablisticNode):
    def __init__(self, node_list, node_weight_keys):
        self.node_list = node_list
        self.node_weight_keys = node_weight_keys
        self.parent = None
        for node in node_list: node.parent = self

    def _get_dist(self, weight_dict):
        weights = [weight_dict[key] for key in self.node_weight_keys]
        probs = torch.nn.functional.softmax(torch.FloatTensor(weights),dim=0)
        return torch.distributions.categorical.Categorical(probs=probs)

    def sample(self, count, weight_dict):
        actions = []
        dist = self._get_dist(weight_dict)
        for idx in dist.sample((count,)):
            actions.extend(self.node_list[idx].sample(1, weight_dict))
        return actions

    def log_prob(self, node, weight_dict):
        dist = self._get_dist(weight_dict)
        # Compute the probability of picking this index
        idx = self.node_list.index(node)
        log_prob = dist.log_prob(torch.IntTensor([idx]))
        # If our parent exists, ask it for the logprob of choosing us.
        return log_prob + (self.parent.log_prob(self, weight_dict) if self.parent else 0)

    def __str__(self, level=0):
        ret = "\t"*level+"node"+"\n"
        for child in self.node_list:
            ret += child.__str__(level=level+1)
        return ret
        
# If you want to add a new kind of action, you need to subclass this leaf node.
# All you need to do is override _sample and _log_prob.
class ProbabilisticLeaf(ProbablisticNode):
    def sample(self, count, weight_dict):
        return [self._sample(weight_dict) for _ in range(count)]
    # Compute the log prob by walking back up the tree.
    def log_prob(self, action, weight_dict):
        # If our parent exists, ask it for the logprob of choosing us.
        return self._log_prob(action, weight_dict) + (self.parent.log_prob(self, weight_dict) if self.parent else 0)

    # Sample a single action using a weight dictionary.
    def _sample(self, weight_dict): raise NotImplementedError("You need to implement this")
    # Given a single action, compute the log probability that the action came from the current leaf.
    def _log_prob(self, action, weight_dict): raise NotImplementedError("You need to implement this")
    def __str__(self, level=0):
        ret = "  "*level+self.__class__.__name__+"\n"
        return ret

