__copyright__ = "Copyright (c) 2020-2021 Jina AI Limited. All rights reserved."
__license__ = "Apache-2.0"

from pathlib import Path
from typing import List

import pytest
from jina import Document, DocumentArray, Executor

from ...spacy_text_encoder import SpacyTextEncoder

_EMBEDDING_DIM = 96


@pytest.fixture(scope='session')
def basic_encoder() -> SpacyTextEncoder:
    return SpacyTextEncoder()


def test_config():
    ex = Executor.load_config(str(Path(__file__).parents[2] / 'config.yml'))
    assert ex.__class__.__name__ == 'SpacyTextEncoder'


def test_encoding_cpu():
    enc = SpacyTextEncoder(require_gpu=False)
    input_data = DocumentArray([Document(text='hello world')])

    enc.encode(docs=input_data, parameters={})

    assert input_data[0].embedding.shape == (_EMBEDDING_DIM,)


@pytest.mark.gpu
def test_encoding_gpu():
    enc = SpacyTextEncoder(require_gpu=True)
    input_data = DocumentArray([Document(text='hello world')])

    enc.encode(docs=input_data, parameters={})

    assert input_data[0].embedding.shape == (_EMBEDDING_DIM,)


@pytest.mark.parametrize(
    'model_name, emb_dim',
    [
        ('en_core_web_sm', 96),
        ('en_core_web_lg', 300),
        ('es_core_news_sm', 96),
    ],
)
def test_models(model_name: str, emb_dim: int):
    encoder = SpacyTextEncoder(model_name)
    input_data = DocumentArray([Document(text='hello world')])

    encoder.encode(docs=input_data, parameters={})

    assert input_data[0].embedding.shape == (emb_dim,)


@pytest.mark.parametrize(
    'traversal_paths, counts',
    [
        (['r'], [['r', 1], ['c', 0], ['cc', 0]]),
        (['c'], [['r', 0], ['c', 3], ['cc', 0]]),
        (['cc'], [['r', 0], ['c', 0], ['cc', 2]]),
        (['cc', 'r'], [['r', 1], ['c', 0], ['cc', 2]]),
    ],
)
def test_traversal_path(
    traversal_paths: List[str], counts: List, basic_encoder: SpacyTextEncoder
):
    text = 'blah'
    docs = DocumentArray([Document(id='root1', text=text)])
    docs[0].chunks = [
        Document(id='chunk11', text=text),
        Document(id='chunk12', text=text),
        Document(id='chunk13', text=text),
    ]
    docs[0].chunks[0].chunks = [
        Document(id='chunk111', text=text),
        Document(id='chunk112', text=text),
    ]

    basic_encoder.encode(docs=docs, parameters={'traversal_paths': traversal_paths})
    for path, count in counts:
        embeddings = docs.traverse_flat([path]).get_attributes('embedding')
        assert len(list(filter(lambda x: x is not None, embeddings))) == count


@pytest.mark.parametrize('batch_size', [1, 2, 4, 8])
def test_batch_size(basic_encoder: SpacyTextEncoder, batch_size: int):
    docs = DocumentArray([Document(text='hello there') for _ in range(32)])
    basic_encoder.encode(docs, parameters={'batch_size': batch_size})

    for doc in docs:
        assert doc.embedding.shape == (_EMBEDDING_DIM,)


def test_quality_embeddings(basic_encoder: SpacyTextEncoder):
    docs = DocumentArray(
        [
            Document(id='A', text='a furry animal that with a long tail'),
            Document(id='B', text='a domesticated mammal with four legs'),
            Document(id='C', text='a type of aircraft that uses rotating wings'),
            Document(id='D', text='flying vehicle that has fixed wings and engines'),
        ]
    )

    basic_encoder.encode(DocumentArray(docs), {})

    # assert semantic meaning is captured in the encoding
    docs.match(docs)
    matches = ['B', 'A', 'D', 'C']
    for i, doc in enumerate(docs):
        assert doc.matches[1].id == matches[i]

__copyright__ = "Copyright (c) 2020-2021 Jina AI Limited. All rights reserved."
__license__ = "Apache-2.0"

import os
from pathlib import Path

import pytest
import spacy

from jina import Document, DocumentArray, Executor

from ...spacy_text_encoder import SpacyTextEncoder


def test_config():
    ex = Executor.load_config(str(Path(__file__).parents[2] / 'config.yml'))
    assert ex.lang == 'en_core_web_sm'


def test_spacy_text_encoder():
    # Input
    docs = DocumentArray(
        [
            Document(text='Han likes eating pizza'),
            Document(text='Han likes pizza'),
            Document(text='Jina rocks'),
        ]
    )

    # Encoder embedding
    encoder = SpacyTextEncoder()
    encoder.encode(docs, parameters={})
    # Compare with ouptut
    assert len(docs) == 3
    for doc in docs:
        assert doc.embedding.shape == (96,)


def test_spacy_text_encoder_traversal_paths():
    # Input
    docs = DocumentArray(
        [
            Document(
                chunks=[
                    Document(text='Han likes eating pizza'),
                    Document(text='Han likes pizza'),
                ]
            ),
            Document(chunks=[Document(text='Jina rocks')]),
        ]
    )

    # Encoder embedding
    encoder = SpacyTextEncoder()
    encoder.encode(docs, parameters={'traversal_paths': ['c']})
    # Compare with ouptut
    assert len(docs) == 2
    assert len(docs[0].chunks) == 2
    for chunk in docs[0].chunks:
        assert chunk.embedding.shape == (96,)
    assert len(docs[1].chunks) == 1
    for chunk in docs[1].chunks:
        assert chunk.embedding.shape == (96,)


def test_unsupported_lang(tmpdir):
    dummy1 = spacy.blank('xx')
    dummy1_dir_path = os.path.join(tmpdir, 'xx1')
    dummy1.to_disk(dummy1_dir_path)
    dummy2 = spacy.blank('xx')
    dummy2_dir_path = os.path.join(tmpdir, 'xx2')
    dummy2.to_disk(dummy2_dir_path)
    # No available language
    with pytest.raises(IOError):
        SpacyTextEncoder('abcd')

    # Language does not have DependencyParser should thrown an error
    # when try to use default encoder
    with pytest.raises(ValueError):
        SpacyTextEncoder(dummy1_dir_path, use_default_encoder=True)

    # And should be fine when 'parser' pipeline is added
    dummy1.add_pipe('parser')
    dummy1.to_disk(dummy1_dir_path)
    SpacyTextEncoder(dummy1_dir_path, use_default_encoder=True)

    # Language does not have SentenceRecognizer should thrown an error
    # when try to use non default encoder
    with pytest.raises(ValueError):
        SpacyTextEncoder(dummy2_dir_path, use_default_encoder=False)

    # And should be fine when 'senter' pipeline is added
    dummy2.add_pipe('tok2vec')
    dummy2.to_disk(dummy2_dir_path)
    SpacyTextEncoder(dummy2_dir_path, use_default_encoder=False)
