/*
 * Copyright 2018 Dmitry Ustalov
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package org.nlpub.watset.mcl;

import org.jgrapht.Graph;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.impl.transforms.Pow;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.ops.transforms.Transforms;
import org.nlpub.watset.graph.Clustering;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

import static java.util.Objects.requireNonNull;

/**
 * This is an implementation of the Markov Clustering (MCL) algorithm.
 * It assumes processing of relatively small graphs due to the lack of
 * pruning optimizations.
 *
 * @param <V> vertices
 * @param <E> edges
 * @see <a href="https://micans.org/mcl/">van Dongen (2000)</a>
 */
public class MarkovClustering<V, E> implements Clustering<V> {
    public static final <V, E> Function<Graph<V, E>, Clustering<V>> provider(int e, double r) {
        return graph -> new MarkovClustering<>(graph, e, r);
    }

    public static final Integer ITERATIONS = 20;

    protected final Graph<V, E> graph;
    protected final int e;
    protected final double r;
    protected INDArray matrix;
    protected Map<V, Integer> index;

    public MarkovClustering(Graph<V, E> graph, int e, double r) {
        this.graph = requireNonNull(graph);
        this.e = e;
        this.r = r;
    }

    @Override
    public void run() {
        index = null;
        matrix = null;

        if (graph.vertexSet().isEmpty()) return;

        index = buildIndex(graph);
        matrix = buildMatrix(graph, index);

        normalize(matrix);

        for (int i = 0; i < ITERATIONS; i++) {
            final INDArray previous = matrix.unsafeDuplication();

            expand(matrix);
            inflate(matrix);

            if (matrix.equals(previous)) break;
        }
    }

    @Override
    public Collection<Collection<V>> getClusters() {
        if (graph.vertexSet().isEmpty()) return Collections.emptySet();

        final Map<Integer, V> inverted = index.entrySet().stream().
                collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));

        final Set<Collection<V>> clusters = new HashSet<>();

        for (int r = 0; r < matrix.rows(); r++) {
            final INDArray row = matrix.getRow(r);

            final Set<V> cluster = new HashSet<>();

            for (int c = 0; c < row.length(); c++) {
                if (row.getDouble(c) > 0) cluster.add(inverted.get(c));
            }

            if (!cluster.isEmpty()) clusters.add(cluster);
        }

        return clusters;
    }

    /**
     * Index the nodes of the input graph.
     *
     * @param graph an input graph.
     * @return a node index.
     */
    protected Map<V, Integer> buildIndex(Graph<V, E> graph) {
        final Map<V, Integer> index = new HashMap<>();

        int i = 0;

        for (final V vertex : graph.vertexSet()) {
            index.put(vertex, i++);
        }

        return index;
    }

    /**
     * Construct an adjacency matrix for a given graph.
     *
     * @param graph a graph.
     * @param index a node index for the graph.
     * @return an adjacency matrix.
     */
    protected INDArray buildMatrix(Graph<V, E> graph, Map<V, Integer> index) {
        final INDArray matrix = Nd4j.create(graph.vertexSet().size(), graph.vertexSet().size());

        for (final E edge : graph.edgeSet()) {
            final int i = index.get(graph.getEdgeSource(edge)), j = index.get(graph.getEdgeTarget(edge));
            matrix.put(i, j, graph.getEdgeWeight(edge));
            matrix.put(j, i, graph.getEdgeWeight(edge));
        }

        for (int r = 0; r < matrix.shape()[0]; r++) matrix.put(r, r, 1);

        return matrix;
    }

    protected void normalize(INDArray matrix) {
        final INDArray sums = matrix.sum(0);
        matrix.diviRowVector(sums);
    }

    protected void expand(INDArray matrix) {
        Transforms.mpow(matrix, e, false);
    }

    protected void inflate(INDArray matrix) {
        Nd4j.getExecutioner().exec(new Pow(matrix, r));
        normalize(matrix);
    }
}

@@ -3451,30 +3451,16 @@
 exSet())
- %7B%0A           
  index.p
@@ -3475,26 +3475,16 @@
 x, i++);
-%0A        %7D
 %0A%0A      
@@ -3565,24 +3565,57 @@
 given graph.
+ Note that the loops are ignored.
 %0A     *%0A    
@@ -3851,41 +3851,12 @@
 d4j.
-create(graph.vertexSet().size(), 
+eye(
 grap
@@ -4024,32 +4024,33 @@
 eTarget(edge));%0A
+%0A
             matr
@@ -4049,24 +4049,59 @@
     
-matrix.put(i, j,
+if (i != j) %7B%0A                final double weight =
  gra
@@ -4114,33 +4114,32 @@
 EdgeWeight(edge)
-)
 ;%0A            ma
@@ -4132,24 +4132,28 @@
             
+    
 matrix.put(j
@@ -4155,39 +4155,20 @@
 put(
-j, i, graph.getEdgeWeight(edge)
+i, j, weight
 );%0A 
@@ -4178,62 +4178,15 @@
     
-%7D%0A%0A
        
- for (int r = 0; r %3C matrix.shape()%5B0%5D; r++)
  mat
@@ -4197,17 +4197,46 @@
 put(
-r, r, 1);
+j, i, weight);%0A            %7D%0A        %7D
 %0A%0A  

Simplify matrix building in MCL, again