import random

import fitz
import numpy as np
import pytest

from image_prediction.estimator.preprocessor.utils import images_to_batch_tensor
from image_prediction.extraction import extract_images_from_pdf
from image_prediction.image_extractor.extractors.parsable import extract_pages


@pytest.mark.parametrize("extractor_type", ["mock"])
@pytest.mark.parametrize("batch_size", [1, 2, 16])
def test_image_extractor_mock(image_extractor, images):
    images_extracted, metadata = map(list, zip(*image_extractor(images)))
    assert images_extracted == images


@pytest.mark.parametrize("extractor_type", ["parsable_pdf", "default"])
@pytest.mark.parametrize("input_size", [{"depth": 3, "width": 170, "height": 220}], indirect=["input_size"])
@pytest.mark.parametrize("alpha", [False, True])
def test_parsable_pdf_image_extractor(image_extractor, pdf, images, metadata, input_size, alpha):
    images_extracted, metadata_extracted = map(list, extract_images_from_pdf(pdf, image_extractor))
    if not alpha:
        assert np.allclose(images_to_batch_tensor(images_extracted), images_to_batch_tensor(images))
    assert list(metadata_extracted) == metadata


@pytest.mark.parametrize("batch_size", [1, 2, 16])
def test_extract_pages(pdf):
    doc = fitz.Document(stream=pdf)

    max_index = max(0, doc.page_count - 1)
    i = random.randint(0, max(0, max_index - 1))
    j = random.randint(i + 1, max_index) if max_index > 0 else 0

    page_range = range(i, j)

    pages = list(extract_pages(doc, page_range))
    assert all((isinstance(p, fitz.Page) for p in pages))
    assert len(pages) == len(page_range)