import numpy as np
import pytest

from doctr.models.detection.differentiable_binarization.base import DBPostProcessor
from doctr.models.detection.linknet.base import LinkNetPostProcessor


def test_dbpostprocessor():
    postprocessor = DBPostProcessor(assume_straight_pages=True)
    r_postprocessor = DBPostProcessor(assume_straight_pages=False)
    with pytest.raises(AssertionError):
        postprocessor(np.random.rand(2, 512, 512).astype(np.float32))
    mock_batch = np.random.rand(2, 512, 512, 1).astype(np.float32)
    out = postprocessor(mock_batch)
    r_out = r_postprocessor(mock_batch)
    # Batch composition
    assert isinstance(out, list)
    assert len(out) == 2
    assert all(isinstance(sample, list) and all(isinstance(v, np.ndarray) for v in sample) for sample in out)
    assert all(all(v.shape[1] == 5 for v in sample) for sample in out)
    assert all(all(v.shape[1] == 4 and v.shape[2] == 2 for v in sample) for sample in r_out)
    # Relative coords
    assert all(all(np.all(np.logical_and(v[:, :4] >= 0, v[:, :4] <= 1)) for v in sample) for sample in out)
    assert all(all(np.all(np.logical_and(v[:, :4] >= 0, v[:, :4] <= 1)) for v in sample) for sample in r_out)
    # Repr
    assert repr(postprocessor) == "DBPostProcessor(bin_thresh=0.3, box_thresh=0.1)"
    # Edge case when the expanded points of the polygon has two lists
    issue_points = np.array(
        [
            [869, 561],
            [923, 581],
            [925, 595],
            [915, 583],
            [889, 583],
            [905, 593],
            [882, 601],
            [901, 595],
            [904, 604],
            [876, 608],
            [915, 614],
            [911, 605],
            [925, 601],
            [930, 616],
            [911, 617],
            [900, 636],
            [931, 637],
            [904, 649],
            [932, 649],
            [932, 628],
            [918, 627],
            [934, 624],
            [935, 573],
            [909, 569],
            [934, 562],
        ],
        dtype=np.int32,
    )
    out = postprocessor.polygon_to_box(issue_points)
    r_out = r_postprocessor.polygon_to_box(issue_points)
    assert isinstance(out, tuple) and len(out) == 4
    assert isinstance(r_out, np.ndarray) and r_out.shape == (4, 2)


def test_linknet_postprocessor():
    postprocessor = LinkNetPostProcessor()
    r_postprocessor = LinkNetPostProcessor(assume_straight_pages=False)
    with pytest.raises(AssertionError):
        postprocessor(np.random.rand(2, 512, 512).astype(np.float32))
    mock_batch = np.random.rand(2, 512, 512, 1).astype(np.float32)
    out = postprocessor(mock_batch)
    r_out = r_postprocessor(mock_batch)
    # Batch composition
    assert isinstance(out, list)
    assert len(out) == 2
    assert all(isinstance(sample, list) and all(isinstance(v, np.ndarray) for v in sample) for sample in out)
    assert all(all(v.shape[1] == 5 for v in sample) for sample in out)
    assert all(all(v.shape[1] == 4 and v.shape[2] == 2 for v in sample) for sample in r_out)
    # Relative coords
    assert all(all(np.all(np.logical_and(v[:4] >= 0, v[:4] <= 1)) for v in sample) for sample in out)
