import os
import unittest

import ir_datasets

from tira.third_party_integrations import (
    ensure_pyterrier_is_loaded,
    load_ir_datasets,
    register_rerank_data_to_ir_datasets,
)

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


class TestIRDatasets(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        register_rerank_data_to_ir_datasets("tests/sample-tirex-rerank-data", "tirex-sample-dataset", "vaswani")

    def test_for_existing_document_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        dataset = ir_datasets.load("tirex-sample-dataset")

        expected_text = (
            "microwave spectroscopy  includes chapters on spectroscope technique\n"
            "and design on measurements on gases liquids and solids on nuclear\n"
            "properties on molecular structure and on further possible applications\n"
            "of microwaves\n"
        )

        assert expected_text == dataset.docs_store().get("8172").text
        assert "8172" == dataset.docs_store().get("8172").doc_id

    def test_for_existing_document_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        dataset = ir_datasets.load("tirex-sample-dataset")

        expected_text = (
            "microwave spectroscopy  includes chapters on spectroscope technique\n"
            "and design on measurements on gases liquids and solids on nuclear\n"
            "properties on molecular structure and on further possible applications\n"
            "of microwaves\n"
        )

        assert expected_text == dataset.docs_store().get("8172").text
        assert "8172" == dataset.docs_store().get("8172").doc_id

    def test_for_non_existing_document_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        dataset = ir_datasets.load("tirex-sample-dataset")

        assert dataset.docs_store().get("does-not-exist") is None

    def test_for_non_existing_document_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        dataset = ir_datasets.load("tirex-sample-dataset")

        assert dataset.docs_store().get("does-not-exist") is None

    def test_for_queries_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        dataset = ir_datasets.load("tirex-sample-dataset")

        assert len(list(dataset.queries_iter())) == 1
        assert (
            "MEASUREMENT OF DIELECTRIC CONSTANT OF LIQUIDS BY THE USE OF MICROWAVE TECHNIQUES\n"
            == list(dataset.queries_iter())[0].text
        )

    def test_for_queries_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        dataset = ir_datasets.load("tirex-sample-dataset")

        assert len(list(dataset.queries_iter())) == 1
        assert (
            "MEASUREMENT OF DIELECTRIC CONSTANT OF LIQUIDS BY THE USE OF MICROWAVE TECHNIQUES\n"
            == list(dataset.queries_iter())[0].text
        )

    def test_for_scoreddocs_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        dataset = ir_datasets.load("tirex-sample-dataset")
        print(list(dataset.scoreddocs))
        assert len(list(dataset.scoreddocs)) == 10
        assert "8172" == list(dataset.scoreddocs)[0].doc_id
        assert "1" == list(dataset.scoreddocs)[0].query_id

    def test_for_scoreddocs_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        dataset = ir_datasets.load("tirex-sample-dataset")
        print(list(dataset.scoreddocs))
        assert len(list(dataset.scoreddocs)) == 10
        assert "8172" == list(dataset.scoreddocs)[0].doc_id
        assert "1" == list(dataset.scoreddocs)[0].query_id

    def test_for_scoreddocs_within_tira(self):
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/re-ranking-outputs"
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        dataset = ir_datasets.load("msmarco-passage/eval/small")
        print(list(dataset.scoreddocs))
        del os.environ["TIRA_INPUT_DATASET"]
        assert len(list(dataset.scoreddocs)) == 3
        assert "doc-1" == list(dataset.scoreddocs)[0].doc_id
        assert "1" == list(dataset.scoreddocs)[0].query_id
        assert 10 == list(dataset.scoreddocs)[0].score
        register_rerank_data_to_ir_datasets("tests/sample-tirex-rerank-data", "tirex-sample-dataset", None)

    def test_loading_raw_ir_datasets_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        ir_datasets = load_ir_datasets()
        dataset = ir_datasets.load("cranfield")
        queries = {str(i.query_id): i.text for i in dataset.queries_iter()}

        assert len(list(dataset.queries_iter())) == 225
        assert queries["269"] == "has a criterion been established for determining the axial compressor\nchoking line ."

    def test_loading_raw_ir_datasets_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        ir_datasets = load_ir_datasets()
        dataset = ir_datasets.load("cranfield")
        queries = {str(i.query_id): i.text for i in dataset.queries_iter()}

        assert len(list(dataset.queries_iter())) == 225
        assert queries["269"] == "has a criterion been established for determining the axial compressor\nchoking line ."

    def test_loading_queries_from_ir_datasets_from_custom_directory_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")
        queries = {i.query_id: i.text for i in dataset.queries_iter()}

        assert len(list(dataset.queries_iter())) == 2
        assert queries["1"] == "fox jumps above animal"
        self.assertEqual("static_ir_dataset-tests/resources/sample-input-full-rank", dataset.dataset_id())

    def test_loading_queries_from_ir_datasets_from_custom_directory_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")
        queries = {i.query_id: i.text for i in dataset.queries_iter()}

        assert len(list(dataset.queries_iter())) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_loading_queries_from_ir_datasets_from_custom_directory_2_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank-without-documents"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")
        queries = {i.query_id: i.text for i in dataset.queries_iter()}

        assert len(list(dataset.queries_iter())) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_loading_queries_from_ir_datasets_from_custom_directory_2_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank-without-documents"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")
        queries = {i.query_id: i.text for i in dataset.queries_iter()}

        assert len(list(dataset.queries_iter())) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_loading_documents_from_ir_datasets_from_custom_directory_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")
        docs = {i.doc_id: i.text for i in dataset.docs_iter()}

        assert len(list(dataset.docs_iter())) == 5
        assert docs["pangram-02"] == "Quick fox jumps nightly above wizard."

    def test_loading_documents_from_ir_datasets_from_custom_directory_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")
        docs = {i.doc_id: i.text for i in dataset.docs_iter()}

        assert len(list(dataset.docs_iter())) == 5
        assert docs["pangram-02"] == "Quick fox jumps nightly above wizard."

    def test_loading_docs_store_from_ir_datasets_from_custom_directory_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")

        assert dataset.docs_store()["pangram-03"].text == "The jay, pig, fox, zebra and my wolves quack!"

    def test_loading_docs_store_from_ir_datasets_from_custom_directory_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        ir_datasets = load_ir_datasets()
        del os.environ["TIRA_INPUT_DATASET"]
        dataset = ir_datasets.load("does-not-exist-and-is-not-used")

        assert dataset.docs_store()["pangram-03"].text == "The jay, pig, fox, zebra and my wolves quack!"

    def test_no_patching_for_pyterrier_datasets_01(self):
        ensure_pyterrier_is_loaded()
        import pyterrier as pt

        dataset = pt.get_dataset("irds:cranfield")
        queries = {str(i["qid"]): i["query"] for _, i in dataset.get_topics().iterrows()}

        assert len(dataset.get_topics()) == 225
        assert queries["269"] == "has a criterion been established for determining the axial compressor choking line"

    def test_no_patching_for_pyterrier_datasets_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        import pyterrier as pt

        dataset = pt.get_dataset("irds:cranfield")
        queries = {str(i["qid"]): i["query"] for _, i in dataset.get_topics().iterrows()}

        assert len(dataset.get_topics()) == 225
        assert queries["269"] == "has a criterion been established for determining the axial compressor choking line"

    def test_patching_for_pyterrier_datasets_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        import ir_datasets

        ir_load = ir_datasets.load
        ensure_pyterrier_is_loaded()
        import pyterrier as pt

        dataset = pt.get_dataset("irds:does-not-exist-and-is-not-used")
        ir_datasets.load = ir_load
        del os.environ["TIRA_INPUT_DATASET"]

        queries = {i["qid"]: i["query"] for _, i in dataset.get_topics().iterrows()}

        assert len(list(dataset.get_topics("title"))) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_patching_for_pyterrier_datasets_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        import ir_datasets

        ir_load = ir_datasets.load
        ensure_pyterrier_is_loaded()
        import pyterrier as pt

        dataset = pt.get_dataset("irds:does-not-exist-and-is-not-used")
        ir_datasets.load = ir_load
        del os.environ["TIRA_INPUT_DATASET"]

        queries = {i["qid"]: i["query"] for _, i in dataset.get_topics().iterrows()}

        assert len(list(dataset.get_topics("title"))) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_patching_for_pyterrier_datasets_for_text_queries_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        import ir_datasets

        ir_load = ir_datasets.load
        ensure_pyterrier_is_loaded()
        import pyterrier as pt

        dataset = pt.get_dataset("irds:does-not-exist-and-is-not-used")
        ir_datasets.load = ir_load
        del os.environ["TIRA_INPUT_DATASET"]

        queries = {i["qid"]: i["query"] for _, i in dataset.get_topics("text").iterrows()}

        assert len(list(dataset.get_topics("text"))) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_patching_for_pyterrier_datasets_for_text_queries_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        import ir_datasets

        ir_load = ir_datasets.load
        ensure_pyterrier_is_loaded()
        import pyterrier as pt

        dataset = pt.get_dataset("irds:does-not-exist-and-is-not-used")
        ir_datasets.load = ir_load
        del os.environ["TIRA_INPUT_DATASET"]

        queries = {i["qid"]: i["query"] for _, i in dataset.get_topics("text").iterrows()}

        assert len(list(dataset.get_topics("text"))) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_patching_for_pyterrier_datasets_for_title_queries_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        import ir_datasets

        ir_load = ir_datasets.load
        ensure_pyterrier_is_loaded()
        import pyterrier as pt

        dataset = pt.get_dataset("irds:does-not-exist-and-is-not-used")
        ir_datasets.load = ir_load
        del os.environ["TIRA_INPUT_DATASET"]

        queries = {i["qid"]: i["query"] for _, i in dataset.get_topics("title").iterrows()}

        assert len(list(dataset.get_topics("title"))) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_patching_for_pyterrier_datasets_for_title_queries_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/sample-input-full-rank"
        import ir_datasets

        ir_load = ir_datasets.load
        ensure_pyterrier_is_loaded()
        import pyterrier as pt

        dataset = pt.get_dataset("irds:does-not-exist-and-is-not-used")
        ir_datasets.load = ir_load
        del os.environ["TIRA_INPUT_DATASET"]

        queries = {i["qid"]: i["query"] for _, i in dataset.get_topics("title").iterrows()}

        assert len(list(dataset.get_topics("title"))) == 2
        assert queries["1"] == "fox jumps above animal"

    def test_loading_ir_dataset_via_rest_api_from_tira_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert dataset.has_docs(), "dataset has no documents"
        assert dataset.has_queries(), "dataset has no queries"
        assert dataset.has_qrels(), "dataset has no qrels"

    def test_loading_ir_dataset_via_rest_api_from_tira_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert dataset.has_docs(), "dataset has no documents"
        assert dataset.has_queries(), "dataset has no queries"
        assert dataset.has_qrels(), "dataset has no qrels"

    def test_loading_qrels_via_rest_api_from_tira_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert len([i for i in dataset.qrels_iter()]) == 2635
        assert (
            "TrecQrel(query_id='1', doc_id='2005.ipm_journal-ir0anthology0volumeA41A1.7', relevance=1, iteration='0')"
            == str([i for i in dataset.qrels_iter()][0])
        )

    def test_loading_qrels_via_rest_api_from_tira_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert len([i for i in dataset.qrels_iter()]) == 2635
        assert (
            "TrecQrel(query_id='1', doc_id='2005.ipm_journal-ir0anthology0volumeA41A1.7', relevance=1, iteration='0')"
            == str([i for i in dataset.qrels_iter()][0])
        )

    def test_loading_docs_via_rest_api_from_tira_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert len([i for i in dataset.docs_iter()]) == 53673
        docsstore = dataset.docs_store()

        assert docsstore.get("2005.ipm_journal-ir0anthology0volumeA41A1.7").text.startswith(
            "A probabilistic model for stemmer generation AbstractIn this paper"
        )

    def test_loading_docs_via_rest_api_from_tira_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert len([i for i in dataset.docs_iter()]) == 53673
        docsstore = dataset.docs_store()

        assert docsstore.get("2005.ipm_journal-ir0anthology0volumeA41A1.7").text.startswith(
            "A probabilistic model for stemmer generation AbstractIn this paper"
        )

    def test_loading_topics_via_rest_api_from_tira_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.topics_file("ir-lab-jena-leipzig-wise-2023/training-20231104-training")

        self.assertIn("/training-20231104-training/truth-data/queries.xml", str(dataset))

    def test_loading_topics_via_rest_api_from_tira_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        from tira.third_party_integrations import ir_datasets

        topics_file = ir_datasets.topics_file("ir-lab-jena-leipzig-wise-2023/training-20231104-training")

        self.assertIn("/training-20231104-training/truth-data/queries.xml", str(topics_file))

    def test_loading_topics_via_rest_api_from_rerank_dataset_from_tira_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        os.environ["TIRA_INPUT_DATASET"] = "tests/resources/re-ranking-outputs"
        ir_datasets = load_ir_datasets()
        dataset = ir_datasets.load("workshop-on-open-web-search/re-ranking-20231027-training")
        del os.environ["TIRA_INPUT_DATASET"]

        assert len([i for i in dataset.queries_iter()]) == 2
        assert len([i for i in dataset.docs_iter()]) == 3

    def test_loading_queries_via_rest_api_from_tira_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=False)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert 68 == len(list(dataset.queries_iter()))
        print(str(list(dataset.queries_iter())[0]))
        assert (
            "TirexQuery(query_id='1', text='retrieval system improving effectiveness', title='retrieval system"
            " improving effectiveness', query='retrieval system improving effectiveness', description='What papers"
            " focus on improving the effectiveness of a retrieval system?', narrative='Relevant papers include research"
            " on what makes a retrieval system effective and what improves the effectiveness of a retrieval system."
            " Papers that focus on improving something else or improving the effectiveness of a system that is not a"
            " retrieval system are not relevant.')" == str(list(dataset.queries_iter())[0])
        )

    def test_loading_queries_via_rest_api_from_tira_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        from tira.third_party_integrations import ir_datasets

        dataset = ir_datasets.load("ir-lab-jena-leipzig-sose-2023/iranthology-20230618-training")

        assert 68 == len(list(dataset.queries_iter()))
        print(str(list(dataset.queries_iter())[0]))
        assert (
            "TirexQuery(query_id='1', text='retrieval system improving effectiveness', title='retrieval system"
            " improving effectiveness', query='retrieval system improving effectiveness', description='What papers"
            " focus on improving the effectiveness of a retrieval system?', narrative='Relevant papers include research"
            " on what makes a retrieval system effective and what improves the effectiveness of a retrieval system."
            " Papers that focus on improving something else or improving the effectiveness of a system that is not a"
            " retrieval system are not relevant.')" == str(list(dataset.queries_iter())[0])
        )

    def test_patching_for_pyterrier_datasets_to_tira(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        import pyterrier as pt

        dataset = pt.get_dataset("irds:workshop-on-open-web-search/retrieval-20231027-training")
        queries = {i["qid"]: i["title"] for _, i in dataset.get_topics().iterrows()}

        assert len(dataset.get_topics()) == 3
        assert queries["1"] == "hubble telescope achievements"

    def test_patching_for_pyterrier_datasets_with_qrels_to_tira(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        import pyterrier as pt

        dataset = pt.get_dataset("irds:workshop-on-open-web-search/retrieval-20231027-training")
        qrels = [i.to_dict() for _, i in dataset.get_qrels().iterrows()]

        assert len(dataset.get_qrels()) == 3
        assert qrels[0] == {"docno": "doc-1", "iteration": "0", "label": 1, "qid": "1"}

    def test_ir_datasets_id_is_available(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        import pyterrier as pt

        expected = "workshop-on-open-web-search/retrieval-20231027-training"

        dataset = pt.get_dataset("irds:workshop-on-open-web-search/retrieval-20231027-training")
        actual = dataset.irds_ref().dataset_id()

        self.assertEqual(expected, actual)

    def test_ir_datasets_get_many_iter_is_available_01(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        import pyterrier as pt

        expected = []

        dataset = pt.get_dataset("irds:workshop-on-open-web-search/retrieval-20231027-training")
        actual = [i for i in dataset.irds_ref().docs_store().get_many_iter(["1", "2"]) if i]

        self.assertEqual(expected, actual)

    def test_ir_datasets_get_many_iter_is_available_02(self):
        ensure_pyterrier_is_loaded(patch_ir_datasets=True)
        import pyterrier as pt

        expected = ["doc-2", "doc-3"]

        dataset = pt.get_dataset("irds:workshop-on-open-web-search/retrieval-20231027-training")
        actual = [i.doc_id for i in dataset.irds_ref().docs_store().get_many_iter(["doc-2", "doc-3"]) if i]

        self.assertEqual(expected, actual)
