
""" ----------------------------------------------------------------
# Force-Directed Edge Bundling
#
# Performance-optimised script version
# ------------------------------------------------------------------
"""

##Edge bundling=group
##input_layer=vector
##cluster_field=field input_layer
##use_clustering_result=boolean True
##initial_step_size=number 100.0
##compatibility=number 0.5
##cycles=number 6
##iterations=number 90
##bundled_edges=output vector

import math
import numpy as np
from datetime import datetime
from PyQt4.QtCore import QVariant

from qgis.core import QgsFeature, QgsPoint, QgsVector, QgsGeometry, QgsField, QGis

import processing


# I'm repeating the parameters here so Eclipse doesn't 
# show errors further down where they are used
initial_step_size = initial_step_size
iterations = iterations
cycles = cycles
use_clustering_result = use_clustering_result
compatibility = compatibility
input_layer = input_layer
cluster_field = cluster_field
progress = progress
bundled_edges = bundled_edges


def forcecalcx(x, y, d) :
    if abs(x) > eps and abs(y) > eps :
        x *= 1.0 / d
    else :
        x = 0.0
    return x
    
def forcecalcy(x, y, d) :
    if abs(x) > eps and abs(y) > eps :
        y *= 1.0 / d
    else :
        y = 0.0
    return y
        
# ------------------------------------ MISC ------------------------------------ #

class MiscUtils:

    @staticmethod
    def project_point_on_line(pt, line):
        """ Projects point onto line, needed for compatibility computation """
        v0 = line.vertexAt(0)
        v1 = line.vertexAt(1)
        length = max(line.length(), 10**(-6))
        r = ((v0.y() - pt.y()) * (v0.y() - v1.y()) - 
             (v0.x() - pt.x()) * (v1.x() - v0.x())) / (length**2)
        return QgsPoint(v0.x() + r * (v1.x() - v0.x()), v0.y() + r * (v1.y() - v0.y()))


# ------------------------------------ EDGE-CLUSTER  ------------------------------------ #

