import unittest

import pandas as pd

from tira.local_client import Client

# .. todo:: this file still uses "unclean" assertions and should be converted to use unittest assertions


def retrieval_submission(queries):
    queries = pd.DataFrame([{"qid": str(i)} for i in queries])
    tira = Client("tests/resources/")
    retriever = tira.pt.from_retriever_submission(
        "ir-benchmarks/tira-ir-starters/retriever-for-re-ranking", dataset="d1"
    )

    ret = retriever(queries)
    return [{"qid": i["qid"], "docno": i["docno"], "score": i["score"], "text": i["text"]} for _, i in ret.iterrows()]


class TestPtFromRetrieverForReRankingSubmissionLoader(unittest.TestCase):
    def test_for_multiple_queries(self):
        expected = [
            {"qid": "1", "docno": "doc-1", "score": 10, "text": "Text of doc-1"},
            {"qid": "1", "docno": "doc-2", "score": 9, "text": "Text of doc-2"},
            {"qid": "3", "docno": "doc-3", "score": 1, "text": "Text of doc-3"},
        ]

        actual = retrieval_submission([1, 3])

        assert len(actual) == 3
        assert actual[0] == expected[0]
        assert actual[1] == expected[1]
        assert actual[2] == expected[2]

    def test_for_single_query_no_1(self):
        expected = [
            {"qid": "1", "docno": "doc-1", "score": 10, "text": "Text of doc-1"},
            {"qid": "1", "docno": "doc-2", "score": 9, "text": "Text of doc-2"},
        ]

        actual = retrieval_submission([1])

        assert len(actual) == 2
        assert actual[0] == expected[0]
        assert actual[1] == expected[1]

    def test_for_single_query_no_3(self):
        expected = [
            {"qid": "3", "docno": "doc-3", "score": 1, "text": "Text of doc-3"},
        ]

        actual = retrieval_submission([3])

        assert len(actual) == 1
        assert actual[0] == expected[0]
