# Licensed to the LF AI & Data foundation under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional, TypeVar, Union

import numpy as np
import pytest
import torch
from pydantic import parse_obj_as

from docarray import BaseDoc, DocList
from docarray.typing import ImageUrl, NdArray, TorchTensor
from docarray.utils._internal.misc import is_tf_available

tf_available = is_tf_available()
if tf_available:
    import tensorflow as tf

    from docarray.typing import TensorFlowTensor


@pytest.fixture()
def da():
    class Text(BaseDoc):
        text: str

    return DocList[Text]([Text(text=f'hello {i}') for i in range(10)])


def test_iterate(da):
    for doc, doc2 in zip(da, da):
        assert doc.id == doc2.id


def test_append():
    class Text(BaseDoc):
        text: str

    da = DocList[Text]([])

    da.append(Text(text='hello', id='1'))

    assert len(da) == 1
    assert da[0].id == '1'


def test_extend():
    class Text(BaseDoc):
        text: str

    da = DocList[Text]([Text(text='hello', id=str(i)) for i in range(10)])

    da.extend([Text(text='hello', id=str(10 + i)) for i in range(10)])

    assert len(da) == 20
    for da, i in zip(da, range(20)):
        assert da.id == str(i)


def test_extend_itself():
    class Text(BaseDoc):
        text: str

    da = DocList[Text]([Text(text='hello', id=str(i)) for i in range(10)])

    da.extend(da)

    assert len(da) == 20


def test_slice(da):
    da2 = da[0:5]
    assert type(da2) == da.__class__
    assert len(da2) == 5


def test_document_array():
    class Text(BaseDoc):
        text: str

    da = DocList([Text(text='hello') for _ in range(10)])

    assert len(da) == 10


def test_empty_array():
    da = DocList()
    len(da) == 0


def test_document_array_fixed_type():
    class Text(BaseDoc):
        text: str

    da = DocList[Text]([Text(text='hello') for _ in range(10)])

    assert len(da) == 10


def test_ndarray_equality():
    class Text(BaseDoc):
        tensor: NdArray

    arr1 = Text(tensor=np.zeros(5))
    arr2 = Text(tensor=np.zeros(5))
    arr3 = Text(tensor=np.ones(5))
    arr4 = Text(tensor=np.zeros(4))

    assert arr1 == arr2
    assert arr1 != arr3
    assert arr1 != arr4


def test_tensor_equality():
    class Text(BaseDoc):
        tensor: TorchTensor

    torch1 = Text(tensor=torch.zeros(128))
    torch2 = Text(tensor=torch.zeros(128))
    torch3 = Text(tensor=torch.zeros(126))
    torch4 = Text(tensor=torch.ones(128))

    assert torch1 == torch2
    assert torch1 != torch3
    assert torch1 != torch4


def test_documentarray():
    class Text(BaseDoc):
        text: str

    da1 = DocList([Text(text='hello')])
    da2 = DocList([Text(text='hello')])

    assert da1 == da2
    assert da1 == [Text(text='hello') for _ in range(len(da1))]
    assert da2 == [Text(text='hello') for _ in range(len(da2))]


@pytest.mark.tensorflow
def test_tensorflowtensor_equality():
    class Text(BaseDoc):
        tensor: TensorFlowTensor

    tensor1 = Text(tensor=tf.constant([1, 2, 3, 4, 5, 6]))
    tensor2 = Text(tensor=tf.constant([1, 2, 3, 4, 5, 6]))
    tensor3 = Text(tensor=tf.constant([[1.0, 2.0], [3.0, 5.0]]))
    tensor4 = Text(tensor=tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]))

    assert tensor1 == tensor2
    assert tensor1 != tensor3
    assert tensor1 != tensor4


def test_text_tensor():
    class Text1(BaseDoc):
        tensor: NdArray

    class Text2(BaseDoc):
        tensor: TorchTensor

    arr_tensor1 = Text1(tensor=np.zeros(2))
    arr_tensor2 = Text2(tensor=torch.zeros(2))

    assert arr_tensor1 == arr_tensor2


