package no.ssb.vtl.script.operations;

import com.codepoetics.protonpack.StreamUtils;
import com.google.common.collect.Maps;
import no.ssb.vtl.model.AbstractUnaryDatasetOperation;
import no.ssb.vtl.model.Component;
import no.ssb.vtl.model.DataPoint;
import no.ssb.vtl.model.DataStructure;
import no.ssb.vtl.model.Dataset;
import no.ssb.vtl.model.Order;
import no.ssb.vtl.model.VTLNumber;
import no.ssb.vtl.model.VTLObject;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Stream;

import static com.google.common.base.Preconditions.*;

public class AggregationOperation extends AbstractUnaryDatasetOperation {

    private final List<Component> groupBy;
    private final Component aggregationComponent;
    private final Function<List<VTLNumber>, VTLNumber> aggregationFunction;

    public AggregationOperation(Dataset child, List<Component> groupBy, Component aggregationComponent, Function<List<VTLNumber>, VTLNumber> aggregationFunction) {
        super(child);
        this.groupBy = groupBy;
        this.aggregationComponent = checkNotNull(aggregationComponent);
        this.aggregationFunction = aggregationFunction;
        
    }

    @Override
    protected DataStructure computeDataStructure() {
        DataStructure.Builder newDataStructure = DataStructure.builder();
        for (Map.Entry<String, Component> entry : getChild().getDataStructure().entrySet()) {
            if (groupBy.contains(entry.getValue()) || aggregationComponent.equals(entry.getValue())) {
                newDataStructure.put(entry);
            }
        }
        return newDataStructure.build();
    }

    @Override
    public Stream<DataPoint> getData() {
        Order.Builder builder = Order.create(getDataStructure());
        groupBy.forEach(component -> builder.put(component, Order.Direction.ASC));
        Order order = builder.build();
        
        Stream<List<DataPoint>> groupedDataPoints = StreamUtils.aggregate(
                getChild().getData(order).orElse(getChild().getData().sorted(order)), (dataPoint, dataPoint2) -> {
                    Map<Component, VTLObject> map1 = getDataStructure().asMap(dataPoint);
                    Map<Component, VTLObject> map2 = getDataStructure().asMap(dataPoint2);
                    return groupBy.stream()
                            .allMatch(component -> map1.get(component).compareTo(map2.get(component)) == 0);
                });
    
        @SuppressWarnings("UnnecessaryLocalVariable")
        Stream<DataPoint> aggregatedDataPoints = groupedDataPoints.map(dataPoints -> {
            Map<Component, VTLObject> resultAsMap = Maps.newHashMap();
            List<VTLNumber> aggregationValues = new ArrayList<>();
            dataPoints.stream()
                    .map(dataPoint ->  getChild().getDataStructure().asMap(dataPoint).entrySet())
                    .flatMap(Collection::stream)
                    .forEach(entry -> {
                        if (entry.getKey().equals(aggregationComponent)) {
                            aggregationValues.add(VTLNumber.of((Number) entry.getValue().get()));
                        } else if (groupBy.contains(entry.getKey())) {
                            resultAsMap.put(entry.getKey(), entry.getValue());
                        }
                    });
            resultAsMap.put(aggregationComponent, aggregationFunction.apply(aggregationValues));
            return getDataStructure().fromMap(resultAsMap);
        });
        
        return aggregatedDataPoints;
    }

    /**
     * Returns the count of unique values by column.
     */
    @Override
    public Optional<Map<String, Integer>> getDistinctValuesCount() {
        return null;
    }

    /**
     * Return the amount of {@link DataPoint} the stream obtained by the
     * method {@link Dataset#getData()} will return.
     */
    @Override
    public Optional<Long> getSize() {
        return null;
    }
}

@@ -1975,28 +1975,24 @@
 uild();%0A    
-    
 %0A        Str
@@ -1995,21 +1995,16 @@
  Stream%3C
-List%3C
 DataPoin
@@ -2009,68 +2009,15 @@
 int%3E
-%3E groupedDataPoints = StreamUtils.aggregate(%0A               
+ data =
  get
@@ -2085,364 +2085,170 @@
 er))
-, (dataPoint, dataPoint2) -%3E %7B%0A                    Map%3CComponent, VTLObject%3E map1 = getDataStructure().asMap(dataPoint);%0A                    Map%3CComponent, VTLObject%3E map2 = getDataStructure().asMap(dataPoint2);%0A                    return groupBy.stream()%0A                            .allMatch(component -%3E map1.get(component).compareTo(map2.get(component)
+;%0A        Stream%3CList%3CDataPoint%3E%3E groupedDataPoints = StreamUtils.aggregate(data,%0A                (dataPoint1, dataPoint2) -%3E order.compare(dataPoint1, dataPoint2
 ) ==
@@ -2252,36 +2252,16 @@
  == 0);%0A
-                %7D);%0A
     %0A   

Use order to group dataPoints