test_index_protocol.py

"""
Tests for the 'Index' class and protocol. All Index classes should support
this functionality.
"""

import pytest
import glob

import sourmash
from sourmash import SourmashSignature
from sourmash.index import (
    LinearIndex,
    ZipFileLinearIndex,
    LazyLinearIndex,
    MultiIndex,
    StandaloneManifestIndex,
    IndexSearchResult,
)
from sourmash.index import CounterGather
from sourmash.index.sqlite_index import SqliteIndex
from sourmash.index.revindex import RevIndex
from sourmash.sbt import SBT, GraphFactory
from sourmash.manifest import CollectionManifest, BaseCollectionManifest
from sourmash.lca.lca_db import LCA_Database, load_single_database
from sourmash.minhash import flatten_and_intersect_scaled, flatten_and_downsample_scaled

import sourmash_tst_utils as utils


def _load_three_sigs():
    # utility function - load & return these three sigs.
    sig2 = utils.get_test_data("2.fa.sig")
    sig47 = utils.get_test_data("47.fa.sig")
    sig63 = utils.get_test_data("63.fa.sig")

    ss2 = sourmash.load_one_signature(sig2, ksize=31)
    ss47 = sourmash.load_one_signature(sig47)
    ss63 = sourmash.load_one_signature(sig63)

    return [ss2, ss47, ss63]


def build_linear_index(runtmp):
    ss2, ss47, ss63 = _load_three_sigs()

    lidx = LinearIndex()
    lidx.insert(ss2)
    lidx.insert(ss47)
    lidx.insert(ss63)

    return lidx


def build_lazy_linear_index(runtmp):
    lidx = build_linear_index(runtmp)
    return LazyLinearIndex(lidx)


def build_sbt_index(runtmp):
    ss2, ss47, ss63 = _load_three_sigs()

    factory = GraphFactory(5, 100, 3)
    root = SBT(factory, d=2)

    root.insert(ss2)
    root.insert(ss47)
    root.insert(ss63)

    return root


def build_sbt_index_save_load(runtmp):
    root = build_sbt_index(runtmp)
    out = runtmp.output("xyz.sbt.zip")
    root.save(out)

    return sourmash.load_file_as_index(out)


def build_zipfile_index(runtmp):
    from sourmash.save_load import SaveSignatures_ZipFile

    location = runtmp.output("index.zip")
    with SaveSignatures_ZipFile(location) as save_sigs:
        for ss in _load_three_sigs():
            save_sigs.add(ss)

    idx = ZipFileLinearIndex.load(location)
    return idx


def build_multi_index(runtmp):
    siglist = _load_three_sigs()
    lidx = LinearIndex(siglist)

    mi = MultiIndex.load([lidx], [None], None)
    return mi


def build_standalone_manifest_index(runtmp):
    sig2 = utils.get_test_data("2.fa.sig")
    sig47 = utils.get_test_data("47.fa.sig")
    sig63 = utils.get_test_data("63.fa.sig")

    ss2 = sourmash.load_one_signature(sig2, ksize=31)
    ss47 = sourmash.load_one_signature(sig47)
    ss63 = sourmash.load_one_signature(sig63)

    siglist = [(ss2, sig2), (ss47, sig47), (ss63, sig63)]

    rows = []
    rows.extend((CollectionManifest.make_manifest_row(ss, loc) for ss, loc in siglist))
    mf = CollectionManifest(rows)
    mf_filename = runtmp.output("mf.csv")

    mf.write_to_filename(mf_filename)

    idx = StandaloneManifestIndex.load(mf_filename)
    return idx


def build_lca_index(runtmp):
    siglist = _load_three_sigs()
    db = LCA_Database(31, 1000, "DNA")
    for ss in siglist:
        db.insert(ss)

    return db


def build_lca_index_save_load(runtmp):
    db = build_lca_index(runtmp)
    outfile = runtmp.output("db.lca.json")
    db.save(outfile)

    return sourmash.load_file_as_index(outfile)


def build_sqlite_index(runtmp):
    filename = runtmp.output("idx.sqldb")
    db = SqliteIndex.create(filename)

    siglist = _load_three_sigs()
    for ss in siglist:
        db.insert(ss)

    return db


def build_revindex(runtmp):
    ss2, ss47, ss63 = _load_three_sigs()

    lidx = RevIndex(template=ss2.minhash)
    lidx.insert(ss2)
    lidx.insert(ss47)
    lidx.insert(ss63)

    return lidx


