import collections
import itertools

import numpy as np
import pytest
from spacy.tokens import Doc

from textacy import Corpus
from textacy import load_spacy_lang


@pytest.fixture(scope="module")
def langs(lang_en):
    return ["en_core_web_sm", lang_en]


@pytest.fixture(scope="module")
def corpus(lang_en, docs_en):
    return Corpus(lang_en, data=docs_en)


class TestCorpusInit:
    def test_init_lang(self, langs):
        for lang in langs:
            assert isinstance(Corpus(lang), Corpus)

    @pytest.mark.parametrize("lang", [b"en_core_web_sm", None])
    def test_init_bad_lang(self, lang):
        with pytest.raises(TypeError):
            Corpus(lang)

    def test_init_texts(self, lang_en, texts_en):
        limit = 3
        texts = texts_en[:limit]
        corpus = Corpus(lang_en, data=texts)
        assert len(corpus) == len(corpus.docs) == limit
        assert all(doc.vocab is corpus.spacy_lang.vocab for doc in corpus)
        assert all(text == doc.text for text, doc in zip(texts, corpus))

    def test_init_records(self, lang_en, records_en):
        limit = 3
        records = records_en[:limit]
        corpus = Corpus(lang_en, data=records)
        assert len(corpus) == len(corpus.docs) == limit
        assert all(doc.vocab is corpus.spacy_lang.vocab for doc in corpus)
        assert all(record[0] == doc.text for record, doc in zip(records, corpus))
        assert all(record[1] == doc._.meta for record, doc in zip(records, corpus))

    def test_init_docs(self, lang_en, docs_en):
        limit = 3
        docs = docs_en[:limit]
        corpus = Corpus(lang_en, data=docs)
        assert len(corpus) == len(corpus.docs) == limit
        assert all(doc.vocab is corpus.spacy_lang.vocab for doc in corpus)
        assert all(doc1 is doc2 for doc1, doc2 in zip(docs, corpus))

    def test_init_no_parser(self):
        spacy_lang = load_spacy_lang("en_core_web_sm", disable=("parser",))
        corpus = Corpus(spacy_lang, data=(spacy_lang("This is a sentence in a doc."),))
        assert len(corpus) == 1
        assert corpus.n_sents == 0


class TestCorpusDunder:
    def test_str(self, corpus):
        cstr = str(corpus)
        assert cstr.startswith("Corpus")
        assert all(f"{n}" in cstr for n in [corpus.n_docs, corpus.n_tokens])

    def test_len(self, corpus):
        assert isinstance(len(corpus), int)
        assert len(corpus) == len(corpus.docs) == corpus.n_docs

    def test_iter(self, corpus):
        assert isinstance(corpus, collections.abc.Iterable)

    def test_getitem(self, corpus):
        assert isinstance(corpus[0], Doc)
        assert all(isinstance(doc, Doc) for doc in corpus[0:2])

    def test_delitem(self, corpus):
        del corpus[-1]
        del corpus[-1:]
        with pytest.raises(TypeError):
            del corpus["foo"]


class TestCorpusProperties:
    def test_vectors(self, corpus):
        vectors = corpus.vectors
        assert isinstance(vectors, np.ndarray)
        assert len(vectors.shape) == 2
        assert vectors.shape[0] == len(corpus)

    def test_vector_norms(self, corpus):
        vector_norms = corpus.vector_norms
        assert isinstance(vector_norms, np.ndarray)
        assert len(vector_norms.shape) == 2
        assert vector_norms.shape[0] == len(corpus)


