package org.tensorics.core.tensor.lang;

import static java.util.Collections.singleton;

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

import org.tensorics.core.lang.Tensorics;
import org.tensorics.core.tensor.Position;
import org.tensorics.core.tensor.Tensor;
import org.tensorics.core.tensor.Tensor.Entry;
import org.tensorics.core.tensor.operations.TensorInternals;

import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Sets;
import com.google.common.collect.Sets.SetView;

public class OngoingOrderedFlattening<S, C extends Comparable<C>> {

    private final Tensor<S> tensor;
    private final Class<C> dimensionToFlatten;

    public OngoingOrderedFlattening(Tensor<S> tensor, Class<C> dimension) {
        super();
        this.tensor = tensor;
        this.dimensionToFlatten = dimension;
    }

    public Tensor<List<S>> intoTensorOfLists() {
        return Tensorics.fromMap(remainingDimensions(), Multimaps.asMap(intoListMultimap()));
    }

    private ListMultimap<Position, S> intoListMultimap() {
        ImmutableListMultimap.Builder<Position, S> builder = ImmutableListMultimap.builder();
        Tensor<Map<C, S>> maps = TensorInternals.mapOut(tensor).inDirectionOf(dimensionToFlatten);
        for (Entry<Map<C, S>> entry : maps.entrySet()) {
            Position newPosition = entry.getPosition();
            Map<C, S> map = entry.getValue();
            List<C> keys = new ArrayList<>(map.keySet());
            Collections.sort(keys);
            for (C key : keys) {
                builder.put(newPosition, map.get(key));
            }
        }
        return builder.build();
    }

    private SetView<Class<?>> remainingDimensions() {
        return Sets.difference(tensor.shape().dimensionSet(), singleton(dimensionToFlatten));
    }
}

@@ -1780,32 +1780,796 @@
 r.build();%0A    %7D
+%0A    %0A    public Tensor%3CList%3CC%3E%3E intoTensorOfKeyLists() %7B%0A        return Tensorics.fromMap(remainingDimensions(), Multimaps.asMap(intoKeyListMultimap()));%0A    %7D%0A%0A    private ListMultimap%3CPosition, C%3E intoKeyListMultimap() %7B%0A        ImmutableListMultimap.Builder%3CPosition, C%3E builder = ImmutableListMultimap.builder();%0A        Tensor%3CMap%3CC, S%3E%3E maps = TensorInternals.mapOut(tensor).inDirectionOf(dimensionToFlatten);%0A        for (Entry%3CMap%3CC, S%3E%3E entry : maps.entrySet()) %7B%0A            Position newPosition = entry.getPosition();%0A            Map%3CC, S%3E map = entry.getValue();%0A            List%3CC%3E keys = new ArrayList%3C%3E(map.keySet());%0A            Collections.sort(keys);%0A            builder.putAll(newPosition, keys);%0A        %7D%0A        return builder.build();%0A    %7D
 %0A%0A    private Se

flattening into keyList