from pathlib import Path
from typing import List

import pytest
from jina import Document, DocumentArray, Executor

from ...flair_text import FlairTextEncoder

_EMBEDDING_DIM = 100


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


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


def test_no_document(basic_encoder: FlairTextEncoder):
    basic_encoder.encode(None, {})


def test_empty_documents(basic_encoder: FlairTextEncoder):
    docs = DocumentArray([])
    basic_encoder.encode(docs, {})
    assert len(docs) == 0


def test_no_text_documents(basic_encoder: FlairTextEncoder):
    docs = DocumentArray([Document()])
    basic_encoder.encode(docs, {})
    assert len(docs) == 1
    assert docs[0].embedding is None


def test_embeddings_str_error():
    with pytest.raises(ValueError, match='embeddings'):
        FlairTextEncoder(embeddings='word:glove')


def test_pooling_strategy_error():
    with pytest.raises(ValueError, match='pooling_strategy'):
        FlairTextEncoder(pooling_strategy='wrong')


def test_unknown_model_error():
    with pytest.raises(ValueError, match='The model name wrong'):
        FlairTextEncoder(embeddings=['wrong:glove'])


def test_encoding_cpu():
    docs = DocumentArray([Document(text='hello there')])
    encoder = FlairTextEncoder(device='cpu')
    encoder.encode(docs, {})

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


@pytest.mark.gpu
def test_encoding_gpu():
    docs = DocumentArray([Document(text='hello there')])
    encoder = FlairTextEncoder(device='cuda')
    encoder.encode(docs, {})

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


@pytest.mark.parametrize(
    'embeddings, dim',
    (
        (['flair:news-forward'], 2048),
        (['flair:news-forward', 'flair:news-backward'], 4096),
        (['word:glove', 'flair:news-backward'], 2148),
        (['byte-pair:en'], 100),
    ),
)
def test_encoding_models(embeddings: List[str], dim: int):
    docs = DocumentArray([Document(text='hello there')])
    encoder = FlairTextEncoder(embeddings=embeddings)
    encoder.encode(docs, {})

    assert docs[0].embedding.shape == (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: FlairTextEncoder
):
    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: FlairTextEncoder, 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: FlairTextEncoder):
    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]

from pathlib import Path

import numpy as np
import pytest
from jina import DocumentArray, Document, Executor

from ...flair_text import FlairTextEncoder


@pytest.fixture()
def docs_generator():
    return DocumentArray((Document(text='random text') for _ in range(30)))


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


def test_flair_batch(docs_generator):
    encoder = FlairTextEncoder(pooling_strategy='mean')
    docs = docs_generator
    encoder.encode(docs, parameters={'batch_size': 10, 'traversal_paths': ['r']})

    assert len(docs.get_attributes('embedding')) == 30
    assert docs[0].embedding.shape == (100,)


def test_traversal_path():
    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),
    ]

    encoder = FlairTextEncoder()
    encoder.encode(docs, parameters={'batch_size': 10, 'traversal_paths': ['c']})

    for path, count in [[['r'], 0], [['c'], 3], [['cc'], 0]]:
        assert len(docs.traverse_flat(path).get_attributes('embedding')) == count
        if count > 0:
            assert docs.traverse_flat(path).get_attributes('embedding')[0].shape == (
                100,
            )

    encoder.encode(docs, parameters={'batch_size': 10, 'traversal_paths': ['cc']})
    for path, count in [[['r'], 0], [['c'], 3], [['cc'], 2]]:
        assert len(docs.traverse_flat(path).get_attributes('embedding')) == count
        if count > 0:
            assert docs.traverse_flat(path).get_attributes('embedding')[0].shape == (
                100,
            )


def test_no_documents():
    encoder = FlairTextEncoder()
    docs = []
    encoder.encode(docs, parameters={'batch_size': 10, 'traversal_paths': ['r']})
    assert not docs


def test_flair_word_encode():
    docs = []
    words = ['apple', 'banana1', 'banana2', 'studio', 'satelite', 'airplane']
    for word in words:
        docs.append(Document(text=word))

    text_encoder = FlairTextEncoder()
    text_encoder.encode(DocumentArray(docs), {})

    txt_to_ndarray = {}
    for d in docs:
        txt_to_ndarray[d.text] = d.embedding

    def dist(a, b):
        nonlocal txt_to_ndarray
        a_embedding = txt_to_ndarray[a]
        b_embedding = txt_to_ndarray[b]
        return np.linalg.norm(a_embedding - b_embedding)

    # assert semantic meaning is captured in the encoding
    small_distance = dist('banana1', 'banana2')
    assert small_distance < dist('banana1', 'airplane')
    assert small_distance < dist('banana1', 'satelite')
    assert small_distance < dist('banana1', 'studio')
    assert small_distance < dist('banana2', 'airplane')
