from __future__ import division

from functools import cmp_to_key
from itertools import cycle

import param
import numpy as np

from ..core.dimension import Dimension
from ..core.data import Dataset
from ..core.operation import Operation
from ..core.util import OrderedDict, unique_array, RecursionError, get_param_values
from .graphs import Graph, Nodes, EdgePaths, redim_graph
from .util import quadratic_bezier


class _layout_sankey(Operation):
    """
    Computes a Sankey diagram from a Graph element for internal use in
    the Sankey element constructor.

    Adapted from d3-sankey under BSD-3 license.
    """

    bounds = param.NumericTuple(default=(0, 0, 1000, 500))

    node_width = param.Number(default=15, doc="""
        Width of the nodes.""")

    node_padding = param.Integer(default=10, doc="""
        Number of pixels of padding relative to the bounds.""")

    iterations = param.Integer(default=32, doc="""
        Number of iterations to run the layout algorithm.""")

    def _process(self, element, key=None):
        nodes, edges, graph = self.layout(element, **self.p)
        params = get_param_values(element)
        return Sankey((element.data, nodes, edges), sankey=graph, **params)

    def layout(self, element, **params):
        self.p = param.ParamOverrides(self, params)
        graph = {'nodes': [], 'links': []}
        self.computeNodeLinks(element, graph)
        self.computeNodeValues(graph)
        self.computeNodeDepths(graph)
        self.computeNodeBreadths(graph)
        self.computeLinkBreadths(graph)
        paths = self.computePaths(graph)

        node_data = []
        for node in graph['nodes']:
            node_data.append((np.mean([node['x0'], node['x1']]),
                              np.mean([node['y0'], node['y1']]),
                              node['index'])+tuple(node['values']))
        if element.nodes.ndims == 3:
            kdims = element.nodes.kdims
        elif element.nodes.ndims:
            kdims = element.node_type.kdims[:2] + element.nodes.kdims[-1:]
        else:
            kdims = element.node_type.kdims
        nodes = element.node_type(node_data, kdims=kdims, vdims=element.nodes.vdims)
        edges = element.edge_type(paths)
        return nodes, edges, graph


    def computePaths(self, graph):
        paths = []
        for link in graph['links']:
            source, target = link['source'], link['target']
            x0, y0 = source['x1'], link['y0']
            x1, y1 = target['x0'], link['y1']
            start = np.array([(x0, link['width']+y0),
                              (x0, y0)])
            src = (x0, y0)
            ctr1 = ((x0+x1)/2., y0)
            ctr2 = ((x0+x1)/2., y1)
            tgt = (x1, y1)
            bottom = quadratic_bezier(src, tgt, ctr1, ctr2)
            mid = np.array([(x1, y1),
                            (x1, y1+link['width'])])

            xmid = (x0+x1)/2.
            y0 = y0+link['width']
            y1 = y1+link['width']
            src = (x1, y1)
            ctr1 = (xmid, y1)
            ctr2 = (xmid, y0)
            tgt = (x0, y0)
            top = quadratic_bezier(src, tgt, ctr1, ctr2)
            spline = np.concatenate([start, bottom, mid, top])
            paths.append(spline)
        return paths


    @classmethod
    def weightedSource(cls, link):
        return cls.nodeCenter(link['source']) * link['value']

    @classmethod
    def weightedTarget(cls, link):
        return cls.nodeCenter(link['target']) * link['value']

    @classmethod
    def nodeCenter(cls, node):
        return (node['y0'] + node['y1']) / 2

    @classmethod
    def ascendingBreadth(cls, a, b):
        return int(a['y0'] - b['y0'])

    @classmethod
    def ascendingSourceBreadth(cls, a, b):
        return cls.ascendingBreadth(a['source'], b['source']) | a['index'] - b['index']

    @classmethod
    def ascendingTargetBreadth(cls, a, b):
        return cls.ascendingBreadth(a['target'], b['target']) | a['index'] - b['index']

    @classmethod
    def computeNodeLinks(cls, element, graph):
        """
        Populate the sourceLinks and targetLinks for each node.
        Also, if the source and target are not objects, assume they are indices.
        """
        index = element.nodes.kdims[-1]
        node_map = {}
        if element.nodes.vdims:
            values = zip(*(element.nodes.dimension_values(d)
                           for d in element.nodes.vdims))
        else:
            values = cycle([tuple()])
        for index, vals in zip(element.nodes.dimension_values(index), values):
            node = {'index': index, 'sourceLinks': [], 'targetLinks': [], 'values': vals}
            graph['nodes'].append(node)
            node_map[index] = node

        links = [element.dimension_values(d) for d in element.dimensions()[:3]]
        for i, (src, tgt, value) in enumerate(zip(*links)):
            source, target = node_map[src], node_map[tgt]
            link = dict(index=i, source=source, target=target, value=value)
            graph['links'].append(link)
            source['sourceLinks'].append(link)
            target['targetLinks'].append(link)

    @classmethod
    def computeNodeValues(cls, graph):
        """
        Compute the value (size) of each node by summing the associated links.
        """
        for node in graph['nodes']:
            source_val = np.sum([l['value'] for l in node['sourceLinks']])
            target_val = np.sum([l['value'] for l in node['targetLinks']])
            node['value'] = max([source_val, target_val])

    def computeNodeDepths(self, graph):
        """
        Iteratively assign the depth (x-position) for each node.
        Nodes are assigned the maximum depth of incoming neighbors plus one;
        nodes with no incoming links are assigned depth zero, while
        nodes with no outgoing links are assigned the maximum depth.
        """
        nodes = graph['nodes']
        depth = 0
        while nodes:
            next_nodes = []
            for node in nodes:
                node['depth'] = depth
                for link in node['sourceLinks']:
                    if not any(link['target'] is node for node in next_nodes):
                        next_nodes.append(link['target'])
            nodes = next_nodes
            depth += 1
            if depth > 10000:
                raise RecursionError('Sankey diagrams only support acyclic graphs.')

        nodes = graph['nodes']
        depth = 0
        while nodes:
            next_nodes = []
            for node in nodes:
                node['height'] = depth
                for link in node['targetLinks']:
                    if not any(link['source'] is node for node in next_nodes):
                        next_nodes.append(link['source'])
            nodes = next_nodes
            depth += 1
            if depth > 10000:
                raise RecursionError('Sankey diagrams only support acyclic graphs.')

        x0, _, x1, _ = self.p.bounds
        dx = self.p.node_width
        kx = (x1 - x0 - dx) / (depth - 1)
        for node in graph['nodes']:
            d  = node['depth'] if node['sourceLinks'] else depth - 1
            node['x0'] = x0 + max([0, min([depth-1, np.floor(d)]) * kx])
            node['x1'] = node['x0'] + dx

    def computeNodeBreadths(self, graph):
        node_map = OrderedDict()
        for n in graph['nodes']:
            if n['x0'] not in node_map:
                node_map[n['x0']] = []
            node_map[n['x0']].append(n)

        _, y0, _, y1 = self.p.bounds
        py = self.p.node_padding

        def initializeNodeBreadth():
            kys = []
            for nodes in node_map.values():
                nsum = np.sum([node['value'] for node in nodes])
                ky = (y1 - y0 - (len(nodes)-1) * py) / nsum
                kys.append(ky)
            ky = np.min(kys)

            for nodes in node_map.values():
                for i, node in enumerate(nodes):
                    node['y0'] = i
                    node['y1'] = i + node['value'] * ky

            for link in graph['links']:
                link['width'] = link['value'] * ky

        def relaxLeftToRight(alpha):
            for nodes in node_map.values():
                for node in nodes:
                    if not node['targetLinks']:
                        continue
                    weighted = sum([self.weightedSource(l) for l in node['targetLinks']])
                    tsum = sum([l['value'] for l in node['targetLinks']])
                    center = self.nodeCenter(node)
                    dy = (weighted/tsum - center)*alpha
                    node['y0'] += dy
                    node['y1'] += dy


        def relaxRightToLeft(alpha):
            for nodes in list(node_map.values())[::-1]:
                for node in nodes:
                    if not node['sourceLinks']:
                        continue
                    weighted = sum([self.weightedTarget(l) for l in node['sourceLinks']])
                    tsum = sum([l['value'] for l in node['sourceLinks']])
                    center = self.nodeCenter(node)
                    dy = (weighted/tsum - center)*alpha
                    node['y0'] += dy
                    node['y1'] += dy


        def resolveCollisions():
            for nodes in node_map.values():
                y = y0
                nodes.sort(key=cmp_to_key(self.ascendingBreadth))
                for node in nodes:
                    dy = y-node['y0']
                    if dy > 0:
                        node['y0'] += dy
                        node['y1'] += dy
                    y = node['y1'] + py

                dy = y-py-y1
                if dy > 0:
                    node['y0'] -= dy
                    node['y1'] -= dy
                    y = node['y0']
                    for node in nodes[:-1][::-1]:
                        dy = node['y1'] + py - y;
                        if dy>0:
                            node['y0'] -= dy
                            node['y1'] -= dy
                        y = node['y0']

        initializeNodeBreadth()
        resolveCollisions()
        alpha = 1
        for _ in range(self.p.iterations):
            alpha = alpha*0.99
            relaxRightToLeft(alpha)
            resolveCollisions()
            relaxLeftToRight(alpha)
            resolveCollisions()

    @classmethod
    def computeLinkBreadths(cls, graph):
        for node in graph['nodes']:
            node['sourceLinks'].sort(key=cmp_to_key(cls.ascendingTargetBreadth))
            node['targetLinks'].sort(key=cmp_to_key(cls.ascendingSourceBreadth))

        for node in graph['nodes']:
            y0 = y1 = node['y0']
            for link in node['sourceLinks']:
                link['y0'] = y0
                y0 += link['width']
            for link in node['targetLinks']:
                link['y1'] = y1
                y1 += link['width']



