# copyright ############################### #
# This file is part of the Xdeps Package.   #
# Copyright (c) CERN, 2021.                 #
# ######################################### #

from dataclasses import dataclass, field
from collections import defaultdict
import logging
from copy import deepcopy

from .refs import ARef, CallRef, Ref, ObjectAttrRef, RefCount
from .utils import os_display_png, mpl_display_png, ipy_display_png
from .utils import AttrDict
from .sorting import toposort

logger = logging.getLogger(__name__)


def dct_merge(dct1, dct2):
    return {**dct1, **dct2}


def _check_root_owner(t, ref):
    if hasattr(t, "_owner"):
        if t._owner is ref:
            return True
        else:
            return _check_root_owner(t._owner, ref)
    else:
        return False


class FuncWrapper:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        return CallRef(self.func, args, tuple(kwargs.items()))


class Task:
    taskid: object
    targets: set
    dependencies: set

    def run(self):
        raise NotImplemented


class GenericTask(Task):
    taskid: object
    action: object
    targets: set
    dependencies: set

    def __repr__(self):
        return f"<Task {self.taskid}:{self.dependencies}=>{self.targets}>"

    def run(self, *args):
        return self.action(*args)


class ExprTask(Task):
    def __init__(self, target, expr):
        self.taskid = target
        self.targets = target._get_dependencies()
        self.dependencies = expr._get_dependencies()
        self.expr = expr

    def __repr__(self):
        return f"{self.taskid} = {self.expr}"

    def run(self):
        value = self.expr._get_value()
        self.taskid._set_value(value)

    def info(self):
        print(f"#  {self.taskid}._expr._get_dependencies()")
        for pp in self.expr._get_dependencies():
            print(f"   {pp} = {pp._get_value()}")
        print()


class InheritanceTask(Task):
    def __init__(self, children, parents):
        self.taskid = children
        self.targets = set([children])
        self.dependencies = set(parents)

    def __repr__(self):
        return f"{self.taskid} <- {self.parents}"

    def run(self, event):
        key, value, isattr = event
        for target in self.targets:
            if isattr:
                getattr(target, key)._set_value(value)
            else:
                target[key]._set_value(value)


class DepEnv:
    __slots__ = ("_data", "_")

    def __init__(self, data, ref):
        object.__setattr__(self, "_data", data)
        object.__setattr__(self, "_", ref)

    def __getattr__(self, key):
        return getattr(self._data, key)

    def __getitem__(self, key):
        return self._data[key]

    def __setattr__(self, key, value):
        self._[key] = value

    def __setitem__(self, key, value):
        self._[key] = value

    def _eval(self, expr):
        return self._._eval(expr)


