package htsjdk.beta.codecs.reads.cram;

import com.google.common.jimfs.Configuration;
import com.google.common.jimfs.Jimfs;
import htsjdk.HtsjdkTest;
import htsjdk.beta.io.IOPathUtils;
import htsjdk.beta.plugin.HtsVersion;
import htsjdk.beta.plugin.interval.HtsIntervalUtils;
import htsjdk.beta.plugin.interval.HtsQueryRule;
import htsjdk.beta.plugin.reads.ReadsDecoder;
import htsjdk.beta.plugin.reads.ReadsDecoderOptions;
import htsjdk.beta.plugin.registry.HtsDefaultRegistry;
import htsjdk.io.HtsPath;
import htsjdk.io.IOPath;
import htsjdk.samtools.CRAMCRAIIndexer;
import htsjdk.samtools.QueryInterval;
import htsjdk.samtools.SAMRecord;
import htsjdk.samtools.SAMSequenceDictionary;
import htsjdk.samtools.SamFiles;
import htsjdk.samtools.ValidationStringency;
import htsjdk.samtools.seekablestream.SeekableFileStream;
import htsjdk.samtools.util.CloseableIterator;
import org.testng.Assert;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.FileSystem;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.util.Iterator;
import java.util.Optional;
import java.util.function.Function;

// These tests are derived from the tests in CRAMIndexQueryTest, which are used to test the (combined)
// CRAM 2.1/30 CRAMFileReader implementation. The tests here are modified to use the plugin framework APIs.
// Any tests in CRAMIndexQueryTest that use a CRAM 2.1 test file are disabled here, since this file is
// testing the 3.0 codec implementation.
public class HtsCRAMCodec30And21QueryTest extends HtsjdkTest {
    private static final IOPath TEST_DATA_DIR = new HtsPath("src/test/resources/htsjdk/samtools/cram/");

    private static final IOPath cramQueryWithBAI = new HtsPath(TEST_DATA_DIR + "cramQueryWithBAI.cram");
    private static final IOPath cramQueryWithCRAI = new HtsPath(TEST_DATA_DIR + "cramQueryWithCRAI.cram");
    private static IOPath cramQueryWithLocalCRAI = null; // generated  by @BeforeClass from cramQueryWithCRAI
    private static final IOPath cramQueryReference = new HtsPath(TEST_DATA_DIR + "human_g1k_v37.20.21.10M-10M200k.fasta");

    private static final IOPath cramQueryReadsWithBAI = new HtsPath(TEST_DATA_DIR +  "cramQueryTest.cram");
    private static IOPath cramQueryReadsWithLocalCRAI = null; // generated  by @BeforeClass from cramQueryReadsWithBAI

    private static final IOPath cramQueryTestEmptyWithBAI = new HtsPath(TEST_DATA_DIR + "cramQueryTestEmpty.cram");
    private static IOPath cramQueryTestEmptyWithLocalCRAI = null; // generated  by @BeforeClass from cramQueryTestEmptyWithBAI
    private static final IOPath cramQueryReadsReference = new HtsPath(TEST_DATA_DIR + "../hg19mini.fasta");

    // htsjdk currently generates .bai index files instead of .crai due to https://github.com/samtools/htsjdk/issues/531;
    // however we still want to test the htsjdk CRAI creation code. In order to run these index query tests against CRAI
    // files generated by htsjdk, we copy the checked-in test files into new CRAMs  and then manually generate .CRAI files
    // to run to use as additional test cases
    @BeforeClass
    public void createLocallyGeneratedCRAIFiles() throws IOException {
        cramQueryWithLocalCRAI = IOPathUtils.createTempPath("cramQueryWithLocalCRAI.", ".cram");
        IOPath tempCRAIOut = new HtsPath(cramQueryWithLocalCRAI.getURIString() + ".crai");
        tempCRAIOut.toPath().toFile().deleteOnExit();
        createLocalCRAMAndCRAI(
                cramQueryWithCRAI,
                cramQueryWithLocalCRAI,
                tempCRAIOut
        );

        cramQueryReadsWithLocalCRAI = IOPathUtils.createTempPath("cramQueryReadsWithLocalCRAI.", ".cram");
        tempCRAIOut = new HtsPath(cramQueryReadsWithLocalCRAI.getURIString() + ".crai");
        tempCRAIOut.toPath().toFile().deleteOnExit();
        cramQueryReadsWithLocalCRAI.toPath().toFile().deleteOnExit();
        createLocalCRAMAndCRAI(
                cramQueryReadsWithBAI,
                cramQueryReadsWithLocalCRAI,
                tempCRAIOut
        );

        cramQueryTestEmptyWithLocalCRAI = IOPathUtils.createTempPath("cramQueryTestEmptyWithLocalCRAI.", ".cram");
        tempCRAIOut = new HtsPath(cramQueryTestEmptyWithLocalCRAI.getURIString() +".crai");
        tempCRAIOut.toPath().toFile().deleteOnExit();
        cramQueryTestEmptyWithLocalCRAI.toPath().toFile().deleteOnExit();
        createLocalCRAMAndCRAI(
                cramQueryTestEmptyWithBAI,
                cramQueryTestEmptyWithLocalCRAI,
                tempCRAIOut
        );
    }

