# -*- coding: utf-8 -*-
# Copyright (c) Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
"""
Force-Directed Graph Layout
===========================

This module contains implementations for a force-directed layout, where the
graph is modelled like a collection of springs or as a collection of
particles attracting and repelling each other. The whole graph tries to
reach a state which requires the minimum energy.
"""

import numpy as np

from ..util import _straight_line_vertices, _rescale_layout, issparse


class fruchterman_reingold(object):
    r"""Fruchterman-Reingold implementation adapted from NetworkX.

    In the Fruchterman-Reingold algorithm, the whole graph is modelled as a
    collection of particles, it runs a simplified particle simulation to
    find a nice layout for the graph.

    Parameters
    ----------
    optimal : number
        Optimal distance between nodes. Defaults to :math:`1/\\sqrt{N}` where
        N is the number of nodes.
    iterations : int
        Number of iterations to perform for layout calculation.
    pos : array
        Initial positions of the nodes

    Notes
    -----
    The algorithm is explained in more detail in the original paper [1]_.

    .. [1] Fruchterman, Thomas MJ, and Edward M. Reingold. "Graph drawing by
       force-directed placement." Softw., Pract. Exper. 21.11 (1991),
       1129-1164.
    """

    def __init__(self, optimal=None, iterations=50, pos=None):
        self.dim = 2
        self.optimal = optimal
        self.iterations = iterations
        self.num_nodes = None
        self.pos = pos

    def __call__(self, adjacency_mat, directed=False):
        """
        Starts the calculation of the graph layout.

        This is a generator, and after each iteration it yields the new
        positions for the nodes, together with the vertices for the edges
        and the arrows.

        There are two solvers here: one specially adapted for SciPy sparse
        matrices, and the other for larger networks.

        Parameters
        ----------
        adjacency_mat : array
            The graph adjacency matrix.
        directed : bool
            Wether the graph is directed or not. If this is True,
            it will draw arrows for directed edges.

        Yields
        ------
        layout : tuple
            For each iteration of the layout calculation it yields a tuple
            containing (node_vertices, line_vertices, arrow_vertices). These
            vertices can be passed to the `MarkersVisual` and `ArrowVisual`.
        """
        if adjacency_mat.shape[0] != adjacency_mat.shape[1]:
            raise ValueError("Adjacency matrix should be square.")

        self.num_nodes = adjacency_mat.shape[0]

        if issparse(adjacency_mat):
            # Use the sparse solver
            solver = self._sparse_fruchterman_reingold
        else:
            solver = self._fruchterman_reingold

        for result in solver(adjacency_mat, directed):
            yield result

    def _fruchterman_reingold(self, adjacency_mat, directed=False):
        if self.optimal is None:
            self.optimal = 1 / np.sqrt(self.num_nodes)

        if self.pos is None:
            # Random initial positions
            pos = np.asarray(
                np.random.random((self.num_nodes, self.dim)),
                dtype=np.float32
            )
        else:
            pos = self.pos.astype(np.float32)

        # Yield initial positions
        line_vertices, arrows = _straight_line_vertices(adjacency_mat, pos,
                                                        directed)
        yield pos, line_vertices, arrows

        # The initial "temperature"  is about .1 of domain area (=1x1)
        # this is the largest step allowed in the dynamics.
        t = 0.1

        # Simple cooling scheme.
        # Linearly step down by dt on each iteration so last iteration is
        # size dt.
        dt = t / float(self.iterations+1)
        # The inscrutable (but fast) version
        # This is still O(V^2)
        # Could use multilevel methods to speed this up significantly
        for iteration in range(self.iterations):
            delta_pos = _calculate_delta_pos(adjacency_mat, pos, t,
                                             self.optimal)
            pos += delta_pos
            _rescale_layout(pos)

            # cool temperature
            t -= dt

            # Calculate edge vertices and arrows
            line_vertices, arrows = _straight_line_vertices(adjacency_mat,
                                                            pos, directed)

            yield pos, line_vertices, arrows

    def _sparse_fruchterman_reingold(self, adjacency_mat, directed=False):
        # Optimal distance between nodes
        if self.optimal is None:
            self.optimal = 1 / np.sqrt(self.num_nodes)

        # Change to list of list format
        # Also construct the matrix in COO format for easy edge construction
        adjacency_arr = adjacency_mat.toarray()
        adjacency_coo = adjacency_mat.tocoo()

        if self.pos is None:
            # Random initial positions
            pos = np.asarray(
                np.random.random((self.num_nodes, self.dim)),
                dtype=np.float32
            )
        else:
            pos = self.pos.astype(np.float32)

        # Yield initial positions
        line_vertices, arrows = _straight_line_vertices(adjacency_coo, pos,
                                                        directed)
        yield pos, line_vertices, arrows

        # The initial "temperature"  is about .1 of domain area (=1x1)
        # This is the largest step allowed in the dynamics.
        t = 0.1
        # Simple cooling scheme.
        # Linearly step down by dt on each iteration so last iteration is
        # size dt.
        dt = t / float(self.iterations+1)
        for iteration in range(self.iterations):
            delta_pos = _calculate_delta_pos(adjacency_arr, pos, t,
                                             self.optimal)
            pos += delta_pos
            _rescale_layout(pos)

            # Cool temperature
            t -= dt

            # Calculate line vertices
            line_vertices, arrows = _straight_line_vertices(adjacency_coo,
                                                            pos, directed)

            yield pos, line_vertices, arrows


def _calculate_delta_pos(adjacency_arr, pos, t, optimal):
    """Helper to calculate the delta position"""
    # XXX eventually this should be refactored for the sparse case to only
    # do the necessary pairwise distances
    delta = pos[:, np.newaxis, :] - pos

    # Distance between points
    distance2 = (delta*delta).sum(axis=-1)
    # Enforce minimum distance of 0.01
    distance2 = np.where(distance2 < 0.0001, 0.0001, distance2)
    distance = np.sqrt(distance2)
    # Displacement "force"
    displacement = np.zeros((len(delta), 2))
    for ii in range(2):
        displacement[:, ii] = (
            delta[:, :, ii] *
            ((optimal * optimal) / (distance*distance) -
             (adjacency_arr * distance) / optimal)).sum(axis=1)

    length = np.sqrt((displacement**2).sum(axis=1))
    length = np.where(length < 0.01, 0.1, length)
    delta_pos = displacement * t / length[:, np.newaxis]
    return delta_pos