class TestCorpusMethods:

    # NOTE: this test must run before we add more docs that don't have valid metadata
    @pytest.mark.parametrize(
        "name, agg_func, exp_type",
        [
            ("num_sections", sum, int),
            ("title", max, str),
            ("categories", lambda x: list(itertools.chain(x)), list),
        ],
    )
    def test_agg_metadata(self, name, agg_func, exp_type, corpus):
        result = corpus.agg_metadata(name, agg_func)
        assert isinstance(result, exp_type)

    def test_add(self, corpus, lang_en):
        datas = (
            "This is an english sentence.",
            ("This is an english sentence.", {"foo": "bar"}),
            lang_en("This is an english sentence."),
            ["This is one sentence.", "This is another sentence."],
            [
                ("This is sentence #1.", {"foo": "bar"}),
                ("This is sentence #2.", {"bat": "baz"}),
            ],
            [
                lang_en("This is sentence #1"),
                lang_en("This is sentence #2"),
            ],
        )
        n_docs = corpus.n_docs
        for data in datas:
            corpus.add(data)
            assert corpus.n_docs > n_docs
            n_docs = corpus.n_docs

    @pytest.mark.xfail(reason="there seems to be bug in spacy for nprocess > 1 ...")
    def test_add_nprocess(self, corpus):
        datas = (
            ["This is one sentence.", "This is another sentence."],
            [
                ("This is sentence #1.", {"foo": "bar"}),
                ("This is sentence #2.", {"bat": "baz"}),
            ],
        )
        n_docs = corpus.n_docs
        for data in datas:
            corpus.add(data, n_process=2)
            assert corpus.n_docs > n_docs
            n_docs = corpus.n_docs

    def test_add_typeerror(self, corpus):
        datas = (
            b"This is a byte string.",
            [b"This is a byte string.", b"This is another byte string."],
        )
        for data in datas:
            with pytest.raises(TypeError):
                corpus.add(data)

    def test_get(self, corpus):
        match_funcs = (
            lambda doc: True,
            lambda doc: "Pollution" in doc._.meta.get("categories", []),
        )
        for match_func in match_funcs:
            assert len(list(corpus.get(match_func))) > 0
            assert len(list(corpus.get(match_func, limit=1))) == 1

    def test_remove(self, corpus):
        match_funcs = (
            lambda doc: doc._.meta.get("foo") == "bar",
            lambda doc: len(doc) < 10,
        )
        n_docs = corpus.n_docs
        for match_func in match_funcs[:1]:
            corpus.remove(match_func)
            assert corpus.n_docs < n_docs
            assert not any([match_func(doc) for doc in corpus])
            n_docs = corpus.n_docs

    def test_word_counts(self, corpus):
        abs_counts = corpus.word_counts(weighting="count", by="lower")
        rel_counts = corpus.word_counts(weighting="freq", by="lower")
        assert isinstance(abs_counts, dict)
        assert all(isinstance(count, int) for count in abs_counts.values())
        assert min(abs_counts.values()) > 0
        assert isinstance(rel_counts, dict)
        assert all(isinstance(count, float) for count in rel_counts.values())
        assert min(rel_counts.values()) > 0 and max(rel_counts.values()) <= 1

    def test_word_counts_error(self, corpus):
        with pytest.raises(ValueError):
            corpus.word_counts(weighting="foo")

    def test_word_doc_counts(self, corpus):
        abs_counts = corpus.word_doc_counts(weighting="count", by="lower")
        rel_counts = corpus.word_doc_counts(weighting="freq", by="lower")
        inv_counts = corpus.word_doc_counts(weighting="idf", by="lower")
        assert isinstance(abs_counts, dict)
        assert all(isinstance(count, int) for count in abs_counts.values())
        assert min(abs_counts.values()) > 0
        assert isinstance(rel_counts, dict)
        assert all(isinstance(count, float) for count in rel_counts.values())
        assert min(rel_counts.values()) > 0 and max(rel_counts.values()) <= 1
        assert isinstance(inv_counts, dict)
        assert min(inv_counts.values()) > 0

    def test_word_doc_counts_error(self, corpus):
        with pytest.raises(ValueError):
            corpus.word_doc_counts(weighting="foo")

    def test_corpus_save_and_load(self, corpus, tmpdir, lang_en):
        filepath = str(tmpdir.join("test_corpus_save_and_load.bin"))
        corpus.save(filepath)
        loaded_corpus = Corpus.load(lang_en, filepath)
        assert isinstance(loaded_corpus, Corpus)
        assert len(loaded_corpus) == len(corpus)
        assert loaded_corpus.spacy_lang.meta == corpus.spacy_lang.meta
        assert loaded_corpus.spacy_lang.pipe_names == corpus.spacy_lang.pipe_names
        assert corpus[0].user_data == loaded_corpus[0].user_data