def build_lca_index_save_load_sql(runtmp):
    db = build_lca_index(runtmp)
    outfile = runtmp.output("db.lca.json")
    db.save(outfile, format="sql")

    x = load_single_database(outfile)
    db_load = x[0]

    return db_load


#
# create a fixture 'index_obj' that is parameterized by all of these
# building functions.
#


@pytest.fixture(
    params=[
        build_linear_index,
        build_lazy_linear_index,
        build_sbt_index,
        build_zipfile_index,
        build_multi_index,
        build_standalone_manifest_index,
        build_lca_index,
        build_sbt_index_save_load,
        build_lca_index_save_load,
        build_sqlite_index,
        build_lca_index_save_load_sql,
        #                        build_revindex,
    ]
)
def index_obj(request, runtmp):
    build_fn = request.param

    # build on demand
    return build_fn(runtmp)


###
### generic Index tests go here
###


def test_index_search_exact_match(index_obj):
    # search for an exact match
    ss2, ss47, ss63 = _load_three_sigs()

    sr = index_obj.search(ss2, threshold=1.0)
    print([s[1].name for s in sr])
    assert len(sr) == 1
    assert sr[0].signature.minhash == ss2.minhash
    assert sr[0].score == 1.0


def test_index_search_lower_threshold(index_obj):
    # search at a lower threshold/multiple results with ss47
    ss2, ss47, ss63 = _load_three_sigs()

    sr = index_obj.search(ss47, threshold=0.1)
    print([s[1].name for s in sr])
    assert len(sr) == 2
    sr.sort(key=lambda x: -x[0])
    assert sr[0].signature.minhash == ss47.minhash
    assert sr[0].score == 1.0
    assert sr[1].signature.minhash == ss63.minhash
    assert round(sr[1].score, 2) == 0.32


def test_index_search_lower_threshold_2(index_obj):
    # search at a lower threshold/multiple results with ss63
    ss2, ss47, ss63 = _load_three_sigs()

    sr = index_obj.search(ss63, threshold=0.1)
    print([s[1].name for s in sr])
    assert len(sr) == 2
    sr.sort(key=lambda x: -x[0])
    assert sr[0].signature.minhash == ss63.minhash
    assert sr[0].score == 1.0
    assert sr[1].signature.minhash == ss47.minhash
    assert round(sr[1].score, 2) == 0.32


def test_index_search_higher_threshold_2(index_obj):
    # search at a higher threshold/one match
    ss2, ss47, ss63 = _load_three_sigs()

    # search for sig63 with high threshold => 1 match
    sr = index_obj.search(ss63, threshold=0.8)
    print([s[1].name for s in sr])
    assert len(sr) == 1
    sr.sort(key=lambda x: -x[0])
    assert sr[0].signature.minhash == ss63.minhash
    assert sr[0].score == 1.0


def test_index_search_containment(index_obj):
    # search for containment at a low threshold/multiple results with ss63
    ss2, ss47, ss63 = _load_three_sigs()

    sr = index_obj.search(ss63, do_containment=True, threshold=0.1)
    print([s[1].name for s in sr])
    assert len(sr) == 2
    sr.sort(key=lambda x: -x[0])
    assert sr[0].signature.minhash == ss63.minhash
    assert sr[0].score == 1.0
    assert sr[1].signature.minhash == ss47.minhash
    assert round(sr[1].score, 2) == 0.48


def test_index_signatures(index_obj):
    # signatures works?
    siglist = list(index_obj.signatures())

    ss2, ss47, ss63 = _load_three_sigs()
    assert len(siglist) == 3

    # check md5sums, since 'in' doesn't always work
    md5s = set(ss.md5sum() for ss in siglist)
    assert ss2.md5sum() in md5s
    assert ss47.md5sum() in md5s
    assert ss63.md5sum() in md5s


def test_index_signatures_with_location(index_obj):
    # signatures_with_location works?
    siglist = list(index_obj.signatures_with_location())

    ss2, ss47, ss63 = _load_three_sigs()
    assert len(siglist) == 3

    # check md5sums, since 'in' doesn't always work
    md5s = set((ss.md5sum() for ss, loc in siglist))
    assert ss2.md5sum() in md5s
    assert ss47.md5sum() in md5s
    assert ss63.md5sum() in md5s


