import logging

import numpy as np
from pymupdf import Document

from yadt.document_il import il_version_1
from yadt.translation_config import TranslationConfig

logger = logging.getLogger(__name__)


class LayoutParser:
    stage_name = "解析页面布局"

    def __init__(self, translation_config: TranslationConfig):
        self.translation_config = translation_config
        self.model = translation_config.doc_layout_model

    def process(self, docs: il_version_1.Document, mupdf_doc: Document):
        """Generate layouts for all pages that need to be translated."""
        # Get pages that need to be translated
        pages_to_translate = [
            page
            for page in docs.page
            if self.translation_config.should_translate_page(page.page_number + 1)
        ]
        total = len(pages_to_translate)
        with self.translation_config.progress_monitor.stage_start(
            self.stage_name, total
        ) as progress:

            # Process pages in batches
            batch_size = 48
            for i in range(0, total, batch_size):
                self.translation_config.raise_if_cancelled()
                batch_pages = pages_to_translate[i : i + batch_size]

                # Prepare batch images
                batch_images = []
                for page in batch_pages:
                    pix = mupdf_doc[page.page_number].get_pixmap()
                    image = np.fromstring(pix.samples, np.uint8).reshape(
                        pix.height, pix.width, 3
                    )[:, :, ::-1]
                    batch_images.append(image)

                # Get predictions for the batch
                layouts_batch = self.model.predict(batch_images, batch_size=batch_size)

                # Process predictions for each page
                for page, layouts in zip(batch_pages, layouts_batch):
                    page_layouts = []
                    for layout in layouts.boxes:
                        # Convert the coordinate system from the picture coordinate
                        # system to the il coordinate system
                        x0, y0, x1, y1 = layout.xyxy
                        pix = mupdf_doc[page.page_number].get_pixmap()
                        h, w = pix.height, pix.width
                        x0, y0, x1, y1 = (
                            np.clip(int(x0 - 1), 0, w - 1),
                            np.clip(int(h - y1 - 1), 0, h - 1),
                            np.clip(int(x1 + 1), 0, w - 1),
                            np.clip(int(h - y0 + 1), 0, h - 1),
                        )
                        page_layout = il_version_1.PageLayout(
                            id=len(page_layouts) + 1,
                            box=il_version_1.Box(
                                x0.item(), y0.item(), x1.item(), y1.item()
                            ),
                            conf=layout.conf.item(),
                            class_name=layouts.names[layout.cls],
                        )
                        page_layouts.append(page_layout)

                    page.page_layout = page_layouts
                    progress.advance(1)

            return docs
