"""Compare extracted answers against the same Poppler view used for labels."""

import tempfile
import unittest
from pathlib import Path

import numpy as np
import pymupdf
from pdf2image import convert_from_path

from copienator.commands.splitting_int import (
    ANSWER_TOP_PADDING_POINTS,
    _render_split_outputs,
)


class SplittingGeometryTests(unittest.TestCase):
    def make_source(self, path, rotation, cropped=True):
        with pymupdf.open() as document:
            page = document.new_page(width=800, height=1000)
            # Asymmetric colours exercise both translation and orientation.
            for row in range(20):
                for column in range(8):
                    colour = (row / 20, column / 8, (row + column) % 7 / 7)
                    page.draw_rect(
                        pymupdf.Rect(column * 100, row * 50,
                                     (column + 1) * 100, (row + 1) * 50),
                        color=None, fill=colour,
                    )
            page.insert_text((200, 500), "Middle of the answer")
            if cropped:
                page.set_cropbox(pymupdf.Rect(20, 80, 760, 920))
            page.set_rotation(rotation)
            document.save(path)

    def assert_matches_preview(self, answer, preview, bounds):
        actual = np.asarray(convert_from_path(answer, dpi=72)[0]).astype(float)
        expected = np.asarray(preview.crop(bounds)).astype(float)
        self.assertEqual(actual.shape, expected.shape)
        self.assertLess(np.abs(actual - expected).mean(), 0.5)

    def test_cropped_and_uncropped_pages_at_every_rotation(self):
        visible = {
            0: (20, 80, 760, 920),
            90: (80, 20, 920, 760),
            180: (40, 80, 780, 920),
            270: (80, 40, 920, 780),
        }
        for rotation in (0, 90, 180, 270):
            for cropped in (False, True):
                for full_answer in (False, True):
                    with self.subTest(rotation=rotation, cropped=cropped, full=full_answer):
                        with tempfile.TemporaryDirectory() as temporary:
                            root = Path(temporary)
                            source = root / "copy.pdf"
                            self.make_source(source, rotation, cropped)
                            before = source.read_bytes()
                            preview = convert_from_path(source, dpi=72, use_cropbox=False)[0]
                            width, height = preview.size
                            bounds = visible[rotation] if cropped else (0, 0, width, height)
                            if full_answer:
                                coordinates = [("A", 0, 0, 10, 0, 100)]
                            else:
                                coordinates = [("A", 0, 250, 260, 0, 100),
                                               ("_", 0, 711, 721, 0, 100)]
                                answer_top = int(height // 4 - ANSWER_TOP_PADDING_POINTS)
                                bounds = (bounds[0], max(bounds[1], answer_top),
                                          bounds[2], min(bounds[3], height * 3 // 4))
                            _render_split_outputs(source, coordinates, root)
                            self.assert_matches_preview(root / "A.pdf", preview, bounds)
                            self.assertEqual(source.read_bytes(), before)

    def test_answer_continues_to_bottom_then_next_page(self):
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            source = root / "copy.pdf"
            first = root / "first.pdf"
            second = root / "second.pdf"
            self.make_source(first, 180)
            self.make_source(second, 0)
            with pymupdf.open() as document:
                for path in (first, second):
                    with pymupdf.open(path) as part:
                        document.insert_pdf(part)
                document.save(source)
            previews = convert_from_path(source, dpi=72, use_cropbox=False)
            _render_split_outputs(source, [("A", 0, 700, 720, 0, 100),
                                           ("_", 1, 211, 230, 500, 600)], root)
            rendered = convert_from_path(root / "A.pdf", dpi=72)
            self.assertEqual(len(rendered), 2)
            for actual, preview, bounds in zip(rendered, previews,
                                               [(40, 688, 780, 920), (20, 80, 760, 250)]):
                expected = np.asarray(preview.crop(bounds)).astype(float)
                actual = np.asarray(actual).astype(float)
                self.assertEqual(actual.shape, expected.shape)
                self.assertLess(np.abs(actual - expected).mean(), 0.5)