from docarray import dataclass, Document, DocumentArray
from docarray.typing import Image, Text


@dataclass
class Page:
    main_text: Text
    image: Image
    description: Text


query_page = Page(
    main_text='Hello world',
    image='testflow.jpg',
    description='This is the image of an apple',
)

query = Document(query_page)  # our query Document

da = DocumentArray(
    [
        Document(
            Page(
                main_text='First page',
                image='apple.png',
                description='This is the image of an apple',
            )
        ),
        Document(
            Page(
                main_text='Second page',
                image='testflow.jpg',
                description='This is an image of a pear',
            )
        ),
    ],
    subindex_configs = {'@.[image]': None})  # specify subindices at docarray creation time


from torchvision.models import resnet50

img_model = resnet50(pretrained=True)

# embed query
query.image.set_image_tensor_shape(shape=(224, 224)).set_image_tensor_channel_axis(
    original_channel_axis=-1, new_channel_axis=0
).set_image_tensor_normalization(channel_axis=0).embed(img_model)

# embed dataset
da['@.[image]'].apply(
    lambda d: d.set_image_tensor_shape(shape=(224, 224))
    .set_image_tensor_channel_axis(original_channel_axis=-1, new_channel_axis=0)
    .set_image_tensor_normalization(channel_axis=0)
).embed(img_model)


closest_match_img = da.find(query.image, on='@.[image]')[0][0]  # search through subindex using `on=`
print('CLOSEST IMAGE:')
closest_match_img.summary()
print('PAGE WITH THE CLOSEST IMAGE:')
closest_match_page = da[closest_match_img.parent_id]
closest_match_page.summary()