/**
 *
 * Copyright (c) Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See License.txt in the project root for license information.
 *
 */

package com.microsoft.azure;

import org.apache.commons.lang3.tuple.MutablePair;
import rx.Observable;
import rx.functions.Action1;
import rx.functions.Func1;
import rx.functions.Func2;
import rx.functions.FuncN;
import rx.schedulers.Schedulers;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * The base implementation of TaskGroup interface.
 *
 * @param <T> the result type of the tasks in the group
 * @param <U> the task item
 */
public abstract class TaskGroupBase<T, U extends TaskItem<T>>
    implements TaskGroup<T, U> {
    private DAGraph<U, DAGNode<U>> dag;

    /**
     * Creates TaskGroupBase.
     *
     * @param rootTaskItemId the id of the root task in this task group
     * @param rootTaskItem the root task
     */
    public TaskGroupBase(String rootTaskItemId, U rootTaskItem) {
        this.dag = new DAGraph<>(new DAGNode<>(rootTaskItemId, rootTaskItem));
    }

    @Override
    public DAGraph<U, DAGNode<U>> dag() {
        return dag;
    }

    @Override
    public boolean isPreparer() {
        return dag.isPreparer();
    }

    @Override
    public void merge(TaskGroup<T, U> parentTaskGroup) {
        dag.merge(parentTaskGroup.dag());
    }

    @Override
    public void prepare() {
        if (isPreparer()) {
            dag.prepare();
        }
    }

    @Override
    public void execute() throws Exception {
        DAGNode<U> nextNode = dag.getNext();
        while (nextNode != null) {
            nextNode.data().execute();
            this.dag().reportedCompleted(nextNode);
            nextNode = dag.getNext();
        }
    }

    @Override
    public Observable<T> executeAsync() {
        return executeReadyTasksAsync();
    }

    @Override
    public T taskResult(String taskId) {
        return dag.getNodeData(taskId).result();
    }

    /**
     * Executes all runnable tasks, a task is runnable when all the tasks its depends
     * on are finished running.
     */
    private Observable<T> executeReadyTasksAsync() {
        DAGNode<U> nextNode = dag.getNext();
        Observable<T> rootObservable = null;
        final List<Observable<T>> observables = new ArrayList<>();
        while (nextNode != null) {
            final DAGNode<U> thisNode = nextNode;
            if (dag().isRootNode(nextNode)) {
                rootObservable = nextNode.data().executeAsync()
                        .doOnNext(new Action1<T>() {
                            @Override
                            public void call(T t) {
                                dag().reportedCompleted(thisNode);
                            }
                        });
            } else {
                Observable<T> nextNodeObservable = nextNode.data().executeAsync()
                        .flatMap(new Func1<T, Observable<T>>() {
                            @Override
                            public Observable<T> call(T t) {
                                dag().reportedCompleted(thisNode);
                                return executeReadyTasksAsync();
                            }
                        });
                observables.add(nextNodeObservable);
            }
            nextNode = dag.getNext();
        }
        if (rootObservable != null) {
            return rootObservable;
        }
        else {
            return Observable.merge(observables).last();
        }
    }
}

@@ -206,315 +206,100 @@
 ort 
-org.apache.commons.lang3.tuple.MutablePair;%0Aimport rx.Observable;%0Aimport rx.functions.Action1;%0Aimport rx.functions.Func1;%0Aimport rx.functions.Func2;%0Aimport rx.functions.FuncN;%0Aimport rx.schedulers.Schedulers;%0A%0Aimport java.util.ArrayList;%0Aimport java.util.HashMap;%0Aimport java.util.List;%0Aimport java.util.Map
+rx.Observable;%0Aimport rx.functions.Func1;%0A%0Aimport java.util.ArrayList;%0Aimport java.util.List
 ;%0A%0A/
@@ -2038,53 +2038,8 @@
 ();%0A
-        Observable%3CT%3E rootObservable = null;%0A
     
@@ -2202,63 +2202,9 @@
     
-if (dag().isRootNode(nextNode)) %7B%0A                rootO
+o
 bser
@@ -2200,35 +2200,38 @@
       observable
- = 
+s.add(
 nextNode.data().
@@ -2269,45 +2269,49 @@
     
-    .doOnNext(new Action1%3CT
+.flatMap(new Func1%3CT, Observable%3CT%3E
 %3E() %7B%0A
-    
     
@@ -2364,31 +2364,36 @@
         
-    
 public 
-void
+Observable%3CT%3E
  call(T 
@@ -2389,36 +2389,32 @@
 %3CT%3E call(T t) %7B%0A
-    
                 
@@ -2492,371 +2492,137 @@
     
-%7D%0A                        %7D);%0A            %7D else %7B%0A                Observable%3CT%3E nextNodeObservable = nextNode.data().executeAsync()%0A                        .flatMap(new Func1%3CT, Observable%3CT%3E%3E() %7B%0A                            @Override%0A                            public Observable%3CT%3E call(T t) %7B%0A                                dag().reportedCompleted(thisNode);
+if (dag().isRootNode(thisNode)) %7B%0A                                return Observable.just(t);%0A                            %7D else %7B
 %0A   
@@ -2738,18 +2738,16 @@
        %7D
-);
 %0A       
@@ -2759,58 +2759,16 @@
     
-observables.add(nextNodeObservable);%0A        
     %7D
+));
 %0A   
@@ -2816,110 +2816,8 @@
   %7D%0A
-        if (rootObservable != null) %7B%0A            return rootObservable;%0A        %7D%0A        else %7B%0A    
     
@@ -2861,30 +2861,20 @@
 es).last();%0A
-        %7D%0A
     %7D%0A%7D%0A

Create VM works