def test_index_len(index_obj):
    # len works?
    assert len(index_obj) == 3


def test_index_bool(index_obj):
    # bool works?
    assert bool(index_obj)


def test_index_location(index_obj):
    # location works?
    assert str(index_obj.location)


def test_index_manifest(index_obj):
    # manifest is either None or a BaseCollectionManifest
    manifest = index_obj.manifest
    if manifest is not None:
        assert isinstance(manifest, BaseCollectionManifest)


def test_index_select_basic(index_obj):
    # select does the basic thing ok
    idx = index_obj.select(
        ksize=31,
        moltype="DNA",
        abund=False,
        containment=True,
        scaled=1000,
        num=0,
        picklist=None,
    )

    assert len(idx) == 3
    siglist = list(idx.signatures())
    assert len(siglist) == 3

    # check md5sums, since 'in' doesn't always work
    md5s = set(ss.md5sum() for ss in siglist)
    ss2, ss47, ss63 = _load_three_sigs()
    assert ss2.md5sum() in md5s
    assert ss47.md5sum() in md5s
    assert ss63.md5sum() in md5s


def test_index_select_nada(index_obj):
    # select works ok when nothing matches!

    # CTB: currently this EITHER raises a ValueError OR returns an empty
    # Index object, depending on implementation. :think:
    # See: https://github.com/sourmash-bio/sourmash/issues/1940
    try:
        idx = index_obj.select(ksize=21)
    except ValueError:
        idx = LinearIndex([])

    assert len(idx) == 0
    siglist = list(idx.signatures())
    assert len(siglist) == 0


def test_index_prefetch(index_obj):
    # test basic prefetch
    ss2, ss47, ss63 = _load_three_sigs()

    # search for ss2
    results = []
    for result in index_obj.prefetch(ss2, threshold_bp=0):
        results.append(result)

    assert len(results) == 1
    assert results[0].signature.minhash == ss2.minhash

    # search for ss47 - expect two results
    results = []
    for result in index_obj.prefetch(ss47, threshold_bp=0):
        results.append(result)

    assert len(results) == 2
    assert results[0].signature.minhash == ss47.minhash
    assert results[1].signature.minhash == ss63.minhash


def test_index_best_containment(index_obj):
    # test basic containment search
    ss2, ss47, ss63 = _load_three_sigs()

    match = index_obj.best_containment(ss2)
    assert match
    assert match.score == 1.0
    assert match.signature.minhash == ss2.minhash

    match = index_obj.best_containment(ss47)
    assert match
    assert match.score == 1.0
    assert match.signature.minhash == ss47.minhash


def test_index_best_containment_threshold_1(index_obj):
    # test best_containment() method, in some detail
    ss2, ss47, ss63 = _load_three_sigs()

    # now construct query signatures with specific numbers of hashes --
    # note, these signatures all have scaled=1000.

    mins = list(sorted(ss2.minhash.hashes))
    new_mh = ss2.minhash.copy_and_clear()

    # query with empty hashes
    assert not new_mh
    with pytest.raises(ValueError):
        index_obj.best_containment(SourmashSignature(new_mh))

    # add one hash
    new_mh.add_hash(mins.pop())
    assert len(new_mh) == 1

    result = index_obj.best_containment(SourmashSignature(new_mh))
    assert result
    containment, match_sig, name = result
    assert containment == 1.0
    assert match_sig.minhash == ss2.minhash

    # check with a threshold -> should be no results.
    with pytest.raises(ValueError):
        index_obj.best_containment(SourmashSignature(new_mh), threshold_bp=5000)

    # add three more hashes => length of 4
    new_mh.add_hash(mins.pop())
    new_mh.add_hash(mins.pop())
    new_mh.add_hash(mins.pop())
    assert len(new_mh) == 4

    result = index_obj.best_containment(SourmashSignature(new_mh))
    assert result
    containment, match_sig, name = result
    assert containment == 1.0
    assert match_sig.minhash == ss2.minhash

    # check with a too-high threshold -> should be no results.
    with pytest.raises(ValueError):
        index_obj.best_containment(SourmashSignature(new_mh), threshold_bp=5000)