class Manager:
    """
    Value dependency manager:

    tasks: taskid -> task
    rdeps: ref -> set of all refs that depends on `ref`
    rtasks: taskid -> all tasks whose dependencies are affected by taskid
    deptasks: ref -> all tasks that has ref as dependency
    tartasks: ref -> all tasks that has ref as target
    containers: label -> controlled container
    """

    def __init__(self):
        self.tasks = {}
        self.containers = {}
        self.rdeps = defaultdict(RefCount)
        self.rtasks = defaultdict(RefCount)
        self.deptasks = defaultdict(RefCount)
        self.tartasks = defaultdict(RefCount)
        self._tree_frozen = False

    def ref(self, container=None, label="_"):
        """Return a ref to an instance (or dict) associated to a label.

        Label must be unique.
        """
        if container is None:
            container = AttrDict()
        objref = Ref(container,  label, self)
        assert label not in self.containers
        self.containers[label] = objref
        return objref

    def set_value(self, ref, value):
        """Set a value pointed by a ref and execute all tasks that depends on ref.

        If the value is a Ref, create a new task from the ref.
        """
        logger.info("set_value %s %s", ref, value)
        if ref in self.tasks:
            self.unregister(ref)
        if isinstance(value, ARef):  # value is an expression
            self.register(ExprTask(ref, value))
            value = value._get_value()  # to be updated
        ref._set_value(value)
        self._run_tasks(self.find_tasks(ref._get_dependencies()))

    def _run_tasks(self, tasks):
        for task in tasks:
            logger.info("Run %s", task)
            task.run()

    def register(self, task):
        """Register a new task identified by taskid"""
        # logger.info("register %s",taskid)
        if self._tree_frozen:
            raise ValueError("Expressions and references cannot be changed "
                             "because the tree is frozen (e.g. because "
                             "a variables cache is active)")
        taskid = task.taskid
        self.tasks[taskid] = task
        for dep in task.dependencies:
            # logger.info("%s have an impact on %s",dep,task.targets)
            self.rdeps[dep].extend(task.targets)
            # logger.info("%s is used by T:%s",dep,taskid)
            self.deptasks[dep].append(taskid)
            for deptask in self.tartasks[dep]:
                # logger.info("%s modifies deps of T:%s",deptask,taskid)
                self.rtasks[deptask].append(taskid)
        for tar in task.targets:
            # logger.info("%s is modified by T:%s",tar,taskid)
            self.tartasks[tar].append(taskid)
            for deptask in self.deptasks[tar]:
                # logger.info("T:%s modifies deps of T:%s",taskid,deptask)
                self.rtasks[taskid].append(deptask)

    def unregister(self, taskid):
        """Unregister the task identified by taskid"""
        if self._tree_frozen:
            raise ValueError("Expressions and references cannot be changed "
                             "because the tree is frozen (e.g. because "
                             "a variables cache is active)")
        task = self.tasks[taskid]
        for dep in task.dependencies:
            for target in task.targets:
                if target in self.rdeps[dep]:
                    self.rdeps[dep].remove(target)
            if taskid in self.rtasks[dep]:
                self.rtasks[dep].remove(taskid)
            if taskid in self.deptasks[dep]:
                self.deptasks[dep].remove(taskid)
        for tar in task.targets:
            self.tartasks[tar].remove(taskid)
            for deptask in self.deptasks[tar]:
                if taskid in self.rtasks[deptask]:
                    self.rtasks[taskid].remove(deptask)
        if taskid in self.rtasks:
            del self.rtasks[taskid]
        del self.tasks[taskid]

    def find_deps(self, start_set):
        """Find all refs that depends on ref in start_seps"""
        assert type(start_set) in (list, tuple, set)
        deps = toposort(self.rdeps, start_set)
        return deps

    def find_taskids_from_tasks(self, start_tasks=None):
        """Find all taskids whose dependencies are affected by the tasks in start_tasks"""
        if start_tasks is None:
            start_tasks = self.rtasks
        tasks = toposort(self.rtasks, start_tasks)
        return tasks

    def find_taskids(self, start_deps=None):
        """Find all taskids that depend on the refs in start_deps"""
        if start_deps is None:
            start_deps = self.rdeps
        start_tasks = set()
        for dep in start_deps:
            start_tasks.update(self.deptasks[dep])
        tasks = toposort(self.rtasks, start_tasks)
        return tasks

    def find_tasks(self, start_deps=None):
        """Find all tasks that depend on the refs in start_deps"""
        if start_deps is None:
            start_deps = self.rdeps
        return [self.tasks[taskid] for taskid in self.find_taskids(start_deps)]

    def iter_expr_tasks_owner(self, ref):
        """Return all ExprTask defintions that write registered container"""
        for t in self.find_tasks():
            # TODO check for all targets or limit to ExprTask
            if _check_root_owner(t.taskid, ref):
                yield str(t.taskid), str(t.expr)

    def copy_expr_from(self, mgr, name, bindings=None):
        """
        Copy expression from another manager

        name: one of toplevel container in mgr
        bindings: dictionary mapping old container names into new container refs
        """
        ref = mgr.containers[name]
        if bindings is None:
            cmbdct = self.containers
        else:
            cmbdct = dct_merge(self.containers, bindings)
        self.load(mgr.iter_expr_tasks_owner(ref), cmbdct)

    def mk_fun(self, name, **kwargs):
        """Write a python function that executes a set of tasks in order of dependencies:
        name: name of the functions
        kwards:
            the keys are used to defined the argument name of the functions
            the values are the refs that will be set
        """
        varlist, start = list(zip(*kwargs.items()))
        tasks = self.find_tasks(start)
        fdef = [f"def {name}({','.join(varlist)}):"]
        for vname, vref in kwargs.items():
            fdef.append(f"  {vref} = {vname}")
        for tt in tasks:
            fdef.append(f"  {tt}")
        fdef = "\n".join(fdef)
        return fdef

    def gen_fun(self, name, **kwargs):
        """Return a python function that executes a set of tasks in order of dependencies:
        name: name of the functions
        kwards:
            the keys are used to defined the argument name of the functions
            the values are the refs that will be set
        """
        fdef = self.mk_fun(name, **kwargs)
        gbl = {}
        lcl = {}
        gbl.update((k, r._owner) for k, r in self.containers.items())
        exec(fdef, gbl, lcl)
        return lcl[name]

    def plot_deps(self, start=None, backend="ipy"):
        """Plot a graph of task and target dependencies from start.

        Possible backend:
            mpl: generate a figure in matplotlib
            os: generate a file /tmp/out.png and use `display` to show it
            ipy: use Ipython facility for Jupyter notebooks
        """
        from pydot import Dot, Node, Edge

        if start is None:
            start = list(self.rdeps)
        pdot = Dot("g", graph_type="digraph", rankdir="LR")
        for task in self.find_tasks(start):
            tn = Node(" " + str(task.taskid), shape="circle")
            pdot.add_node(tn)
            for tt in task.targets:
                pdot.add_node(Node(str(tt), shape="square"))
                pdot.add_edge(Edge(tn, str(tt), color="blue"))
            for tt in task.dependencies:
                pdot.add_node(Node(str(tt), shape="square"))
                pdot.add_edge(Edge(str(tt), tn, color="blue"))
        png = pdot.create_png()
        if backend == "mpl":
            mpl_display_png(png)
        elif backend == "os":
            os_display_png(png)
        elif backend == "ipy":
            ipy_display_png(png)
        return pdot

    def plot_tasks(self, start=None, backend="ipy"):
        """Plot a graph of task dependencies

        Possible backend:
            mpl: generate a figure in matplotlib
            os: generate a file /tmp/out.png and use `display` to show it
            ipy: use Ipython facility for Jupyter notebooks
        """
        from pydot import Dot, Node, Edge

        if start is None:
            start = list(self.rdeps)
        pdot = Dot("g", graph_type="digraph", rankdir="LR")
        for task in self.find_tasks(start):
            tn = Node(str(task.taskid), shape="circle")
            pdot.add_node(tn)
            for dep in task.dependencies:
                for tt in self.tartasks[dep]:
                    pdot.add_edge(Edge(str(tt), tn, color="blue"))
        png = pdot.create_png()
        if backend == "mpl":
            mpl_display_png(png)
        elif backend == "os":
            os_display_png(png)
        elif backend == "ipy":
            ipy_display_png(png)
        return pdot

    def dump(self):
        """Dump in json all ExprTask defined in the manager"""
        data = [
            (str(tt.taskid), str(tt.expr))
            # for t in self.find_tasks(self.rdeps)
            for tt in self.tasks.values()
            if isinstance(tt, ExprTask)
        ]
        return data

    def load(self, dump, dct=None):
        """Reload the expressions in dump  using container in dct

        dump: list of (lhs,rhs) pairs
        dct: dictionary of named references of containers,
             self containers by default

        """
        if dct is None:
            dct = self.containers
        for lhs, rhs in dump:
            lhs = eval(lhs, {}, dct)
            rhs = eval(rhs, {}, dct)
            task = ExprTask(lhs, rhs)
            if lhs in self.tasks:
                self.unregister(lhs)
            self.register(task)

    def newenv(self, label="_", data=None):
        "Experimental"
        if data is None:
            data = AttrDict()
        ref = self.ref(data, label=label)
        return DepEnv(data, ref)

    def refattr(self, container=None, label="_"):
        "Experimental"
        if container is None:
            container = AttrDict()
        objref = ObjectAttrRef(container, label, self)
        assert label not in self.containers
        self.containers[label] = objref
        return objref

    def cleanup(self):
        """
        Remove empty sets from dicts
        """
        for dct in self.rdeps, self.rtasks, self.deptasks, self.tartasks:
            for kk, ss in list(dct.items()):
                if len(ss) == 0:
                    del dct[kk]

    def copy(self):
        """
        Create a copy of in new manager
        """
        other = Manager()
        other.containers = deepcopy(self.containers)
        other.tasks = deepcopy(self.tasks)
        other.rdeps = deepcopy(self.rdeps)
        other.rtasks = deepcopy(self.rtasks)
        other.deptasks = deepcopy(self.deptasks)
        other.tartasks = deepcopy(self.tartasks)
        return other

    def clone(self):
        """
        Regenerate a new manager
        """
        other = Manager()
        other.containers.update(self.containers)
        for task in self.tasks.values():
            other.register(task)
        other.cleanup()
        return other

    def verify(self, dcts=("rdeps", "rtasks", "deptasks", "tartasks")):
        self.cleanup()
        other = self.clone()
        for dct in dcts:
            odct = getattr(other, dct)
            sdct = getattr(self, dct)
            for kk, ss in list(sdct.items()):
                if set(ss) != set(odct[kk]):
                    print(f"{dct}[{kk}] not consistent")
                    print(f"{dct}[{kk}] self - check:", set(ss) - set(odct[kk]))
                    print(f"{dct}[{kk}] check - self:", set(odct[kk]) - set(ss))
                    raise (ValueError(f"{self} is not consistent in {dct}[{kk}]"))

    def refresh(self):
        self.rdeps = defaultdict(RefCount)
        self.rtasks = defaultdict(RefCount)
        self.deptasks = defaultdict(RefCount)
        self.tartasks = defaultdict(RefCount)
        for task in self.tasks.values():
            self.register(task)
        self.cleanup()
