/*
 * 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 com.milaboratory.core.sequence.NucleotideSequence;
import com.milaboratory.test.TestUtil;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;

import static com.milaboratory.util.TempFileDest.folder;

public class UnorderedMergerTest {
    @Rule
    public TemporaryFolder temporaryFolder = new TemporaryFolder();

    public void checkMerger(List<List<String>> inputData, int memoryBudget, MergeStrategyTest.Substring... properties) throws IOException {
        // Building expected value
        Map<List<String>, List<List<String>>> expectedMap = new HashMap<>();
        for (int ci = 0; ci < inputData.size(); ci++)
            for (String cell : inputData.get(ci)) {
                List<String> key = Arrays.stream(properties).map(p -> p.get(cell)).collect(Collectors.toList());
                List<List<String>> row = expectedMap.computeIfAbsent(key, __ -> {
                    ArrayList<List<String>> grpRow = new ArrayList<>(inputData.size());
                    for (int i = 0; i < inputData.size(); i++)
                        grpRow.add(new ArrayList<>());
                    return grpRow;
                });
                row.get(ci).add(cell);
            }

        for (List<List<String>> row : expectedMap.values())
            for (List<String> c : row)
                c.sort(Comparator.naturalOrder());

        Set<List<List<String>>> expected = new HashSet<>(expectedMap.values());

        UnorderedMerger<String> merger = new UnorderedMerger<>(
                String.class,
                inputData.stream()
                        .map(CUtils::asOutputPort)
                        .collect(Collectors.toList()),
                Arrays.asList(properties),
                2,
                folder(temporaryFolder.newFolder(), false),
                2, 2,
                memoryBudget, 10);
        merger.loadData();

        for (List<List<String>> row : CUtils.it(merger)) {
            for (List<String> c : row)
                c.sort(Comparator.naturalOrder());
            Assert.assertTrue(expected.contains(row));
            Assert.assertTrue(expected.remove(row));
            // System.out.println(merger.getProgress());
        }

        Assert.assertTrue(expected.isEmpty());
    }

    @Test
    public void test1() throws IOException {
        MergeStrategyTest.Substring ss1 = new MergeStrategyTest.Substring(2, 4);
        MergeStrategyTest.Substring ss2 = new MergeStrategyTest.Substring(2, 5);

        List<String> strings1 = new ArrayList<>();
        strings1.add("ATTAGAAA");
        strings1.add("GACATAAA");
        strings1.add("TGCGAAAA");
        strings1.add("AGCGGAAA");
        strings1.add("AGCGGAAA");
        strings1.add("AGCGAAAA");
        strings1.add("AATTCAAA");

        List<String> strings2 = new ArrayList<>();
        strings2.add("ATTAGTTT");
        strings2.add("GACATTTT");
        strings2.add("TGCGATTT");
        strings2.add("AGCGGTTT");
        strings2.add("AGCGGTTT");
        strings2.add("AGCGATTT");
        strings2.add("AATTCTTT");

        checkMerger(Arrays.asList(strings1, strings2), 100, ss1);
        checkMerger(Arrays.asList(strings1, strings2), 100, ss2);
        checkMerger(Arrays.asList(strings1, strings2), 100, ss2, ss1);
    }

    @Test
    public void test2() throws IOException {
        int its = TestUtil.its(10, 100);
        ThreadLocalRandom r = ThreadLocalRandom.current();
        for (int __ = 0; __ < its; __++) {
            int sequenceLength = 5 + r.nextInt(15); // test with fixed sequence length

            int numberOfColumns = r.nextInt(10); // zero included
            List<List<String>> data = new ArrayList<>();
            for (int i = 0; i < numberOfColumns; i++) {
                ArrayList<String> col = new ArrayList<>();
                data.add(col);
                int numberOfRecordsInColumn = r.nextInt(10000); // zero included
                for (int i1 = 0; i1 < numberOfRecordsInColumn; i1++)
                    col.add(TestUtil.randomSequence(NucleotideSequence.ALPHABET, sequenceLength, sequenceLength).toString());
            }

            List<MergeStrategyTest.Substring> properties = new ArrayList<>();
            int numberOfProperties = 1 + r.nextInt(2);
            for (int i = 0; i < numberOfProperties; i++) {
                int len = 1 + r.nextInt(sequenceLength - 1);
                int from = r.nextInt(sequenceLength - len);
                properties.add(new MergeStrategyTest.Substring(from, from + len));
            }

            checkMerger(data, 100000, properties.toArray(new MergeStrategyTest.Substring[0]));
        }
    }
}