def test_best_containment_threshold_5(index_obj):
    # test gather() method, in some detail
    ss2, ss47, ss63 = _load_three_sigs()

    # now construct query signatures with specific numbers of hashes --
    # note, these signatures all have scaled=1000.

    mins = list(sorted(ss2.minhash.hashes.keys()))
    new_mh = ss2.minhash.copy_and_clear()

    # add five hashes
    for i in range(5):
        new_mh.add_hash(mins.pop())
        new_mh.add_hash(mins.pop())
        new_mh.add_hash(mins.pop())
        new_mh.add_hash(mins.pop())
        new_mh.add_hash(mins.pop())

    # should get a result with no threshold (any match at all is returned)
    result = index_obj.best_containment(SourmashSignature(new_mh))
    assert result
    containment, match_sig, name = result
    assert containment == 1.0
    assert match_sig.minhash == ss2.minhash

    # now, check with a threshold_bp that should be meet-able.
    result = index_obj.best_containment(SourmashSignature(new_mh), threshold_bp=5000)
    assert result
    containment, match_sig, name = result
    assert containment == 1.0
    assert match_sig.minhash == ss2.minhash


###
### CounterGather tests
###


def create_basic_counter_gather(runtmp):
    "Construct a CounterGather class."
    return CounterGather


class CounterGather_LinearIndex:
    """
    Provides an (inefficient) CounterGather-style class, for
    protocol testing purposes.
    """

    def __init__(self, orig_query):
        "Constructor - take a SourmashSignature that is the original query."
        orig_query_mh = orig_query.minhash
        if orig_query_mh.scaled == 0:
            raise ValueError

        # Index object used to actually track matches.
        self.idx = LinearIndex()
        self.orig_query_mh = orig_query_mh.copy().flatten()
        self.query_started = 0
        self.scaled = orig_query_mh.scaled
        self.locations = {}

    def add(self, ss, *, location=None, require_overlap=True):
        "Insert potential match."
        if self.query_started:
            raise ValueError("cannot add more signatures to counter after peek/consume")

        # skip duplicates
        md5 = ss.md5sum()
        if md5 in self.locations:
            return

        # confirm that this match has an overlap...
        add_mh = ss.minhash.flatten()
        overlap = self.orig_query_mh.count_common(add_mh, downsample=True)

        # ...figure out what scaled we are operating at now...
        if overlap:
            self.downsample(add_mh.scaled)
        elif require_overlap:
            raise ValueError("no overlap between query and signature!?")

        # ...and add to the Index, while also tracking location!
        self.idx.insert(ss)
        self.locations[md5] = location

    def signatures(self):
        "Yield all signatures"
        return self.idx.signatures()

    def downsample(self, scaled):
        "Track highest scaled across all possible matches."
        if scaled > self.scaled:
            self.scaled = scaled
        return self.scaled

    def peek(self, cur_query_mh, *, threshold_bp=0):
        """
        Find best match to current query within this CounterGather object.
        """
        self.query_started = 1

        scaled = self.downsample(cur_query_mh.scaled)
        cur_query_mh = flatten_and_downsample_scaled(cur_query_mh, scaled)

        # no hashes remaining? exit.
        if not self.orig_query_mh or not cur_query_mh:
            return []

        # verify current query is a subset of the original.
        if cur_query_mh.contained_by(self.orig_query_mh, downsample=True) < 1:
            raise ValueError("current query not a subset of original query")

        # did we get a match?
        res = self.idx.peek(cur_query_mh, threshold_bp=threshold_bp)
        if not res:
            return []
        sr, intersect_mh = res

        # got match - replace location & return.
        match = sr.signature
        md5 = match.md5sum()
        location = self.locations[md5]
        return IndexSearchResult(sr.score, match, location), intersect_mh

    def consume(self, *args, **kwargs):
        self.query_started = 1
        return self.idx.consume(*args, **kwargs)


