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<DataPoint> data = getChild().getData(order).orElse(getChild().getData().sorted(order));
        Stream<List<DataPoint>> groupedDataPoints = StreamUtils.aggregate(data,
                (dataPoint1, dataPoint2) -> order.compare(dataPoint1, dataPoint2) == 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;
    }
}

@@ -83,47 +83,8 @@
 ls;%0A
-import com.google.common.collect.Maps;%0A
 impo
@@ -130,24 +130,24 @@
 tOperation;%0A
+
 import no.ss
@@ -382,65 +382,8 @@
 t;%0A%0A
-import java.util.ArrayList;%0Aimport java.util.Collection;%0A
 impo
@@ -482,24 +482,60 @@
 n.Function;%0A
+import java.util.stream.Collectors;%0A
 import java.
@@ -1723,24 +1723,153 @@
 getData() %7B%0A
+        DataStructure childStructure = getChild().getDataStructure();%0A        DataStructure structure = getDataStructure();%0A    %0A
         Orde
@@ -1901,32 +1901,25 @@
 .create(
-getDataS
+s
 tructure
 ());%0A   
@@ -1910,18 +1910,16 @@
 tructure
-()
 );%0A     
@@ -2454,24 +2454,107 @@
 Points -%3E %7B%0A
+            DataPoint firstDataPointOfGroup = DataPoint.create(dataPoints.get(0));%0A
             
@@ -2597,26 +2597,61 @@
 p = 
-Maps.newH
+childStructure.
 as
-h
 Map(
-);
+firstDataPointOfGroup);%0A        
 %0A   
@@ -2698,39 +2698,8 @@
 es =
- new ArrayList%3C%3E();%0A           
  dat
@@ -2753,35 +2753,21 @@
 oint -%3E 
- getChild().getData
+child
 Structur
@@ -2767,18 +2767,16 @@
 tructure
-()
 .asMap(d
@@ -2789,158 +2789,11 @@
 nt).
-entrySet())%0A                    .flatMap(Collection::stream)%0A                    .forEach(entry -%3E %7B%0A                        if (entry.getKey().equals
+get
 (agg
@@ -2811,18 +2811,16 @@
 ponent))
- %7B
 %0A       
@@ -2836,38 +2836,26 @@
     
-        aggregationValues.add(
+.map(vtlObject -%3E 
 VTLN
@@ -2876,32 +2876,25 @@
 er) 
-entry.getValue()
+vtlObject
 .get()))
 ;%0A  
@@ -2889,17 +2889,16 @@
 .get()))
-;
 %0A       
@@ -2914,187 +2914,47 @@
     
-    %7D else if (groupBy.contains(entry.getKey())) %7B%0A                            resultAsMap.put(entry.getKey(), entry.getValue());%0A                        %7D%0A                    %7D);
+.collect(Collectors.toList());%0A        
 %0A   
@@ -3066,32 +3066,25 @@
  return 
-getDataS
+s
 tructure
 ().fromM
@@ -3075,18 +3075,16 @@
 tructure
-()
 .fromMap

Fix potential bug in creation of aggregated DataPoint and utilize streams better to clean up a little