mosaic_generator/mosaic_builder.py

"""
mosaic_builder.py

Reconstructs the final mosaic image by placing pre-scaled tiles into their
corresponding grid-cell positions.
"""

import numpy as np


class MosaicBuilder:
    """
    MosaicBuilder assembles the output mosaic using:
    - FAISS-selected tile indices
    - Pre-resized tiles (from TileManager)
    - Grid/cell dimensions
    """

    def __init__(self, tm):
        """
        Parameters
        ----------
        tm : TileManager
            A TileManager instance containing pre-scaled tiles and FAISS index.
        """
        self.tm = tm

    # -------------------------------------------------------------
    # MAIN MOSAIC RECONSTRUCTION
    # -------------------------------------------------------------
    def build(self, tile_indices, dims, grid_n):
        """
        Construct final mosaic image using selected tile indices.

        Parameters
        ----------
        tile_indices : np.ndarray
            Flattened array of selected tile indices (length = grid_n * grid_n).
        dims : tuple
            (W, H, cell_w, cell_h):
                W, H     → final image width & height
                cell_w   → width of each grid cell
                cell_h   → height of each grid cell
        grid_n : int
            Number of cells per side in the mosaic.

        Returns
        -------
        np.ndarray
            Final mosaic as an RGB array of shape (H, W, 3).

        Raises
        ------
        ValueError
            If tile indices, dims, or pre-scaled tiles are invalid.
        RuntimeError
            If tiles have not been pre-resized by TileManager.
        """

        # ------------------ VALIDATION ------------------
        if tile_indices is None or not isinstance(tile_indices, np.ndarray):
            raise ValueError("tile_indices must be a NumPy array.")

        expected_len = grid_n * grid_n
        if tile_indices.size != expected_len:
            raise ValueError(
                f"Expected {expected_len} tile indices, got {tile_indices.size}."
            )

        if self.tm.pre_scaled_tiles is None:
            raise RuntimeError(
                "Tiles have not been resized. Call TileManager.prepare_scaled_tiles() first."
            )

        if not isinstance(dims, tuple) or len(dims) != 4:
            raise ValueError("dims must be a tuple of (W, H, cell_w, cell_h).")

        w, h, cell_w, cell_h = dims
        if any(x <= 0 for x in [w, h, cell_w, cell_h]):
            raise ValueError(f"Invalid dims values: {dims}")

        # ------------------ OUTPUT CANVAS ------------------
        out = np.zeros((h, w, 3), dtype=np.uint8)

        # ------------------ PLACE TILES ------------------
        k = 0
        for gy in range(grid_n):
            for gx in range(grid_n):
                idx = tile_indices[k]

                if idx < 0 or idx >= len(self.tm.pre_scaled_tiles):
                    raise ValueError(f"Tile index {idx} out of range.")

                tile = self.tm.pre_scaled_tiles[idx]

                out[
                    gy * cell_h:(gy + 1) * cell_h,
                    gx * cell_w:(gx + 1) * cell_w
                ] = tile

                k += 1

        return out