package org.psjava.algo.graph;

import org.psjava.algo.graph.dfs.DFSVisitorBase;
import org.psjava.algo.graph.dfs.SingleSourceDFS;
import org.psjava.ds.array.DynamicArray;
import org.psjava.ds.graph.DirectedEdge;
import org.psjava.ds.graph.RootedTree;
import org.psjava.ds.map.MutableMap;
import org.psjava.ds.map.MutableMapFactory;
import org.psjava.ds.numbersystrem.AddInvert;
import org.psjava.ds.numbersystrem.AddableNumberSystem;
import org.psjava.ds.tree.segmenttree.SegmentTree;
import org.psjava.ds.tree.segmenttree.SegmentTreeByArrayImplementation;
import org.psjava.util.VisitorStopper;

import java.util.function.Function;

// TODO this is not used any problem. solve something to prove this.
public class DistanceCalculatorInRootedTree {

    private final LowestCommonAncestorAlgorithm lca;
    private final MutableMapFactory mapFactory;

    public DistanceCalculatorInRootedTree(LowestCommonAncestorAlgorithm lca, MutableMapFactory mapFactory) {
        this.lca = lca;
        this.mapFactory = mapFactory;
    }

    public <V, W, E extends DirectedEdge<V>> DistanceCalculatorInRootedTreeResult<V, W> calc(RootedTree<V, E> tree, final Function<E, W> weight, final AddableNumberSystem<W> ns) {
        final DynamicArray<W> pathWeights = DynamicArray.create();
        final MutableMap<V, Integer> discoverIndex = mapFactory.create();
        final MutableMap<V, Integer> indexOfWalkingDown = mapFactory.create();
        final MutableMap<V, Integer> indexOfWalkingUp = mapFactory.create();

        SingleSourceDFS.traverse(tree.graph, tree.root, new DFSVisitorBase<V, E>() {

            @Override
            public void onDiscovered(V vertex, int depth, VisitorStopper stopper) {
                int index = pathWeights.size();
                discoverIndex.add(vertex, index);
            }

            @Override
            public void onWalkDown(E outEdge) {
                int index = pathWeights.size();
                indexOfWalkingDown.add(outEdge.to(), index);
                pathWeights.addToLast(weight.apply(outEdge));
            }

            @Override
            public void onWalkUp(E edge) {
                int index = pathWeights.size();
                indexOfWalkingUp.add(edge.to(), index);
                pathWeights.addToLast(AddInvert.calc(ns, weight.apply(edge)));
            }
        });

        final SegmentTree<W> segmentTree = new SegmentTreeByArrayImplementation<>(pathWeights, ns::add);
        final LowestCommonAncestorQuerySession<V> lcaSession = lca.calc(tree);

        return new DistanceCalculatorInRootedTreeResult<V, W>() {
            @Override
            public W getDistance(V v1, V v2) {
                V ancestor = lcaSession.query(v1, v2);
                W r = ns.getZero();
                if (!v1.equals(ancestor))
                    r = ns.add(r, segmentTree.query(discoverIndex.get(ancestor), discoverIndex.get(v1)));
                if (!v2.equals(ancestor))
                    r = ns.add(r, segmentTree.query(discoverIndex.get(ancestor), discoverIndex.get(v2)));
                return r;
            }

            @Override
            public void modifyDistance(V v1, V v2, W w) {
                V child = (discoverIndex.get(v1) < discoverIndex.get(v2)) ? v2 : v1;
                segmentTree.update(indexOfWalkingDown.get(child), w);
                segmentTree.update(indexOfWalkingUp.get(child), AddInvert.calc(ns, w));
            }
        };
    }

}
@@ -128,49 +128,8 @@
 FS;%0A
-import org.psjava.ds.array.DynamicArray;%0A
 impo
@@ -165,16 +165,16 @@
 edEdge;%0A
+
 import o
@@ -550,16 +550,67 @@
 opper;%0A%0A
+import java.util.ArrayList;%0Aimport java.util.List;%0A
 import j
@@ -1228,28 +1228,20 @@
   final 
-DynamicArray
+List
 %3CW%3E path
@@ -1254,27 +1254,23 @@
 s = 
-DynamicArray.create
+new ArrayList%3C%3E
 ();%0A
@@ -2015,22 +2015,16 @@
 ghts.add
-ToLast
 (weight.
@@ -2216,32 +2216,32 @@
 e.to(), index);%0A
+
                 
@@ -2259,14 +2259,8 @@
 .add
-ToLast
 (Add

change deprecated usage - SegmentTreeByArrayImplementation() in DistanceCalculatorInRootedTree