import tensorflow as tf
from .sections import GraphSection, section_from_name
from tensorflow.python.client.session import _FetchMapper
from .utils import VerboseTimer
from contextlib import ExitStack


def _flatten_list(nested_list):
    """
    Given a nested list, returns flat list of all its elements

    :param nested_list: nested list
    :return: flat list
    """
    if not isinstance(nested_list, list):
        return [nested_list]

    res = []
    for sub_list in nested_list:
        res += _flatten_list(sub_list)

    return res


class GraphSectionManager(object):
    """
    Optimizes a graph with GraphSections by running partial backward passes.
    This reduces the memory consumption at the expense of an additional forward pass and multiple data transfers
    between GPU and main memory.
    """
    def __init__(self, graph=None):
        """
        Constructor
        :param graph: graph that is split into sections, defaults to tf.get_default_graph()
        """
        self._graph = tf.get_default_graph() if graph is None else graph

        self._sections = []
        self._cache = {}

        # which tensors to cache during full forward pass
        self._tensors_to_cache_in_fwd = []

        # caches for request op information
        self._req_input_tensors = {}  # requested op -> [list of input tensors]
        self._req_eval_sections = {}  # requested op -> section to evaluate in (or None for forward pass)
        self._req_reduced_inputs = {}  # requested op -> [list of input tensors minus those from eval_section]


    ####################################################################################################################
    # Section handling
    ####################################################################################################################

    def new_section(self):
        """
        Constructs a new graph section and adds it to the manager.

        :return: newly created section
        """
        section = GraphSection(graph=self._graph, manager=self)
        return section

    def add_section(self, section):
        """
        Adds a GraphSection to the list of managed sections.

        :param section: instance of GraphSection
        :return: index of added section
        """
        self._sections.append(section)
        return len(self._sections) - 1

    def get_sections(self):
        """
        Returns a list of the sections
        Sections are ordered in the way there were added.

        :return: list of sections
        """
        return self._sections


    ####################################################################################################################
    # Methods to prepare the graph for training
    ####################################################################################################################

    def add_training_ops(self, optimizer, loss=None, var_list=None, grads=None, global_step=None, verbose=False,
                         summaries=None):
        """
        Constructs a training operation for each section. If `grads` is not given, it is computed by
            grads = optimizer.compute_gradients(loss, var_list)

        Each section's training operations applies the gradients of the section's variables and runs all operations in
        the section's GraphKeys.UPDATE_OPS collection. All variables are assumed to be contained in the section's
        GLOBAL_VARIABLES collection.

        :param optimizer: tensorflow optimizer to use
        :param loss: loss tensor to optimize
        :param var_list: variable list to compute gradients on
        :param grads: gradients as returned by optimizer.compute_gradients, alternative to loss and var_list
        :param global_step: global step tensor to increment after full backward pass
        :param verbose: if True, adds tf.Print operations to log backward passes over sections
        :param summaries: optional list of collections to add gradient histogram summaries to. Defaults to None
        """

        # add gradient computation nodes for all trainable variables
        if grads is None:
            assert loss is not None, 'Either gradients or loss have to be given.'
            grads = optimizer.compute_gradients(loss, var_list=var_list)

        # store references to gradient ops for simple access
        grad_dict = {v: g for g, v in grads}

        for s, section in reversed(list(enumerate(self.get_sections()))):
            # construct gradient application
            xs = section.get_collection(tf.GraphKeys.GLOBAL_VARIABLES).intersection(grad_dict.keys())
            apply_op = None
            if len(xs) > 0:
                cur_grads = [(grad_dict[v], v) for v in xs]

                if summaries is not None:
                    for v in xs:
                        tf.summary.histogram('gradients/%s' + v.name, grad_dict[v], collections=summaries)

                # if we should be verbose, log backward passes
                if verbose:
                    cur_grads[0] = (tf.Print(cur_grads[0][0], [cur_grads[0][0]],
                                             'Running backward pass on section %d' % s), cur_grads[0][1])

                # only increment global step in last partial backward pass
                apply_op = optimizer.apply_gradients(cur_grads, global_step=global_step if s == 0 else None)

                print("Found %d gradient application operations in section %d. Adding to training op."
                      % (len(cur_grads), s))

            # no gradients to apply
            else:
                print("Section %d does not contain gradient application operations." % s)

                # in last section's backward pass -> need to increment global step separately
                if s == 0 and global_step is not None:
                    apply_op = tf.assign_add(global_step, 1)


            # group update operatio
ns
            update_op = None
            update_ops = section.get_collection(tf.GraphKeys.UPDATE_OPS)
            if len(update_ops) > 0:

                print("Found %d update operations in section %d. Adding to training op." % (len(update_ops), s))
                update_op = tf.group(*update_ops)
            else:
                print("Section %d does not contain update operations." % s)

            # construct final training operation
            if apply_op is not None and update_op is not None:
                train_op = tf.group(apply_op, update_op)
            elif apply_op is not None:
                train_op = apply_op
            elif update_op is not None:
                train_op = update_op
            else:
                train_op = tf.no_op()

            section.set_training_op(train_op)

    def prepare_training(self):
        """
        Prepares the partial training by computing metadata about sections and creating training operations.
        Should be run after the full construction of the graph, including training ops.
        """
        assert len(self.get_sections()) > 0, 'There has to be at least one GraphSection in this graph.'

        # compute list of unique tensors to cache in forward pass
        self._tensors_to_cache_in_fwd = list(set([t for s in self.get_sections() for t in s.get_incoming()]))

        # compute tensors to feed into backward passes
        all_tensors_to_feed = set()
        for s, section in reversed(list(enumerate(self.get_sections()))):
            # find list of tensors to feed into training operation
            tensors_to_feed = self._find_feeds_from_other_sections(section.get_training_op(),
                                                                   ignore=[section], given=section.get_incoming())

            # store info in section
            section.set_tensors_to_feed([t[0] for t in tensors_to_feed])
            all_tensors_to_feed.update(tensors_to_feed)

        # tell sections to cache those tensors needed by other sections
        self._mark_tensors_for_caching(all_tensors_to_feed, only_next_run=False)

    def _mark_tensors_for_caching(self, tensors, only_next_run=False):
        """
        Stores which tensors to cac