class CounterGather_LCA:
    """
    Provides an alternative implementation of a CounterGather-style class,
    based on LCA_Database. This is currently just for protocol
    and API testing purposes.
    """

    def __init__(self, query):
        from sourmash.lca.lca_db import LCA_Database

        query_mh = query.minhash
        if query_mh.scaled == 0:
            raise ValueError("must use scaled MinHash")

        self.orig_query_mh = query_mh
        lca_db = LCA_Database(query_mh.ksize, query_mh.scaled, query_mh.moltype)
        self.db = lca_db
        self.siglist = {}
        self.locations = {}
        self.query_started = 0

    def add(self, ss, *, location=None, require_overlap=True):
        "Add this signature into the counter."
        if self.query_started:
            raise ValueError("cannot add more signatures to counter after peek/consume")

        overlap = self.orig_query_mh.count_common(ss.minhash, True)
        if overlap:
            self.downsample(ss.minhash.scaled)
        elif require_overlap:
            raise ValueError("no overlap between query and signature!?")

        self.db.insert(ss)

        md5 = ss.md5sum()
        self.siglist[md5] = ss
        self.locations[md5] = location

    def signatures(self):
        "Yield all signatures."
        yield from self.siglist.values()

    def downsample(self, scaled):
        "Track highest scaled across all possible matches."
        if scaled > self.db.scaled:
            self.db.downsample_scaled(scaled)
        return self.db.scaled

    def peek(self, query_mh, *, threshold_bp=0):
        "Return next possible match."
        from sourmash import SourmashSignature

        self.query_started = 1
        scaled = self.downsample(query_mh.scaled)
        query_mh = query_mh.downsample(scaled=scaled)

        if not self.orig_query_mh or not query_mh:
            return []

        if query_mh.contained_by(self.orig_query_mh, downsample=True) < 1:
            raise ValueError("current query not a subset of original query")

        query_ss = SourmashSignature(query_mh)

        # returns search_result, intersect_mh
        try:
            result = self.db.best_containment(query_ss, threshold_bp=threshold_bp)
        except ValueError:
            result = None

        if not result:
            return []

        cont = result.score
        match = result.signature

        intersect_mh = flatten_and_intersect_scaled(result.signature.minhash, query_mh)

        md5 = result.signature.md5sum()
        location = self.locations[md5]

        new_sr = IndexSearchResult(cont, match, location)
        return [new_sr, intersect_mh]

    def consume(self, intersect_mh):
        self.query_started = 1


@pytest.fixture(
    params=[
        CounterGather,
        CounterGather_LinearIndex,
        CounterGather_LCA,
    ]
)
def counter_gather_constructor(request):
    build_fn = request.param

    # build on demand
    return build_fn