def test_get_bulk_attributes_function():
    class Mmdoc(BaseDoc):
        text: str
        tensor: NdArray

    N = 10

    da = DocList[Mmdoc](
        (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N))
    )

    tensors = da._get_data_column('tensor')

    assert len(tensors) == N
    for tensor in tensors:
        assert tensor.shape == (3, 224, 224)

    texts = da._get_data_column('text')

    assert len(texts) == N
    for i, text in enumerate(texts):
        assert text == f'hello{i}'


def test_set_attributes():
    class InnerDoc(BaseDoc):
        text: str

    class Mmdoc(BaseDoc):
        inner: InnerDoc

    N = 10

    da = DocList[Mmdoc]((Mmdoc(inner=InnerDoc(text=f'hello{i}')) for i in range(N)))

    list_docs = [InnerDoc(text=f'hello{i}') for i in range(N)]
    da._set_data_column('inner', list_docs)

    for doc, list_doc in zip(da, list_docs):
        assert doc.inner == list_doc


def test_get_bulk_attributes():
    class Mmdoc(BaseDoc):
        text: str
        tensor: NdArray

    N = 10

    da = DocList[Mmdoc](
        (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N))
    )

    tensors = da.tensor

    assert len(tensors) == N
    for tensor in tensors:
        assert tensor.shape == (3, 224, 224)

    texts = da.text

    assert len(texts) == N
    for i, text in enumerate(texts):
        assert text == f'hello{i}'


def test_get_bulk_attributes_document():
    class InnerDoc(BaseDoc):
        text: str

    class Mmdoc(BaseDoc):
        inner: InnerDoc

    N = 10

    da = DocList[Mmdoc]((Mmdoc(inner=InnerDoc(text=f'hello{i}')) for i in range(N)))

    assert isinstance(da.inner, DocList)


def test_get_bulk_attributes_optional_type():
    class Mmdoc(BaseDoc):
        text: str
        tensor: Optional[NdArray]

    N = 10

    da = DocList[Mmdoc](
        (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N))
    )

    tensors = da.tensor

    assert len(tensors) == N
    for tensor in tensors:
        assert tensor.shape == (3, 224, 224)

    texts = da.text

    assert len(texts) == N
    for i, text in enumerate(texts):
        assert text == f'hello{i}'


def test_get_bulk_attributes_union_type():
    class Mmdoc(BaseDoc):
        text: str
        tensor: Union[NdArray, TorchTensor]

    N = 10

    da = DocList[Mmdoc](
        (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N))
    )

    tensors = da.tensor

    assert len(tensors) == N
    assert isinstance(tensors, list)
    for tensor in tensors:
        assert tensor.shape == (3, 224, 224)

    texts = da.text

    assert len(texts) == N
    for i, text in enumerate(texts):
        assert text == f'hello{i}'


@pytest.mark.tensorflow
def test_get_bulk_attributes_union_type_nested():
    class MyDoc(BaseDoc):
        embedding: Union[Optional[TorchTensor], Optional[NdArray]]
        embedding2: Optional[Union[TorchTensor, NdArray, TensorFlowTensor]]
        embedding3: Optional[Optional[TorchTensor]]
        embedding4: Union[
            Optional[Union[TorchTensor, NdArray, TensorFlowTensor]], TorchTensor
        ]

    da = DocList[MyDoc](
        [
            MyDoc(
                embedding=torch.rand(10),
                embedding2=torch.rand(10),
                embedding3=torch.rand(10),
                embedding4=torch.rand(10),
            )
            for _ in range(10)
        ]
    )

    for attr in ['embedding', 'embedding2', 'embedding3', 'embedding4']:
        tensors = getattr(da, attr)
        assert len(tensors) == 10
        assert isinstance(tensors, list)
        for tensor in tensors:
            assert tensor.shape == (10,)


def test_get_from_slice():
    class Doc(BaseDoc):
        text: str
        tensor: NdArray

    N = 10

    da = DocList[Doc](
        (Doc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N))
    )

    da_sliced = da[0:10:2]
    assert isinstance(da_sliced, DocList)

    tensors = da_sliced.tensor
    assert len(tensors) == 5
    for tensor in tensors:
        assert tensor.shape == (3, 224, 224)

    texts = da_sliced.text
    assert len(texts) == 5
    for i, text in enumerate(texts):
        assert text == f'hello{i * 2}'


