from collections import defaultdict
from distributed import SchedulerPlugin
from dask.core import reverse_dict
from dask.base import tokenize
from dask.order import graph_metrics, ndependencies
import click
from distributed.diagnostics.plugin import SchedulerPlugin
import numpy as np


def unravel_deps(hlg_deps, name, unravelled_deps=None):
    """Recursively construct a set of all dependencies for a specific task."""

    if unravelled_deps is None:
        unravelled_deps = set()

    for dep in hlg_deps[name]:
        unravelled_deps |= {dep}
        unravel_deps(hlg_deps, dep, unravelled_deps)

    return unravelled_deps


def get_node_depths(dependencies, root_nodes, metrics):
    node_depths = {}

    for k in dependencies.keys():
        # Get dependencies per node.
        deps = unravel_deps(dependencies, k)
        # Associate nodes with root nodes.
        roots = root_nodes & deps
        offset = metrics[k][-1]
        node_depths[k] = max(metrics[r][-1] - offset for r in roots) if roots else 0

    return node_depths


class schedular(SchedulerPlugin):
    def __init__(self, autorestrictor, local_cache):
        self.autorestrictor = autorestrictor
        self.local_cache = local_cache
        super().__init__()

    def add_worker(self, scheduler, worker):
        if self.local_cache:
            # Set the resource label to the ip of the node that the worker is on, so that tasks that require a specific node can be assigned to the correct worker.
            ip = worker[worker.rfind("/") + 1 : worker.rfind(":")]
            scheduler.add_resources(worker=worker, resources={ip: 1})

    def update_graph(self, scheduler, dsk=None, keys=None, restrictions=None, **kw):
        if self.autorestrictor:
            print("Using autorestrictor")
            """Processes dependencies to assign tasks to specific workers."""
            workers = list(scheduler.workers.keys())
            n_worker = len(workers)

            tasks = scheduler.tasks
            dependencies = kw["dependencies"]

            # print('In update_graph :', scheduler, ',*,', dsk, ',*,', keys , ',*,', restrictions , ',*,', kw)
            if dependencies:
                dependents = reverse_dict(dependencies)

                # print('reversed dict:', dependents)

                _, total_dependencies = ndependencies(dependencies, dependents)
                # TODO: Avoid calling graph metrics.
                metrics = graph_metrics(dependencies, dependents, total_dependencies)

                # Terminal nodes have no dependents, root nodes have no dependencies.
                # Horizontal partition nodes are initialized as the terminal nodes.
                part_nodes = {k for (k, v) in dependents.items() if not v}
                root_nodes = {k for (k, v) in dependencies.items() if not v}

                # Figure out the depth of every task. Depth is defined as maximum
                # distance from a root node. TODO: Optimize get_node_depths.

                node_depths = get_node_depths(dependencies, root_nodes, metrics)
                # try:
                max_depth = max(node_depths.values())
                # except:
                #        print('&&&&& dependencies, root_nodes, metrics',node_depths,',*,',dependencies, root_nodes, metrics)

                # If we have fewer partition nodes than workers, we cannot utilise all
                # the workers and are likely dealing with a reduction. We work our way
                # back through the graph, starting at the deepest terminal nodes, and
                # try to find a depth at which there was enough work to utilise all
                # workers.
                while (len(part_nodes) < n_worker) & (max_depth > 0):
                    _part_nodes = part_nodes.copy()
                    for pn in _part_nodes:
                        if node_depths[pn] == max_depth:
                            part_nodes ^= set((pn,))
                            part_nodes |= dependencies[pn]
                    max_depth -= 1
                    if max_depth <= 0:
                        return  # In this case, there in nothing we can do - fall back.

                part_roots = {}
                part_dependencies = {}
                part_dependents = {}

                for pn in part_nodes:
                    # Get dependencies per partition node.
                    part_dependencies[pn] = unravel_deps(dependencies, pn)
                    # Get dependents per partition node.
                    part_dependents[pn] = unravel_deps(dependents, pn)
                    # Associate partition nodes with root nodes.
                    part_roots[pn] = root_nodes & part_dependencies[pn]

                # Create a unique token for each set of partition roots. TODO: This is
                # very strict. What about nodes with very similar roots? Tokenization
                # may be overkill too.
                root_tokens = {tokenize(*sorted(v)): v for v in part_roots.values()}

                hash_map = defaultdict(set)
                group_offset = 0

                # Associate partition roots with a specific group if they are not a
                # subset of another, larger root set.
                for k, v in root_tokens.items():
                    if any(v < vv for vv in root_tokens.values()):  # Strict subset.
                        continue
                    else:
                        hash_map[k] |= set([group_offset])
                        group_offset += 1

                # If roots were a subset, they should share the group of their
                # superset/s.
                for k, v in root_tokens.items():
                    if not v:  # Special case - no dependencies. Handled below.
                        continue
                    shared_roots = {
                        kk: None for kk, vv in root_tokens.items() if v < vv
                    }
                    if shared_roots:
                        hash_map[k] = set().union(
                            *[hash_map[kk] for kk in shared_roots.keys()]
                        )

                task_groups = defaultdict(set)

                for pn in part_nodes:
                    pdp = part_dependencies[pn]
                    pdn = part_dependents[pn]

                    if pdp:
                        groups = hash_map[tokenize(*sorted(part_roots[pn]))]
                    else:  # Special case - no dependencies.
                        groups = {group_offset}
                        group_offset += 1

                    for g in groups:
                        task_groups[g] |= pdp | pdn | {pn}

                worker_loads = {wkr: 0 for wkr in workers}

                for task_group in task_groups.values():
                    assignee = min(worker_loads, key=worker_loads.get)
                    worker_loads[assignee] += len(task_group)

                    for task_name in task_group:
                        try:
                            task = tasks[task_name]
                        except KeyError:  # Keys may not have an assosciated task.
                            continue

                        # print('^^^^^^',dir(task))
                        # if task._worker_restrictions is None:
                        #    task._worker_restrictions = set()
                        # task._worker_restrictions |= {assignee}
                        # task._loose_restrictions = False

                        if task.worker_restrictions is None:
                            task.worker_restrictions = set()
                        task.worker_restrictions |= {assignee}
                        task.loose_restrictions = False


@click.command()
@click.option("--autorestrictor", default=False)
@click.option("--local_cache", default=False)
def dask_setup(scheduler, autorestrictor, local_cache):
    plugin = schedular(autorestrictor, local_cache)
    scheduler.add_plugin(plugin)