def test_counter_get_signatures(counter_gather_constructor):
    # test .signatures() method
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear()
    match_mh_2.add_many(range(10, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear()
    match_mh_3.add_many(range(15, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    siglist = list(counter.signatures())
    assert len(siglist) == 3
    assert match_ss_1 in siglist
    assert match_ss_2 in siglist
    assert match_ss_3 in siglist


def _consume_all(query_mh, counter, threshold_bp=0):
    results = []
    query_mh = query_mh.to_mutable()

    last_intersect_size = None
    while 1:
        result = counter.peek(query_mh, threshold_bp=threshold_bp)
        if not result:
            break

        sr, intersect_mh = result
        print(sr.signature.name, len(intersect_mh))
        if last_intersect_size:
            assert len(intersect_mh) <= last_intersect_size

        last_intersect_size = len(intersect_mh)

        counter.consume(intersect_mh)
        query_mh.remove_many(intersect_mh.hashes)

        results.append((sr, len(intersect_mh)))

    return results


def test_counter_gather_1(counter_gather_constructor):
    # check a contrived set of non-overlapping gather results,
    # generated via CounterGather
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear()
    match_mh_2.add_many(range(10, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear()
    match_mh_3.add_many(range(15, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    results = _consume_all(query_ss.minhash, counter)

    expected = (
        ["match1", 10],
        ["match2", 5],
        ["match3", 2],
    )
    assert len(results) == len(expected), results

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_1_b(counter_gather_constructor):
    # check a contrived set of somewhat-overlapping gather results,
    # generated via CounterGather. Here the overlaps are structured
    # so that the gather results are the same as those in
    # test_counter_gather_1(), even though the overlaps themselves are
    # larger.
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear()
    match_mh_2.add_many(range(7, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear()
    match_mh_3.add_many(range(13, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    results = _consume_all(query_ss.minhash, counter)

    expected = (
        ["match1", 10],
        ["match2", 5],
        ["match3", 2],
    )
    assert len(results) == len(expected), results

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_1_c_with_threshold(counter_gather_constructor):
    # check a contrived set of somewhat-overlapping gather results,
    # generated via CounterGather. Here the overlaps are structured
    # so that the gather results are the same as those in
    # test_counter_gather_1(), even though the overlaps themselves are
    # larger.
    # use a threshold, here.

    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear()
    match_mh_2.add_many(range(7, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear()
    match_mh_3.add_many(range(13, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    results = _consume_all(query_ss.minhash, counter, threshold_bp=3)

    expected = (["match1", 10], ["match2", 5])
    assert len(results) == len(expected), results

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_1_d_diff_scaled(counter_gather_constructor):
    # test as above, but with different scaled.
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear().downsample(scaled=10)
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear().downsample(scaled=20)
    match_mh_2.add_many(range(7, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear().downsample(scaled=30)
    match_mh_3.add_many(range(13, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    results = _consume_all(query_ss.minhash, counter)

    expected = (
        ["match1", 10],
        ["match2", 5],
        ["match3", 2],
    )
    assert len(results) == len(expected), results

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_1_d_diff_scaled_query(counter_gather_constructor):
    # test as above, but with different scaled for QUERY.
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))

    match_mh_1 = query_mh.copy_and_clear().downsample(scaled=10)
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear().downsample(scaled=20)
    match_mh_2.add_many(range(7, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear().downsample(scaled=30)
    match_mh_3.add_many(range(13, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    # downsample query now -
    query_ss = SourmashSignature(query_mh.downsample(scaled=100), name="query")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    results = _consume_all(query_ss.minhash, counter)

    expected = (
        ["match1", 10],
        ["match2", 5],
        ["match3", 2],
    )
    assert len(results) == len(expected), results

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_1_e_abund_query(counter_gather_constructor):
    # test as above, but abund query
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1, track_abundance=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear().flatten()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear().flatten()
    match_mh_2.add_many(range(7, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear().flatten()
    match_mh_3.add_many(range(13, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    # must flatten before peek!
    results = _consume_all(query_ss.minhash.flatten(), counter)

    expected = (
        ["match1", 10],
        ["match2", 5],
        ["match3", 2],
    )
    assert len(results) == len(expected), results

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_1_f_abund_match(counter_gather_constructor):
    # test as above, but abund query
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1, track_abundance=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh.flatten(), name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    match_mh_2 = query_mh.copy_and_clear()
    match_mh_2.add_many(range(7, 15))
    match_ss_2 = SourmashSignature(match_mh_2, name="match2")

    match_mh_3 = query_mh.copy_and_clear()
    match_mh_3.add_many(range(13, 17))
    match_ss_3 = SourmashSignature(match_mh_3, name="match3")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)
    counter.add(match_ss_2)
    counter.add(match_ss_3)

    # must flatten before peek!
    results = _consume_all(query_ss.minhash.flatten(), counter)

    expected = (
        ["match1", 10],
        ["match2", 5],
        ["match3", 2],
    )
    assert len(results) == len(expected), results

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_2(counter_gather_constructor):
    # check basic set of gather results on semi-real data,
    # generated via CounterGather
    testdata_combined = utils.get_test_data("gather/combined.sig")
    testdata_glob = utils.get_test_data("gather/GCF*.sig")
    testdata_sigs = glob.glob(testdata_glob)

    query_ss = sourmash.load_one_signature(testdata_combined, ksize=21)
    subject_sigs = [
        (sourmash.load_one_signature(t, ksize=21), t) for t in testdata_sigs
    ]

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    for ss, loc in subject_sigs:
        counter.add(ss, location=loc)

    results = _consume_all(query_ss.minhash, counter)

    expected = (
        ["NC_003198.1", 487],
        ["NC_000853.1", 192],
        ["NC_011978.1", 169],
        ["NC_002163.1", 157],
        ["NC_003197.2", 152],
        ["NC_009486.1", 92],
        ["NC_006905.1", 76],
        ["NC_011080.1", 59],
        ["NC_011274.1", 42],
        ["NC_006511.1", 31],
        ["NC_011294.1", 7],
        ["NC_004631.1", 2],
    )
    assert len(results) == len(expected)

    for (sr, size), (exp_name, exp_size) in zip(results, expected):
        sr_name = sr.signature.name.split()[0]
        print(sr_name, size)

        assert sr_name == exp_name
        assert size == exp_size


def test_counter_gather_exact_match(counter_gather_constructor):
    # query == match
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    # load up the counter; provide a location override, too.
    counter = counter_gather_constructor(query_ss)
    counter.add(query_ss, location="somewhere over the rainbow")

    results = _consume_all(query_ss.minhash, counter)
    assert len(results) == 1
    (sr, intersect_mh) = results[0]

    assert sr.score == 1.0
    assert sr.signature == query_ss
    assert sr.location == "somewhere over the rainbow"


def test_counter_gather_multiple_identical_matches(counter_gather_constructor):
    # test multiple identical matches being inserted, with only one return
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    # create counter...
    counter = counter_gather_constructor(query_ss)

    # now add multiple identical matches.
    match_mh = query_mh.copy_and_clear()
    match_mh.add_many(range(5, 15))

    for name in "match1", "match2", "match3":
        match_ss = SourmashSignature(match_mh, name=name)
        counter.add(match_ss, location=name)

    results = _consume_all(query_ss.minhash, counter)
    assert len(results) == 1

    sr, overlap_count = results[0]
    assert sr.score == 0.5
    assert overlap_count == 10

    # any one of the three is valid
    assert sr.location in ("match1", "match2", "match3")


def test_counter_gather_add_after_peek(counter_gather_constructor):
    # cannot add after peek or consume
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(query_ss, location="somewhere over the rainbow")

    counter.peek(query_ss.minhash)

    with pytest.raises(ValueError):
        counter.add(query_ss, location="try again")


def test_counter_gather_add_after_consume(counter_gather_constructor):
    # cannot add after peek or consume
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(query_ss, location="somewhere over the rainbow")

    counter.consume(query_ss.minhash)

    with pytest.raises(ValueError):
        counter.add(query_ss, location="try again")


def test_counter_gather_consume_empty_intersect(counter_gather_constructor):
    # check that consume works fine when there is an empty signature.
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(query_ss, location="somewhere over the rainbow")

    # nothing really happens here :laugh:, just making sure there's no error
    counter.consume(query_ss.minhash.copy_and_clear())


def test_counter_gather_empty_initial_query(counter_gather_constructor):
    # check empty initial query
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1, require_overlap=False)

    assert counter.peek(query_ss.minhash) == []


def test_counter_gather_num_query(counter_gather_constructor):
    # check num query
    query_mh = sourmash.MinHash(n=500, ksize=31)
    query_mh.add_many(range(0, 10))
    query_ss = SourmashSignature(query_mh, name="query")

    with pytest.raises(ValueError):
        counter_gather_constructor(query_ss)


def test_counter_gather_empty_cur_query(counter_gather_constructor):
    # test empty cur query
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(query_ss, location="somewhere over the rainbow")

    cur_query_mh = query_ss.minhash.copy_and_clear()
    results = _consume_all(cur_query_mh, counter)
    assert results == []


def test_counter_gather_add_num_matchy(counter_gather_constructor):
    # test add num query
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh = sourmash.MinHash(n=500, ksize=31)
    match_mh.add_many(range(0, 20))
    match_ss = SourmashSignature(match_mh, name="query")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    with pytest.raises(ValueError):
        counter.add(match_ss, location="somewhere over the rainbow")


def test_counter_gather_bad_cur_query(counter_gather_constructor):
    # test cur query that is not subset of original query
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(query_ss, location="somewhere over the rainbow")

    cur_query_mh = query_ss.minhash.copy_and_clear()
    cur_query_mh.add_many(range(20, 30))
    with pytest.raises(ValueError):
        counter.peek(cur_query_mh)


def test_counter_gather_add_no_overlap(counter_gather_constructor):
    # check adding match with no overlap w/query
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 10))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(10, 20))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    with pytest.raises(ValueError):
        counter.add(match_ss_1)

    assert counter.peek(query_ss.minhash) == []


def test_counter_gather_big_threshold(counter_gather_constructor):
    # check 'peek' with a huge threshold
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_mh.add_many(range(0, 20))
    query_ss = SourmashSignature(query_mh, name="query")

    match_mh_1 = query_mh.copy_and_clear()
    match_mh_1.add_many(range(0, 10))
    match_ss_1 = SourmashSignature(match_mh_1, name="match1")

    # load up the counter
    counter = counter_gather_constructor(query_ss)
    counter.add(match_ss_1)

    # impossible threshold:
    threshold_bp = 30 * query_ss.minhash.scaled
    results = counter.peek(query_ss.minhash, threshold_bp=threshold_bp)
    assert results == []


def test_counter_gather_empty_counter(counter_gather_constructor):
    # check empty counter
    query_mh = sourmash.MinHash(n=0, ksize=31, scaled=1)
    query_ss = SourmashSignature(query_mh, name="query")

    # empty counter!
    counter = counter_gather_constructor(query_ss)

    assert counter.peek(query_ss.minhash) == []