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 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;
import org.testng.Assert;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

// 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", 100013L, 2},
            {cramQueryWithLocalCRAI, cramQueryReference, "20", 100013L, 2},
            {cramQueryWithBAI, cramQueryReference, "20", 100013L, 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",
                631L,
                24
            },
            {
                new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTestGATKGen.cram"),
                new HtsPath(TEST_DATA_DIR + "mitoAlignmentStartTest.fa"),
                "Mito",
                631L,
                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;
    }
}
