"""LayoutLMv3-based cell encoder for spreadsheet structure understanding.

Wraps ``microsoft/layoutlmv3-base`` to produce per-cell embeddings that
fuse text content with 2-D spatial layout information.  Each spreadsheet
cell is treated as a single "word" with its bounding box derived from
row/col coordinates.

Supports frozen (feature extraction) and fine-tunable modes.
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn


def _try_load_layoutlm():
    """Lazy import to avoid hard dependency on transformers."""
    try:
        from transformers import LayoutLMv3Model, LayoutLMv3Tokenizer
        return LayoutLMv3Model, LayoutLMv3Tokenizer
    except ImportError:
        return None, None


class LayoutLMCellEncoder(nn.Module):
    """Encode spreadsheet cells using LayoutLMv3.

    Each cell's text and bounding box are passed through the model to
    produce a 768-d (base) or 1024-d (large) embedding per cell.

    Args:
        model_name: HuggingFace model identifier.
        frozen: If True, freeze all LayoutLM parameters (feature extraction).
        max_cells: Maximum cells per sheet (truncate if exceeded).
    """

    MODEL_DIM = 768  # layoutlmv3-base

    def __init__(
        self,
        model_name: str = "microsoft/layoutlmv3-base",
        frozen: bool = True,
        max_cells: int = 512,
    ):
        super().__init__()
        LayoutLMv3Model, LayoutLMv3Tokenizer = _try_load_layoutlm()
        if LayoutLMv3Model is None:
            raise ImportError(
                "LayoutLMv3 requires `pip install transformers`. "
                "Install it to use the LayoutLM cell encoder."
            )

        self.tokenizer = LayoutLMv3Tokenizer.from_pretrained(model_name)
        self.model = LayoutLMv3Model.from_pretrained(model_name)
        self.max_cells = max_cells

        if frozen:
            for param in self.model.parameters():
                param.requires_grad = False
            self.model.eval()

        self._frozen = frozen

    def _cells_to_inputs(
        self,
        texts: list[str],
        rows: list[int],
        cols: list[int],
        max_row: int,
        max_col: int,
    ) -> dict[str, torch.Tensor]:
        """Convert cell data to LayoutLMv3 input format.

        Bounding boxes are normalized to [0, 1000] range (LayoutLM convention).
        Each cell occupies one "word" position.
        """
        n = min(len(texts), self.max_cells)

        row_scale = 1000.0 / max(max_row, 1)
        col_scale = 1000.0 / max(max_col, 1)

        words = []
        boxes = []
        for i in range(n):
            word = texts[i][:50] if texts[i] else "[empty]"
            words.append(word)

            x0 = int(cols[i] * col_scale)
            y0 = int(rows[i] * row_scale)
            x1 = min(x0 + int(col_scale), 1000)
            y1 = min(y0 + int(row_scale), 1000)
            boxes.append([x0, y0, x1, y1])

        encoding = self.tokenizer(
            words,
            boxes=boxes,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=self.max_cells + 2,
            is_split_into_words=True,
        )
        return encoding

    def forward(
        self,
        texts: list[str],
        rows: list[int],
        cols: list[int],
        max_row: int = 100,
        max_col: int = 100,
    ) -> torch.Tensor:
        """Produce (N, 768) cell embeddings.

        Args:
            texts: Cell text values.
            rows: Cell row indices.
            cols: Cell column indices.
            max_row: Maximum row in the sheet.
            max_col: Maximum column in the sheet.

        Returns:
            Tensor of shape (N, 768) with per-cell embeddings.
        """
        n = len(texts)
        inputs = self._cells_to_inputs(texts, rows, cols, max_row, max_col)
        device = next(self.model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        if self._frozen:
            with torch.no_grad():
                outputs = self.model(**inputs)
        else:
            outputs = self.model(**inputs)

        hidden = outputs.last_hidden_state[0]
        word_ids = inputs.get("input_ids", None)

        cell_embeddings = []
        token_idx = 1
        for i in range(min(n, self.max_cells)):
            if token_idx < hidden.size(0) - 1:
                cell_embeddings.append(hidden[token_idx])
                token_idx += 1
            else:
                cell_embeddings.append(torch.zeros(self.MODEL_DIM, device=device))

        while len(cell_embeddings) < n:
            cell_embeddings.append(torch.zeros(self.MODEL_DIM, device=device))

        return torch.stack(cell_embeddings[:n])

    @staticmethod
    def is_available() -> bool:
        """Check if LayoutLMv3 dependencies are installed."""
        m, _ = _try_load_layoutlm()
        return m is not None