class Sankey(Graph):
    """
    Sankey is an acyclic, directed Graph type that represents the flow
    of some quantity between its nodes.
    """

    group = param.String(default='Sankey', constant=True)

    vdims = param.List(default=[Dimension('Value')])

    def __init__(self, data, kdims=None, vdims=None, **params):
        if isinstance(data, tuple):
            data = data + (None,)*(3-len(data))
            edges, nodes, edgepaths = data
        else:
            edges, nodes, edgepaths = data, None, None
        sankey_graph = params.pop('sankey', None)
        compute = not (sankey_graph and isinstance(nodes, Nodes) and isinstance(edgepaths, EdgePaths))
        super(Graph, self).__init__(edges, kdims=kdims, vdims=vdims, **params)
        if compute:
            if nodes is None:
                src = self.dimension_values(0, expanded=False)
                tgt = self.dimension_values(1, expanded=False)
                values = unique_array(np.concatenate([src, tgt]))
                nodes = Dataset(values, 'index')
            elif not isinstance(nodes, Dataset):
                try:
                    nodes = Dataset(nodes)
                except:
                    nodes = Dataset(nodes, 'index')
            if not nodes.kdims:
                raise ValueError('Could not determine index in supplied node data. '
                                 'Ensure data has at least one key dimension, '
                                 'which matches the node ids on the edges.')
            self._nodes = nodes
            nodes, edgepaths, graph = _layout_sankey.instance().layout(self)
            self._nodes = nodes
            self._edgepaths = edgepaths
            self._sankey = graph
        else:
            if not isinstance(nodes, self.node_type):
                raise TypeError("Expected Nodes object in data, found %s."
                                % type(nodes))
            self._nodes = nodes
            if not isinstance(edgepaths, self.edge_type):
                raise TypeError("Expected EdgePaths object in data, found %s."
                                % type(edgepaths))
            self._edgepaths = edgepaths
            self._sankey = sankey_graph
        self._validate()
        self.redim = redim_graph(self, mode='dataset')

    def clone(self, data=None, shared_data=True, new_type=None, link=True,
              *args, **overrides):
        if data is None:
            overrides['sankey'] = self._sankey
        return super(Sankey, self).clone(data, shared_data, new_type, link,
                                         *args, **overrides)