    private void createLocalCRAMAndCRAI(
            final IOPath inputCRAM,
            final IOPath outputCRAM,
            final IOPath outputCRAI) throws IOException
    {
        Files.copy(inputCRAM.toPath(), outputCRAM.toPath(), StandardCopyOption.REPLACE_EXISTING);
        try (final FileOutputStream bos = new FileOutputStream(outputCRAI.toPath().toFile());
             final SeekableFileStream sfs = new SeekableFileStream(outputCRAM.toPath().toFile())) {
            CRAMCRAIIndexer.writeIndex(sfs, bos);
        }
    }

    @DataProvider(name = "singleIntervalOverlapping")
    public Object[][] singleIntervalOverlapping() {
        return new Object[][] {
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 100009, 100009), new String[]{"a", "b", "c"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 100009, 100009), new String[]{"a", "b", "c"}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 100009, 100009), new String[]{"a", "b", "c"}},

                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 100009, 100011), new String[]{"a", "b", "c", "d", "e"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 100009, 100011), new String[]{"a", "b", "c", "d", "e"}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 100009, 100011), new String[]{"a", "b", "c", "d", "e"}},

                // interval with 1 start
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 1, 100010), new String[]{"a", "b", "c", "d"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 1, 100010), new String[]{"a", "b", "c", "d"}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 1, 100010), new String[]{"a", "b", "c", "d"}},

                // interval with 0 end
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 100015, 0), new String[]{"a", "b", "c", "d", "e", "f"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 100015, 0), new String[]{"a", "b", "c", "d", "e", "f"}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 100015, 0), new String[]{"a", "b", "c", "d", "e", "f"}},

                // interval with 1 start and 0 end
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 1, 0), new String[]{"a", "b", "c", "d", "e", "f",  "f"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 1, 0), new String[]{"a", "b", "c", "d", "e", "f",  "f"}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 1, 0), new String[]{"a", "b", "c", "d", "e", "f",  "f"}},

                //totally empty cram file
                {cramQueryTestEmptyWithBAI, cramQueryReadsReference, new QueryInterval(0, 1, 0), new String[]{}},
                {cramQueryTestEmptyWithLocalCRAI, cramQueryReadsReference, new QueryInterval(0, 1, 0), new String[]{}},

                { new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.cram"),
                        new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.fa"),
                        new QueryInterval(0, 631, 631),
                        new String[] {
                                "IL29_4505:7:51:11752:16255#2",
                                "IL29_4505:7:30:11521:4492#2",
                                "IL29_4505:7:78:5750:11942#2",
                                "IL29_4505:7:69:14718:6562#2",
                                "IL29_4505:7:88:5383:14756#2",
                                "IL29_4505:7:51:13987:8014#2",
                                "IL29_4505:7:26:6292:20619#2",
                                "IL29_4505:7:6:12278:3566#2",
                                "IL29_4505:7:7:7226:5753#2",
                                "IL29_4505:7:2:18466:8416#2",
                                "IL29_4505:7:42:15538:18597#2",
                                "IL29_4505:7:76:8383:16887#2",
                                "IL29_4505:7:114:15293:18740#2",
                                "IL29_4505:7:2:12262:7582#2",
                                "IL29_4505:7:29:7803:4644#2",
                                "IL29_4505:7:29:7816:4622#2",
                                "IL29_4505:7:40:16979:10045#2",
                                "IL29_4505:7:44:5928:2053#2",
                                "IL29_4505:7:93:5878:1091#2",
                                "IL29_4505:7:108:3637:10795#2",
                                "IL29_4505:7:13:6479:1635#2",
                                "IL29_4505:7:3:12568:11445#2",
                                "IL29_4505:7:32:8719:10919#2",
                                "IL29_4505:7:51:15171:6725#2",
                                "IL29_4505:7:112:9674:17564#2"
                        },
                }
        };
    }

    @Test(dataProvider="singleIntervalOverlapping")
    public void testQueryOverlappingHtsInterval(
            final IOPath cramFileName,
            final IOPath referenceFileName,
            final QueryInterval interval,
            final String[] expectedNames)
    {
        doQueryTest(
                decoder -> {
                    final SAMSequenceDictionary samDictionary = decoder.getHeader().getSequenceDictionary();
                    return decoder.queryOverlapping(
                            HtsIntervalUtils.fromQueryIntervalArray(new QueryInterval[]{interval}, samDictionary)
                    );
                },
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @Test(dataProvider="singleIntervalOverlapping")
    public void testQueryOverlappingIntervalComponents(
            final IOPath cramFileName,
            final IOPath referenceFileName,
            final QueryInterval interval,
            final String[] expectedNames)
    {
        doQueryTest(
                decoder -> decoder.queryOverlapping(
                        decoder.getHeader().getSequence(interval.referenceIndex).getSequenceName(),
                        interval.start,
                        interval.end
                ),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @Test(dataProvider="singleIntervalOverlapping")
    public void testQueryWithRule(
            final IOPath cramFileName,
            final IOPath referenceFileName,
            final QueryInterval interval,
            final String[] expectedNames)
    {
        doQueryTest(
                decoder -> {
                    final SAMSequenceDictionary samDictionary = decoder.getHeader().getSequenceDictionary();
                    return decoder.query(
                            HtsIntervalUtils.fromQueryIntervalArray(new QueryInterval[]{interval}, samDictionary),
                            HtsQueryRule.OVERLAPPING);
                },
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @DataProvider(name = "singleIntervalContained")
    public Object[][] singleIntervalContained() {
        return new Object[][] {
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 100013, 100070), new String[]{"f", "f",}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 100013, 100070), new String[]{"f", "f",}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 100013, 100070), new String[]{"f", "f"}},

                // interval with 1 start
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 1, 100100), new String[]{"e", "f", "f"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 1, 100100), new String[]{"e", "f", "f"}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 1, 100100), new String[]{"e", "f", "f"}},

                // interval with 0 end
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 100010, 0), new String[]{"d", "e", "f", "f",}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 100010, 0), new String[]{"d", "e", "f", "f",}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 100010, 0), new String[]{"d", "e", "f", "f",}},

                // interval with 1 start and 0 end
                {cramQueryWithCRAI, cramQueryReference, new QueryInterval(0, 1, 0), new String[]{"a", "b", "c", "d", "e", "f",  "f"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new QueryInterval(0, 1, 0), new String[]{"a", "b", "c", "d", "e", "f",  "f"}},
                {cramQueryWithBAI, cramQueryReference, new QueryInterval(0, 1, 0), new String[]{"a", "b", "c", "d", "e", "f",  "f"}},

                //totally empty cram file
                {cramQueryTestEmptyWithBAI, cramQueryReadsReference, new QueryInterval(0, 1, 0), new String[]{}},
                {cramQueryTestEmptyWithLocalCRAI, cramQueryReadsReference, new QueryInterval(0, 1, 0), new String[]{}},

                { new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.cram"),
                        new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.fa"),
                        new QueryInterval(0, 631, 656),
                        new String[] {"IL29_4505:7:88:5383:14756#2"},
                }
        };
    }

    @Test(dataProvider="singleIntervalContained")
    public void testQueryContainedHtsInterval(
            final HtsPath cramFileName,
            final HtsPath referenceFileName,
            final QueryInterval interval,
            final String[] expectedNames)
    {
        doQueryTest(
                decoder -> decoder.queryContained(
                        HtsIntervalUtils.fromQueryIntervalArray(new QueryInterval[]{interval},
                                decoder.getHeader().getSequenceDictionary())),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @Test(dataProvider="singleIntervalContained")
    public void testQueryContainedIntervalComponents(
            final HtsPath cramFileName,
            final HtsPath referenceFileName,
            final QueryInterval interval,
            final String[] expectedNames)
    {
        doQueryTest(
                decoder -> decoder.queryContained(
                        decoder.getHeader().getSequence(interval.referenceIndex).getSequenceName(),
                        interval.start,
                        interval.end
                ),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @Test(dataProvider="singleIntervalContained")
    public void testQueryWithRule(
            final HtsPath cramFileName,
            final HtsPath referenceFileName,
            final QueryInterval interval,
            final String[] expectedNames)
    {
        doQueryTest(
                decoder -> decoder.query(
                        HtsIntervalUtils.fromQueryIntervalArray(
                                new QueryInterval[]{interval}, decoder.getHeader().getSequenceDictionary()
                        ),
                        HtsQueryRule.CONTAINED),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @DataProvider(name = "multipleIntervalOverlapping")
    public Object[][] multipleIntervalOverlapping() {
        return new Object[][]{
                {cramQueryWithCRAI, cramQueryReference,
                        new QueryInterval[]{new QueryInterval(0, 100009, 100009), new QueryInterval(0, 100011, 100011)},
                        new String[]{"a", "b", "c", "d", "e"}},
                {cramQueryWithLocalCRAI, cramQueryReference,
                        new QueryInterval[]{new QueryInterval(0, 100009, 100009), new QueryInterval(0, 100011, 100011)},
                        new String[]{"a", "b", "c", "d", "e"}},
                {cramQueryWithBAI, cramQueryReference,
                        new QueryInterval[]{new QueryInterval(0, 100009, 100009), new QueryInterval(0, 100011, 100011)},
                        new String[]{"a", "b", "c", "d", "e"}},
                // no matching reads
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 300, 310), new QueryInterval(1, 300, 310)},
                        new String[]{}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 300, 310), new QueryInterval(1, 300, 310)},
                        new String[]{}},
                // matching reads from first interval only
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 205, 206), new QueryInterval(3, 300, 301)},
                        new String[]{"a", "b"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 205, 206), new QueryInterval(3, 300, 301)},
                        new String[]{"a", "b"}},
                // matching reads from last interval only
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(3, 700, 701)},
                        new String[]{"k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(3, 700, 701)},
                        new String[]{"k"}},
                //matching reads from each interval
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 205, 206), new QueryInterval(3, 700, 701)},
                        new String[]{"a", "b", "k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 205, 206), new QueryInterval(3, 700, 701)},
                        new String[]{"a", "b", "k"}},
                //matching reads from each interval - 4 intervals
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{
                                new QueryInterval(0, 200, 201), new QueryInterval(1, 500, 501),
                                new QueryInterval(2, 300, 301), new QueryInterval(3, 700, 701)},
                        new String[]{"a", "f", "i", "k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{
                                new QueryInterval(0, 200, 201), new QueryInterval(1, 500, 501),
                                new QueryInterval(2, 300, 301), new QueryInterval(3, 700, 701)},
                        new String[]{"a", "f", "i", "k"}},
                // first read is before the first interval
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(2, 300, 301), new QueryInterval(3, 700, 701)},
                        new String[]{"i", "k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(2, 300, 301), new QueryInterval(3, 700, 701)},
                        new String[]{"i", "k"}},
                // first interval is before the first read
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(0, 200, 201)},
                        new String[]{"a"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(0, 200, 201)},
                        new String[]{"a"}},
                // intervals in reverse order
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(0, 200, 201)},
                        new String[]{"a"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(0, 200, 201)},
                        new String[]{"a"}},
        };
    }

    @Test(dataProvider="multipleIntervalOverlapping")
    public void testQueryOverlappingHtsIntervals(
            final HtsPath cramFileName,
            final HtsPath referenceFileName,
            final QueryInterval[] intervals,
            final String[] expectedNames)
    {
        final QueryInterval[] optimizedIntervals = QueryInterval.optimizeIntervals(intervals);
        Assert.assertTrue(optimizedIntervals.length > 1);

        doQueryTest(
                decoder -> decoder.queryOverlapping(
                        HtsIntervalUtils.fromQueryIntervalArray(
                                optimizedIntervals,
                                decoder.getHeader().getSequenceDictionary()
                        )
                ),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @DataProvider(name = "otherMultipleIntervals")
    public Object[][] otherMultipleIntervals() {
        return new Object[][]{
                // accept an empty QueryIntervalArray
                {cramQueryWithBAI, cramQueryReference,
                        new QueryInterval[]{},
                        new String[]{}},
                // intervals overlapping - optimized to a single interval
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 1000, 1030), new QueryInterval(0, 1020, 1076)},
                        new String[]{"d"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 1000, 1030), new QueryInterval(0, 1020, 1076)},
                        new String[]{"d"}}
        };
    }

    // these are tested separately because we want the other multi-interval tests to
    // assert that the interval list size is > 1 post-optimization to ensure we're
    // using more than one interval; whereas these tests optimize down to 0 or 1 intervals
    @Test(dataProvider="otherMultipleIntervals")
    public void testOtherMultipleIntervals(
            final HtsPath cramFileName,
            final HtsPath referenceFileName,
            final QueryInterval[] intervals,
            final String[] expectedNames)
    {
        final QueryInterval[] optimizedIntervals = QueryInterval.optimizeIntervals(intervals);
        doQueryTest(
                decoder -> decoder.queryContained(HtsIntervalUtils.fromQueryIntervalArray(
                        optimizedIntervals,
                        decoder.getHeader().getSequenceDictionary()
                )),
                cramFileName,
                referenceFileName,
                expectedNames
        );
        doQueryTest(
                decoder -> decoder.queryOverlapping(
                        HtsIntervalUtils.fromQueryIntervalArray(
                                optimizedIntervals,
                            decoder.getHeader().getSequenceDictionary()
                )),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @DataProvider(name = "multipleIntervalContained")
    public Object[][] multipleIntervalContained() {
        return new Object[][]{
                {cramQueryWithCRAI, cramQueryReference,
                        new QueryInterval[]{new QueryInterval(0, 100008, 100008), new QueryInterval(0, 100013, 0)},
                        new String[]{"f", "f"}},
                {cramQueryWithLocalCRAI, cramQueryReference,
                        new QueryInterval[]{new QueryInterval(0, 100008, 100008), new QueryInterval(0, 100013, 0)},
                        new String[]{"f", "f"}},
                {cramQueryWithBAI, cramQueryReference,
                        new QueryInterval[]{new QueryInterval(0, 100008, 100008), new QueryInterval(0, 100013, 0)},
                        new String[]{"f", "f"}},
                // no matching reads
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 300, 310), new QueryInterval(1, 300, 310)},
                        new String[]{}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 300, 310), new QueryInterval(1, 300, 310)},
                        new String[]{}},
                // matching reads from first interval only
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 205, 305), new QueryInterval(3, 300, 301)},
                        new String[]{"b", "c"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 205, 305), new QueryInterval(3, 300, 301)},
                        new String[]{"b", "c"}},
                // matching reads from last interval only
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(3, 700, 776)},
                        new String[]{"k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(3, 700, 776)},
                        new String[]{"k"}},
                //matching reads from each interval
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 200, 281), new QueryInterval(3, 700, 776)},
                        new String[]{"a", "b", "k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 200, 281), new QueryInterval(3, 700, 776)},
                        new String[]{"a", "b", "k"}},
                //matching reads from each interval - 4 intervals
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{
                                new QueryInterval(0, 200, 281), new QueryInterval(1, 500, 576),
                                new QueryInterval(2, 300, 376), new QueryInterval(3, 700, 776)},
                        new String[]{"a", "b", "f", "i", "k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{
                                new QueryInterval(0, 200, 281), new QueryInterval(1, 500, 576),
                                new QueryInterval(2, 300, 376), new QueryInterval(3, 700, 776)},
                        new String[]{"a", "b", "f", "i", "k"}},
                // first read is before the first interval
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(2, 300, 301), new QueryInterval(3, 700, 776)},
                        new String[]{"k"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(2, 300, 301), new QueryInterval(3, 700, 776)},
                        new String[]{"k"}},
                // first interval is before the first read
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(0, 200, 276)},
                        new String[]{"a"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 100, 101), new QueryInterval(0, 200, 276)},
                        new String[]{"a"}},
                // intervals in reverse order
                {cramQueryReadsWithBAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 200, 276), new QueryInterval(0, 100, 101)},
                        new String[]{"a"}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference,
                        new QueryInterval[]{new QueryInterval(0, 200, 276), new QueryInterval(0, 100, 101)},
                        new String[]{"a"}},
        };
    }

    @Test(dataProvider="multipleIntervalContained")
    public void testQueryContainedMultipleIntervals(
            final HtsPath cramFileName,
            final HtsPath referenceFileName,
            final QueryInterval[] intervals,
            final String[] expectedNames)
    {
        final QueryInterval[] optimizedIntervals = QueryInterval.optimizeIntervals(intervals);
        Assert.assertTrue(optimizedIntervals.length > 1);
        doQueryTest(
                decoder -> decoder.queryContained(HtsIntervalUtils.fromQueryIntervalArray(
                        optimizedIntervals,
                        decoder.getHeader().getSequenceDictionary()
                )),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @DataProvider(name = "unmappedQueries")
    public Object[][] unmappedQueries() {
        return new Object[][] {
                {cramQueryWithCRAI, cramQueryReference, new String[]{"g", "h", "h", "i", "i"}},
                {cramQueryWithLocalCRAI, cramQueryReference, new String[]{"g", "h", "h", "i", "i"}},
                {cramQueryWithBAI, cramQueryReference, new String[]{"g", "h", "h", "i", "i"}},
                //no unmapped reads
                {cramQueryReadsWithBAI, cramQueryReadsReference, new String[]{}},
                {cramQueryReadsWithLocalCRAI, cramQueryReadsReference, new String[]{}}
        };
    }

    @Test(dataProvider="unmappedQueries")
    public void testQueryUnmapped(
            final HtsPath cramFileName,
            final HtsPath referenceFileName,
            final String[] expectedNames)
    {
        doQueryTest(
                decoder -> decoder.queryUnmapped(),
                cramFileName,
                referenceFileName,
                expectedNames
        );
    }

    @DataProvider(name = "alignmentStartQueries")
    public Object[][] alignmentStartQueries() {
        return new Object[][] {
                // cram file, reference, query contig, alignment start query, number of reads with matching alignment start
                {cramQueryWithCRAI, cramQueryReference, "20", 100013, 2},
                {cramQueryWithLocalCRAI, cramQueryReference, "20", 100013, 2},
                {cramQueryWithBAI, cramQueryReference, "20", 100013, 2},
                // tests to ensure that query results on inputs that have matching reads that are followed by reads
                // in the same container that don't match the alignment start are properly constrained to only the
                // matching reads
                {new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.cram"),
                        new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.fa"),
                        "Mito", 631, 24},
                {new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTestGATKGen.cram"),
                        new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.fa"),
                        "Mito", 631, 24},
        };
    }

    @Test(dataProvider="alignmentStartQueries")
    public void testQueryAlignmentStart(
            final IOPath cramInputPath,
            final IOPath referencePath,
            final String queryContig,
            final long alignmentStart,
            final int expectedReadCount) {
        final ReadsDecoderOptions readsDecoderOptions =
                new ReadsDecoderOptions()
                        .setValidationStringency(ValidationStringency.LENIENT)
                        .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));

        try (final ReadsDecoder cramDecoder =
                     HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions);
             final CloseableIterator<SAMRecord> it = cramDecoder.queryStart(queryContig, alignmentStart)) {
            int count = 0;
            while (it.hasNext()) {
                it.next();
                count++;
            }
            Assert.assertEquals(count, expectedReadCount);
        }
    }

    @DataProvider(name = "mateQueries")
    public Object[][] mateQueries() {
        return new Object[][] {
                // cram, reference, query contig, query start, name of read mates at alignment start
                {cramQueryWithCRAI, cramQueryReference, "20", 100013, "f"},
                {cramQueryWithLocalCRAI, cramQueryReference, "20", 100013, "f"},
                {cramQueryWithBAI, cramQueryReference, "20", 100013, "f"},
                {new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.cram"),
                        new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.fa"),
                        "Mito", 631, "IL29_4505:7:30:11521:4492#2"},
                {new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTestGATKGen.cram"),
                        new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.fa"),
                        "Mito", 631, "IL29_4505:7:30:11521:4492#2"},
        };
    }

    @Test(dataProvider="mateQueries")
    public void testQueryMate(
            final IOPath cramInputPath,
            final IOPath referencePath,
            final String queryContig,
            final int alignmentStart,
            final String expectedReadName)
    {
        // the test file ("mitoAlignmentStartTest.cram") used here, which was provided by a user who reported
        // a queryMate bug, contains reads that have a mateReferenceName == "*", but with mateAlignmentStart != 0;
        // so we have to use ValidationStringency.LENIENT to suppress validation from throwing when it sees these..
        final ReadsDecoderOptions readsDecoderOptions =
                new ReadsDecoderOptions()
                        .setValidationStringency(ValidationStringency.LENIENT)
                        .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));

        try (final ReadsDecoder cramDecoder =
                     HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions)) {
            SAMRecord firstRecord;
            SAMRecord firstRecordMate;
             try (final CloseableIterator<SAMRecord> it = cramDecoder.queryStart(queryContig, alignmentStart)) {
                 Assert.assertTrue(it.hasNext());
                 firstRecord = it.next();
                 Assert.assertEquals(firstRecord.getReadName(), expectedReadName);
             }

            // get the mate for the first record
            final Optional<SAMRecord> optFirstMate = cramDecoder.queryMate(firstRecord);
            Assert.assertTrue(optFirstMate.isPresent());
            firstRecordMate = optFirstMate.get();
            Assert.assertEquals(firstRecordMate.getReadName(), firstRecord.getReadName());

            final Optional<SAMRecord> optMatesMate = cramDecoder.queryMate(firstRecordMate);
            Assert.assertTrue(optMatesMate.isPresent());
            final SAMRecord matesMate = optMatesMate.get();
            Assert.assertEquals(matesMate, firstRecord);
        }
    }

    @Test(dataProvider="mateQueries")
    public void testQueryMateEnsureInternalIteratorsClosed(
            final IOPath cramInputPath,
            final IOPath referencePath,
            final String queryContig,
            final int alignmentStart,
            final String expectedReadName)
    {
        // the test file ("mitoAlignmentStartTest.cram") used here, which was provided by a user who reported
        // a queryMate bug, contains reads that have a mateReferenceName == "*", but with mateAlignmentStart != 0;
        // so we have to use ValidationStringency.LENIENT to suppress validation from throwing when it sees these..
        final ReadsDecoderOptions readsDecoderOptions =
                new ReadsDecoderOptions()
                        .setValidationStringency(ValidationStringency.LENIENT)
                        .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));

        try (final ReadsDecoder cramDecoder =
                     HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions)) {
            SAMRecord firstRecord;
            SAMRecord firstRecordMate;
            try (final CloseableIterator<SAMRecord> it = cramDecoder.queryStart(queryContig, alignmentStart)) {
                Assert.assertTrue(it.hasNext());
                firstRecord = it.next();
                Assert.assertEquals(firstRecord.getReadName(), expectedReadName);
            }

            // get the mate for the first record
            final Optional<SAMRecord> optFirstMate = cramDecoder.queryMate(firstRecord);
            Assert.assertTrue(optFirstMate.isPresent());
            firstRecordMate = optFirstMate.get();
            Assert.assertEquals(firstRecordMate.getReadName(), firstRecord.getReadName());

            final Optional<SAMRecord> optMatesMate = cramDecoder.queryMate(firstRecordMate);
            Assert.assertTrue(optMatesMate.isPresent());
            final SAMRecord matesMate = optMatesMate.get();
            Assert.assertEquals(matesMate, firstRecord);

            // now ensure that the querymate calls closed their iterators by re-executing a query that returns
            // an iterator, and verify that it doesn't throw
            try (final CloseableIterator<SAMRecord> it = cramDecoder.queryStart(queryContig, alignmentStart)) {
                Assert.assertTrue(it.hasNext());
                firstRecord = it.next();
                Assert.assertEquals(firstRecord.getReadName(), expectedReadName);
            }
        }
    }

    private void doQueryTest(
            final Function<ReadsDecoder, CloseableIterator<SAMRecord>> getIterator,
            final IOPath cramInputPath,
            final IOPath referencePath,
            final String[] expectedNames)
    {
        final ReadsDecoderOptions readsDecoderOptions =
                new ReadsDecoderOptions()
                        .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));
        try (final ReadsDecoder cramDecoder =
                     HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions);
            final CloseableIterator<SAMRecord> it = getIterator.apply(cramDecoder)) {
            int count = 0;
            while (it.hasNext()) {
                final SAMRecord samRec = it.next();
                Assert.assertTrue(count < expectedNames.length);
                Assert.assertEquals(samRec.getReadName(), expectedNames[count]);
                count++;
            }
            Assert.assertEquals(count, expectedNames.length);
        }
    }


    @DataProvider(name = "iteratorStateTests")
    public Object[][] iteratorStateQueries() {
        return new Object[][] {
                {cramQueryWithCRAI, cramQueryReference, 12, 7},
                {cramQueryWithLocalCRAI, cramQueryReference, 12, 7},
                {cramQueryWithBAI, cramQueryReference, 12, 7}
        };
    }

    @Test(dataProvider="iteratorStateTests")
    public void testSerialQueries(
            final IOPath cramInputPath,
            final IOPath referencePath,
            final int expectedTotalCount,
            final int expectedOverlappingCount)
    {
        final ReadsDecoderOptions readsDecoderOptions =
                new ReadsDecoderOptions()
                        .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));

        try (final ReadsDecoder cramDecoder =
                     HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions)) {
            try (final CloseableIterator<SAMRecord> it = cramDecoder.iterator()) {
                Assert.assertEquals(consumeIterator(it), expectedTotalCount);
            }

            try (final CloseableIterator<SAMRecord> it =
                    cramDecoder.queryOverlapping("20", 100013, 100070)) {
                Assert.assertEquals(consumeIterator(it), expectedOverlappingCount);
            }

            try (final CloseableIterator<SAMRecord> it = cramDecoder.iterator()) {
                Assert.assertEquals(consumeIterator(it), expectedTotalCount);
            }
        }
    }

    @Test(dataProvider="iteratorStateTests", expectedExceptions = IllegalStateException.class)
    public void testRejectParallelQueries(
            final IOPath cramInputPath,
            final IOPath referencePath,
            final int unusedExpectedTotalCount,
            final int unusedExpectedOverlappingCount)
    {
        final ReadsDecoderOptions readsDecoderOptions =
                new ReadsDecoderOptions()
                        .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));

        try (final ReadsDecoder cramDecoder =
                     HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions);
             final CloseableIterator<SAMRecord> origIt = cramDecoder.iterator()) {
            // opening a second iterator while the first one is still open should throw
            cramDecoder.queryOverlapping("20", 100013, 100070);
        }
    }

    @DataProvider(name = "unmappedSliceTest")
    public Object[][] unmappedMultiSliceTest() {
        return new Object[][] {
                // the main test feature of these files is that they have several mapped reads followed by
                // some number of unmapped reads, each created with seqs_per_slice = 100 to force the unmapped
                // reads to be distributed over multiple slices (at least for large numbers of unmapped reads)
                // tests the fix to https://github.com/samtools/htsjdk/issues/562
                {new HtsPath(TEST_DATA_DIR + "NA12878.20.21.1-100.100-SeqsPerSlice.0-unMapped.cram"),
                        new HtsPath(TEST_DATA_DIR + "human_g1k_v37.20.21.1-100.fasta"), 0},
                {new HtsPath(TEST_DATA_DIR + "NA12878.20.21.1-100.100-SeqsPerSlice.1-unMapped.cram"),
                        new HtsPath(TEST_DATA_DIR + "human_g1k_v37.20.21.1-100.fasta"), 1},
                {new HtsPath(TEST_DATA_DIR + "NA12878.20.21.1-100.100-SeqsPerSlice.500-unMapped.cram"),
                        new HtsPath(TEST_DATA_DIR + "human_g1k_v37.20.21.1-100.fasta"), 500},
        };
    }

    @Test(dataProvider = "unmappedSliceTest")
    public void testUnmappedMultiSlice(
            final IOPath cramInputPath,
            final IOPath referencePath,
            final int expectedCount)
    {
        final ReadsDecoderOptions readsDecoderOptions =
                new ReadsDecoderOptions()
                        .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));
        int count = 0;
        try (final ReadsDecoder cramDecoder =
                     HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions);
             final CloseableIterator<SAMRecord> it = cramDecoder.queryUnmapped()) {
            Assert.assertEquals(cramDecoder.getVersion(), new HtsVersion(3, 0, 0));
            while (it.hasNext()) {
                it.next();
                count++;
            }
        }
        Assert.assertEquals(count, expectedCount);
    }

    @DataProvider(name = "serialQueries")
    public Object[][] serialQueries() {
        return new Object[][]{
                {cramQueryWithCRAI, cramQueryReference,
                        new QueryInterval(0, 100009, 100009), new String[]{"a", "b", "c"}, 5},
                {cramQueryWithLocalCRAI, cramQueryReference,
                        new QueryInterval(0, 100009, 100009), new String[]{"a", "b", "c"}, 5},
                {cramQueryWithBAI, cramQueryReference,
                        new QueryInterval(0, 100009, 100009), new String[]{"a", "b", "c"}, 5},
        };
    }

    @Test(dataProvider = "serialQueries")
    public void testSerialQueriesOnRemoteFile(
            final IOPath cramInputPath,
            final IOPath referencePath,
            final QueryInterval interval,
            final String[] expectedNames,
            final int nUnmapped) throws IOException {
        final Path cramIndex = SamFiles.findIndex(cramInputPath.toPath());

        try (final FileSystem jimfs = Jimfs.newFileSystem(Configuration.unix())) {
            final Path jimfsCRAM = jimfs.getPath("remotecram.cram");
            final Path jimfsCRAI = jimfs.getPath("remotecram.crai");

            Files.copy(cramInputPath.toPath(), jimfsCRAM);
            Files.copy(cramIndex, jimfsCRAI);

            final ReadsDecoderOptions readsDecoderOptions =
                    new ReadsDecoderOptions()
                            .setCRAMDecoderOptions(new CRAMDecoderOptions().setReferencePath(referencePath));
            try (final ReadsDecoder cramDecoder =
                         HtsDefaultRegistry.getReadsResolver().getReadsDecoder(cramInputPath, readsDecoderOptions)) {
                 try (final CloseableIterator<SAMRecord> firstIterator = cramDecoder.queryOverlapping(
                         HtsIntervalUtils.fromQueryIntervalArray(
                                 new QueryInterval[]{interval},
                                 cramDecoder.getHeader().getSequenceDictionary()))) {

                     int count = 0;
                     while (firstIterator.hasNext()) {
                         final SAMRecord samRec = firstIterator.next();
                         Assert.assertTrue(count < expectedNames.length);
                         Assert.assertEquals(samRec.getReadName(), expectedNames[count]);
                         count++;
                     }
                     Assert.assertEquals(count, expectedNames.length);
                 }

                // now execute a second query on the same cramDecoder, this time for unmapped reads
                try (final CloseableIterator<SAMRecord> secondIterator = cramDecoder.queryUnmapped()) {
                    int count = 0;
                    while (secondIterator.hasNext()) {
                        final SAMRecord samRec = secondIterator.next();
                        Assert.assertTrue(samRec.getReadUnmappedFlag());
                        count++;
                    }
                    Assert.assertEquals(count, nUnmapped);
                }
            }
        }
    }

    private static int consumeIterator(Iterator<?> it) {
        int count = 0;
        while (it.hasNext()) {
            it.next();
            count++;
        }
        return count;
    }
}