def test_del_item(da):
    assert len(da) == 10
    del da[2]
    assert len(da) == 9
    assert da.text == [
        'hello 0',
        'hello 1',
        'hello 3',
        'hello 4',
        'hello 5',
        'hello 6',
        'hello 7',
        'hello 8',
        'hello 9',
    ]
    del da[0:2]
    assert len(da) == 7
    assert da.text == [
        'hello 3',
        'hello 4',
        'hello 5',
        'hello 6',
        'hello 7',
        'hello 8',
        'hello 9',
    ]


def test_generic_type_var():
    T = TypeVar('T', bound=BaseDoc)

    def f(a: DocList[T]) -> DocList[T]:
        return a

    def g(a: DocList['BaseDoc']) -> DocList['BaseDoc']:
        return a

    a = DocList()
    f(a)
    g(a)


def test_construct():
    class Text(BaseDoc):
        text: str

    docs = [Text(text=f'hello {i}') for i in range(10)] + [BaseDoc()]

    da = DocList[Text].construct(docs)

    assert type(da[-1]) == BaseDoc


def test_reverse():
    class Text(BaseDoc):
        text: str

    docs = [Text(text=f'hello {i}') for i in range(10)]

    da = DocList[Text](docs)
    da.reverse()
    assert da[-1].text == 'hello 0'
    assert da[0].text == 'hello 9'


class Image(BaseDoc):
    tensor: Optional[NdArray] = None
    url: ImageUrl


def test_remove():
    images = [Image(url=f'http://url.com/foo_{i}.png') for i in range(3)]
    da = DocList[Image](images)
    da.remove(images[1])
    assert len(da) == 2
    assert da[0] == images[0]
    assert da[1] == images[2]


def test_pop():
    images = [Image(url=f'http://url.com/foo_{i}.png') for i in range(3)]
    da = DocList[Image](images)
    popped = da.pop(1)
    assert len(da) == 2
    assert popped == images[1]
    assert da[0] == images[0]
    assert da[1] == images[2]


def test_sort():
    images = [
        Image(url=f'http://url.com/foo_{i}.png', tensor=NdArray(i)) for i in [2, 0, 1]
    ]
    da = DocList[Image](images)
    da.sort(key=lambda img: len(img.tensor))
    assert len(da) == 3
    assert da[0].url == 'http://url.com/foo_0.png'
    assert da[1].url == 'http://url.com/foo_1.png'


def test_optional_field():
    from typing import Optional

    from docarray import BaseDoc, DocList
    from docarray.typing import ImageUrl, NdArray

    class Nested(BaseDoc):
        tensor: NdArray

    class Image(BaseDoc):
        url: ImageUrl
        features: Optional[Nested] = None

    docs = DocList[Image]([Image(url='http://url.com/foo.png') for _ in range(10)])

    assert docs.features == [None for _ in range(10)]
    assert isinstance(docs.features, list)
    assert not isinstance(docs.features, DocList)


def test_validate_list_dict():
    images = [
        dict(url=f'http://url.com/foo_{i}.png', tensor=NdArray(i)) for i in [2, 0, 1]
    ]

    # docs = DocList[Image]([Image(url=image['url'], tensor=image['tensor']) for image in images])

    docs = parse_obj_as(DocList[Image], images)

    assert docs.url == [
        'http://url.com/foo_2.png',
        'http://url.com/foo_0.png',
        'http://url.com/foo_1.png',
    ]


def test_legacy_doc():
    from docarray.documents.legacy import LegacyDocument

    newDoc = LegacyDocument()
    da = DocList[LegacyDocument]([newDoc])
    da.summary()


def test_parameterize_list():
    from docarray import DocList, BaseDoc

    with pytest.raises(TypeError) as excinfo:
        da = DocList[BaseDoc()]
        assert da is None

    assert str(excinfo.value) == 'Expecting a type, got object instead'


def test_not_double_subcriptable():
    from docarray import DocList
    from docarray.documents import TextDoc

    with pytest.raises(TypeError) as excinfo:
        da = DocList[TextDoc][TextDoc]
        assert da is None