class EdgeCluster():
    
    def __init__(self, edges):
        self.S = initial_step_size  # Weighting factor (needs to be cached, because will be decreased in every cycle)
        self.I = iterations         # Number of iterations per cycle (needs to be cached, because will be decreased in every cycle)
        self.edges = edges          # Edges to bundle in this cluster
        self.edge_lengths = []      # Array to cache edge lenghts
        self.E = len(edges)         # Number of edges
        self.EP = 2                 # Current number of edge points
        self.SP = 0                 # Current number of subdivision points
        self.compatibility_matrix = np.zeros(shape=(self.E,self.E)) # Compatibility matrix
        self.direction_matrix = np.zeros(shape=(self.E,self.E))     # Encodes direction of edge pairs
        self.N = (2**cycles ) + 1                                   # Maximum number of points per edge
        self.epm_x = np.zeros(shape=(self.E,self.N))                # Bundles edges (x-values)
        self.epm_y = np.zeros(shape=(self.E,self.N))                # Bundles edges (y-values)
        
    def compute_compatibilty_matrix(self):
        """
        Compatibility is stored in a matrix (rows = edges, columns = edges). 
        Every coordinate in the matrix tells whether the two edges (r,c)/(c,r) 
        are compatible, or not. The diagonal is always zero, and the other fields 
        are filled with either -1 (not compatible) or 1 (compatible). 
        The matrix is symmetric. 
        """
        edges_as_geom = []
        edges_as_vect = []
        for e_idx, edge in enumerate(self.edges):
            geom = edge.geometry()
            edges_as_geom.append(geom)
            edges_as_vect.append(QgsVector(geom.vertexAt(1).x() - geom.vertexAt(0).x(), 
                                           geom.vertexAt(1).y() - geom.vertexAt(0).y()))
            self.edge_lengths.append(edges_as_vect[e_idx].length())
            
        progress = 0
        
        for i in range(self.E-1):
            for j in range(i+1, self.E):
                # Parameters 
                lavg = (self.edge_lengths[i] + self.edge_lengths[j]) / 2.0
                dot = edges_as_vect[i].normalized() * edges_as_vect[j].normalized()
                    
                # Angle compatibility 
                angle_comp = abs(dot)
                        
                # Scale compatibility 
                scale_comp = 2.0 / (lavg / min(self.edge_lengths[i], 
                                    self.edge_lengths[j]) + max(self.edge_lengths[i], 
                                    self.edge_lengths[j]) / lavg)
                        
                # Position compatibility 
                i0 = edges_as_geom[i].vertexAt(0)
                i1 = edges_as_geom[i].vertexAt(1)
                j0 = edges_as_geom[j].vertexAt(0)
                j1 = edges_as_geom[j].vertexAt(1)
                e1_mid = QgsPoint((i0.x() + i1.x()) / 2.0, (i0.y() + i1.y()) / 2.0)
                e2_mid = QgsPoint((j0.x() + j1.x()) / 2.0, (j0.y() + j1.y()) / 2.0)
                diff = QgsVector(e2_mid.x() - e1_mid.x(), e2_mid.y() - e1_mid.y())
                pos_comp = lavg / (lavg + diff.length())
                    
                # Visibility compatibility 
                mid_E1 = edges_as_geom[i].centroid()
                mid_E2 = edges_as_geom[j].centroid()
                #dist = mid_E1.distance(mid_E2)
                I0 = MiscUtils.project_point_on_line(j0, edges_as_geom[i])
                I1 = MiscUtils.project_point_on_line(j1, edges_as_geom[i])
                mid_I = QgsGeometry.fromPolyline([I0, I1]).centroid()
                dist_I = I0.distance(I1)
                if dist_I == 0.0:
                    visibility1 = 0.0
                else:
                    visibility1 = max(0, 1 - ((2 * mid_E1.distance(mid_I)) / dist_I))
                J0 = MiscUtils.project_point_on_line(i0, edges_as_geom[j])
                J1 = MiscUtils.project_point_on_line(i1, edges_as_geom[j])
                mid_J = QgsGeometry.fromPolyline([J0, J1]).centroid()
                dist_J = J0.distance(J1)
                if dist_J == 0.0:
                    visibility2 = 0.0
                else:
                    visibility2 = max(0, 1 - ((2 * mid_E2.distance(mid_J)) / dist_J))
                visibility_comp = min(visibility1, visibility2)

                # Compatibility score 
                comp_score = angle_comp * scale_comp * pos_comp * visibility_comp
                    
                # Fill values into the matrix (1 = yes, -1 = no) and use matrix symmetry (i/j = j/i) 
                if comp_score >= compatibility:
                    self.compatibility_matrix[i, j] = 1
                    self.compatibility_matrix[j, i] = 1
                else:
                    self.compatibility_matrix[i, j] = -1
                    self.compatibility_matrix[j, i] = -1
                    
                # Store direction 
         
       distStart1 = j0.distance(i0)
                distStart2 = j1.distance(i0)
                if distStart1 > distStart2:
                    self.direction_matrix[i, j] = -1
                    self.direction_matrix[j, i] = -1
                else:
                    self.direction_matrix[i, j] = 1
                    self.direction_matrix[j, i] = 1
    
    
    def force_directed_eb(self):
        """ Force-directed edge bund
ling """ 
        
        # Create compatibility matrix 
        self.compute_compatibilty_matrix()
        
        for e_idx, edge in enumerate(self.edges):
            vertices = edge.geometry().asPolyline()
            self.epm_x[e_idx, 0] = vertices[0].x()
            self.epm_y[e_idx, 0] = vertices[0].y()
            self.epm_x[e_idx, self.N-1] = vertices[1].x()
            self.epm_y[e_idx, self.N-1] = vertices[1].y()
        
        # For each cycle
        for c in range(cycles):
            print 'Cycle {0}'.format(c)
            # New number of subdivision points 
            current_num = self.EP
            currentindeces = []
            for i in range(current_num):
                idx = int((float(i) / float(current_num - 1)) * float(self.N - 1))
                c