/*
 * Grakn - A Distributed Semantic Database
 * Copyright (C) 2016-2018 Grakn Labs Limited
 *
 * Grakn is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Grakn is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Grakn. If not, see <http://www.gnu.org/licenses/gpl.txt>.
 */

package ai.grakn.graql.internal.query.analytics;

import ai.grakn.GraknTx;
import ai.grakn.concept.Label;
import ai.grakn.concept.LabelId;
import ai.grakn.graql.analytics.CountQuery;
import ai.grakn.graql.internal.analytics.CountMapReduce;
import ai.grakn.graql.internal.analytics.CountMapReduceWithAttribute;
import ai.grakn.graql.internal.analytics.CountVertexProgram;
import org.apache.tinkerpop.gremlin.process.computer.ComputerResult;

import java.util.Collection;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static ai.grakn.graql.internal.analytics.GraknMapReduce.RESERVED_TYPE_LABEL_KEY;

class CountQueryImpl extends AbstractComputeQuery<Long> implements CountQuery {

    CountQueryImpl(Optional<GraknTx> graph) {
        this.tx = graph;
    }

    @Override
    public Long execute() {
        LOGGER.info("CountMapReduce is called");
        long startTime = System.currentTimeMillis();

        initSubGraph();
        getAllSubTypes();

        if (!selectedTypesHaveInstance()) {
            LOGGER.debug("Count = 0");
            LOGGER.info("CountMapReduce is done in " + (System.currentTimeMillis() - startTime) + " ms");
            return 0L;
        }

        Set<LabelId> typeLabelIds = convertLabelsToIds(subLabels);
        Map<Integer, Long> count;

        if (includeAttribute) {
            Set<LabelId> rolePlayerLabelIds = getRolePlayerLabelIds();
            rolePlayerLabelIds.addAll(typeLabelIds);

            ComputerResult result = getGraphComputer().compute(
                    new CountVertexProgram(),
                    new CountMapReduceWithAttribute(),
                    rolePlayerLabelIds, false);
            count = result.memory().get(CountMapReduceWithAttribute.class.getName());
        } else {
            ComputerResult result = getGraphComputer().compute(
                    null,
                    new CountMapReduce(),
                    typeLabelIds, false);
            count = result.memory().get(CountMapReduce.class.getName());
        }

        long finalCount = count.keySet().stream()
                .filter(id -> typeLabelIds.contains(LabelId.of(id)))
                .mapToLong(count::get).sum();
        if (count.containsKey(RESERVED_TYPE_LABEL_KEY)) {
            finalCount += count.get(RESERVED_TYPE_LABEL_KEY);
        }

        LOGGER.debug("Count = " + finalCount);
        LOGGER.info("CountMapReduce is done in " + (System.currentTimeMillis() - startTime) + " ms");
        return finalCount;
    }

    @Override
    public boolean isReadOnly() {
        return true;
    }

    @Override
    public CountQuery in(String... subTypeLabels) {
        return (CountQuery) super.in(subTypeLabels);
    }

    @Override
    public CountQuery in(Collection<Label> subLabels) {
        return (CountQuery) super.in(subLabels);
    }

    @Override
    String graqlString() {
        return "count" + subtypeString();
    }

    @Override
    public CountQuery withTx(GraknTx tx) {
        return (CountQuery) super.withTx(tx);
    }
}

@@ -918,65 +918,8 @@
 ry;%0A
-import ai.grakn.graql.internal.analytics.CountMapReduce;%0A
 impo
@@ -1536,25 +1536,20 @@
 ount
-MapReduce is call
+ query start
 ed%22)
@@ -1768,33 +1768,31 @@
 o(%22Count
-MapReduce is done
+ query finished
  in %22 + 
@@ -1982,44 +1982,8 @@
 t;%0A%0A
-        if (includeAttribute) %7B%0A    
     
@@ -2045,36 +2045,32 @@
 lIds();%0A        
-    
 rolePlayerLabelI
@@ -2095,36 +2095,32 @@
 lIds);%0A%0A        
-  
-  
 ComputerResult r
@@ -2147,36 +2147,32 @@
 uter().compute(%0A
-    
                 
@@ -2205,36 +2205,32 @@
                 
-    
 new CountMapRedu
@@ -2260,28 +2260,24 @@
             
- 
-   
 rolePlayerLa
@@ -2284,36 +2284,32 @@
 belIds, false);%0A
-    
         count = 
@@ -2377,282 +2377,8 @@
 ());
-%0A        %7D else %7B%0A            ComputerResult result = getGraphComputer().compute(%0A                    null,%0A                    new CountMapReduce(),%0A                    typeLabelIds, false);%0A            count = result.memory().get(CountMapReduce.class.getName());%0A        %7D
 %0A%0A  
@@ -2748,25 +2748,23 @@
 ount
-MapReduce is done
+ query finished
  in 
@@ -2856,84 +2856,8 @@
  %7D%0A%0A
-    @Override%0A    public boolean isReadOnly() %7B%0A        return true;%0A    %7D%0A%0A
     

Make analytics count return the same result as aggregate count (#2584)