/*
 * Copyright (c) 2022 MiLaboratories Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
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.*;

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,
            TempFileDest tempDest,
            int readerConcurrency,
            int writerConcurrency,
            long memoryBudget,
            long objectSizeInitialGuess) {
        this(clazz, inputs, by, bitsPerStep, tempDest,
                readerConcurrency, writerConcurrency, memoryBudget,
                objectSizeInitialGuess, new PrimitivIOStateBuilder());
    }

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

        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,
                tempDest,
                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();
    }
}
