package com.milaboratory.util.sorting;

import cc.redberry.pipe.CUtils;
import cc.redberry.pipe.OutputPort;
import cc.redberry.pipe.OutputPortCloseable;
import cc.redberry.pipe.util.CountingOutputPort;
import cc.redberry.pipe.util.FlatteningOutputPort;
import com.milaboratory.primitivio.PrimitivIOStateBuilder;
import com.milaboratory.util.CanReportProgress;
import com.milaboratory.util.CanReportProgressAndStage;
import com.milaboratory.util.ProgressAndStage;
import com.milaboratory.util.SmartProgressReporter;

import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;

public final class UnorderedMerger<T> implements OutputPortCloseable<List<List<T>>>, CanReportProgressAndStage {
    private final Object sync = new Object();
    private final List<OutputPort<T>> inputs;
    private final int inputCount;
    private final HashSorter<WithIndex<T>> hashSorter;
    private final Comparator<WithIndex<T>> effectiveComparator;

    private final ProgressAndStage ps = new ProgressAndStage("Initialization");

    private OutputPortCloseable<WithIndex<T>> sorted;
    private WithIndex<T> lastElement;

    public UnorderedMerger(
            Class<T> clazz,
            List<OutputPort<T>> inputs,
            List<? extends SortingProperty<? super T>> by,
            int bitsPerStep,
            Path tempFilePrefix,
            int readerConcurrency,
            int writerConcurrency,
            long memoryBudget,
            long objectSizeInitialGuess) {
        this.inputs = inputs;
        this.inputCount = inputs.size();

        PrimitivIOStateBuilder stateBuilder = new PrimitivIOStateBuilder();
        stateBuilder.registerCustomSerializer(WithIndex.class, new WithIndex.WISerializer<>(clazz));
        @SuppressWarnings("unchecked")
        List<MappedSortingProperty<WithIndex<T>, T>> mappedBy = by.stream()
                .map(prop -> new MappedSortingProperty<WithIndex<T>, T>((SortingProperty<T>) prop, wi -> wi.obj))
                .collect(Collectors.toList());
        //noinspection unchecked,rawtypes
        hashSorter = new HashSorter<WithIndex<T>>(
                (Class) WithIndex.class,
                HashUtils.hashFunctionFor(mappedBy), SortingUtil.combine(mappedBy),
                bitsPerStep,
                tempFilePrefix,
                readerConcurrency, writerConcurrency,
                stateBuilder.getOState(), stateBuilder.getIState(),
                memoryBudget, objectSizeInitialGuess + 4
        );
        this.effectiveComparator = hashSorter.getEffectiveComparator();
    }

    public void loadData() {
        ps.setStage("Reading data for merge");
        ps.delegate(SmartProgressReporter.combinedProgress(inputs.stream()
                .filter(i -> i instanceof CanReportProgress)
                .map(i -> (CanReportProgress) i)
                .collect(Collectors.toList())
        ));

        List<OutputPort<WithIndex<T>>> indexedInputs = new ArrayList<>(inputs.size());
        for (int i = 0; i < inputs.size(); i++) {
            int fi = i;
            indexedInputs.add(CUtils.wrap(inputs.get(i), o -> new WithIndex<>(fi, o)));
        }
        FlatteningOutputPort<WithIndex<T>> flattened = new FlatteningOutputPort<>(CUtils.asOutputPort(indexedInputs));
        CountingOutputPort<WithIndex<T>> hsInput = new CountingOutputPort<>(flattened);
        OutputPortCloseable<WithIndex<T>> sortedOutput = hashSorter.port(hsInput);
        CountingOutputPort<WithIndex<T>> countedSorted = new CountingOutputPort<>(sortedOutput);
        this.sorted = countedSorted;

        ps.setStage("Merging records");
        ps.delegate(SmartProgressReporter.extractProgress(countedSorted, hsInput.getCount()));

        this.lastElement = sorted.take();
    }

    @Override
    public List<List<T>> take() {
        synchronized (sync) {
            if (sorted == null)
                throw new IllegalStateException("Please call UnorderedMerger.loadData() before using it as OutputPort.");
            WithIndex<T> lastElement;
            WithIndex<T> ref = lastElement = this.lastElement;
            if (lastElement == null) {
                ps.finish();
                return null;
            }
            List<List<T>> result = new ArrayList<>(inputCount);
            for (int i = 0; i < inputCount; i++)
                //noinspection unchecked
                result.add(Collections.EMPTY_LIST);
            do {
                List<T> grp = result.get(lastElement.index);
                if (grp == Collections.EMPTY_LIST)
                    result.set(lastElement.index, grp = new ArrayList<>());
                grp.add(lastElement.obj);
            } while ((lastElement = sorted.take()) != null &&
                    effectiveComparator.compare(lastElement, ref) == 0);
            this.lastElement = lastElement;
            return result;
        }
    }

    @Override
    public void close() {
        sorted.close();
    }

    @Override
    public String getStage() {
        return ps.getStage();
    }

    @Override
    public double getProgress() {
        return ps.getProgress();
    }

    @Override
    public boolean isFinished() {
        return ps.isFinished();
    }
}

@@ -1593,144 +1593,685 @@
 this
-.inputs = inputs;%0A        this.inputCount = inputs.size
+(clazz, inputs, by, bitsPerStep, tempFilePrefix,%0A                readerConcurrency, writerConcurrency, memoryBudget,%0A                objectSizeInitialGuess, new PrimitivIOStateBuilder
 ()
+)
 ;%0A
-%0A        PrimitivIOStateBuilder stateBuilder = new PrimitivIOStateBuilder
+    %7D%0A%0A    public UnorderedMerger(%0A            Class%3CT%3E clazz,%0A            List%3COutputPort%3CT%3E%3E inputs,%0A            List%3C? extends SortingProperty%3C? super T%3E%3E by,%0A            int bitsPerStep,%0A            Path tempFilePrefix,%0A            int readerConcurrency,%0A            int writerConcurrency,%0A            long memoryBudget,%0A            long objectSizeInitialGuess,%0A            PrimitivIOStateBuilder stateBuilder) %7B%0A        this.inputs = inputs;%0A        this.inputCount = inputs.size
 ();%0A
+%0A
     

Add PrimitivIOState optional parameter to UnirderedMerger.