# simple_mosaic.py
from pathlib import Path
from typing import List, Tuple, Iterator
import numpy as np
from PIL import Image, ImageDraw
from tile_library import build_cifar10_tile_library, build_cifar100_tile_library

class SimpleMosaicImage:
    def __init__(self, path: str):
        self.path = path
        self.img = Image.open(path).convert("RGB")
        self.width, self.height = self.img.size
        print(f"[INFO] Loaded: {path} | size={self.width}x{self.height}")

    def resize(self, longest_side: int = 4000) -> "SimpleMosaicImage":
        w, h = self.width, self.height
        scale = longest_side / max(w, h)
        if scale < 1.0:
            new_w = int(w * scale)
            new_h = int(h * scale)
            self.img = self.img.resize((new_w, new_h), Image.BICUBIC)
            self.width, self.height = new_w, new_h
            print(f"[INFO] Resized to {new_w}x{new_h}")
        return self

    def quantize_colors(self, n_colors: int = 16) -> "SimpleMosaicImage":
        """Apply color quantization using PIL's built-in algorithm"""
        quantized = self.img.quantize(colors=n_colors, method=Image.MEDIANCUT)
        self.img = quantized.convert('RGB')
        print(f"[INFO] Color quantized to {n_colors} colors")
        return self

    def crop_to_grid(self, grid_size: int = 32) -> "SimpleMosaicImage":
        """Smart boundary handling: preserve original size when possible"""
        # Only crop if loss is minimal (<2%), otherwise preserve original size
        new_w = (self.width // grid_size) * grid_size
        new_h = (self.height // grid_size) * grid_size

        lost_pixels = (self.width - new_w) + (self.height - new_h)
        total_pixels = self.width + self.height
        loss_ratio = lost_pixels / total_pixels

        if loss_ratio < 0.02:  # Only crop if loss < 2%
            self.img = self.img.crop((0, 0, new_w, new_h))
            self.width, self.height = new_w, new_h
            print(f"[INFO] Cropped to {new_w}x{new_h} for grid {grid_size} (loss: {loss_ratio:.1%})")
        else:
            print(f"[INFO] Preserved original size {self.width}x{self.height} (would lose {loss_ratio:.1%})")
        return self

    def _as_array(self):
        return np.asarray(self.img, dtype=np.uint8)

    def iter_cells(self, grid_size: int):
        for y in range(0, self.height, grid_size):
            for x in range(0, self.width, grid_size):
                yield (x, y, grid_size, grid_size)
    
    def draw_cells(self, cells, outline=(255, 0, 0), width=0.1):
        """
        Draw cell borders on the original image, returns a new image.
        outline: border color
        width: border line width
        """
        canvas = self.img.copy()
        draw = ImageDraw.Draw(canvas)
        for (x, y, w, h) in cells:
            # PIL rectangle bottom-right is inclusive, -1 to avoid overflow
            draw.rectangle((x, y, x + w - 1, y + h - 1), outline=outline, width=width)
        return canvas
    
    @staticmethod
    def _cell_mean(arr, x, y, w, h):
        block = arr[y:y+h, x:x+w, :]
        mean = block.mean(axis=(0,1))
        return tuple(int(round(v)) for v in mean)

    @staticmethod
    def _nearest_color(target, palette):
        tr, tg, tb = target
        pal = np.array(palette, dtype=np.int16)
        diff = pal - np.array([tr, tg, tb], dtype=np.int16)
        dist2 = np.sum(diff*diff, axis=1)
        idx = int(np.argmin(dist2))
        return tuple(int(v) for v in pal[idx])
    
    def build_adaptive_cells(
        self,
        start_size: int = 64,
        min_size: int = 16,
        threshold: float = 20.0,     # Use grayscale variance as complexity measure
        ) -> list[tuple[int,int,int,int]]:
        """
        Returns [(x,y,w,h), ...]: Quadtree-style adaptive grid using iterative stack.
        Requirement: Image should be resized/cropped to be divisible by start_size for better alignment.
        """
        arr = self._as_array()
        # Grayscale (BT.601)
        gray = (0.299*arr[...,0] + 0.587*arr[...,1] + 0.114*arr[...,2]).astype(np.float32)

        cells: list[tuple[int,int,int,int]] = []

        # First rough division by start_size, push large blocks to stack
        stack: list[tuple[int,int,int,int]] = []
        for yy in range(0, self.height, start_size):
            for xx in range(0, self.width, start_size):
                ww = min(start_size, self.width  - xx)
                hh = min(start_size, self.height - yy)
                stack.append((xx, yy, ww, hh))

        # Process stack: decide whether to keep or subdivide into 4 blocks
        while stack:
            x, y, w, h = stack.pop()

            # Keep if reached minimum size
            if w <= min_size or h <= min_size:
                cells.append((x, y, w, h))
                continue

            # Complexity: grayscale variance
            region = gray[y:y+h, x:x+w]
            score = float(region.var())

            # Below threshold -> keep without subdivision
            if score < threshold:
                cells.append((x, y, w, h))
                continue

            # Otherwise subdivide into 4 blocks (try to halve), handle boundary remainder
            w2 = max(min_size, w // 2)
            h2 = max(min_size, h // 2)
            # Fallback: keep if cannot subdivide further (avoid infinite loop)
            if w2 == w and h2 == h:
                cells.append((x, y, w, h))
                continue

            # Top-left
            stack.append((x, y, w2, h2))
            # Top-right
            x2 = x + w2
            wR = min(w - w2, self.width - x2)
            if wR > 0:
                stack.append((x2, y, wR, h2))
            # Bottom-left
            y2 = y + h2
            hB = min(h - h2, self.height - y2)
            if hB > 0:
                stack.append((x, y2, w2, hB))
            # Bottom-right
            if wR > 0 and hB > 0:
                stack.append((x2, y2, wR, hB))

        return cells

    def mosaic_average_color_adaptive(self, cells):
        arr = self._as_array()
        out = np.empty_like(arr)
        for (x, y, w, h) in cells:
            color = self._cell_mean(arr, x, y, w, h)
            out[y:y+h, x:x+w, :] = color
        return Image.fromarray(out, mode="RGB")

    def mosaic_with_tiles_adaptive(self, cells, tiles, tile_means: np.ndarray):
        """
        Adaptive grid version: pass in cells from build_adaptive_cells.
        """
        out_img = Image.new("RGB", (self.width, self.height))
        arr = self._as_array()
        means = tile_means.astype(np.float32)

        for (x, y, w, h) in cells:
            block_mean = np.array(self._cell_mean(arr, x, y, w, h), dtype=np.float32)
            diff = means - block_mean[None, :]
            idx = int(np.argmin(np.sum(diff*diff, axis=1)))
            tile = tiles[idx].resize((w, h), Image.BILINEAR)
            out_img.paste(tile, (x, y))
        return out_img
    

    def save(self, image: Image.Image, out_path: str) -> None:
        Path(out_path).parent.mkdir(parents=True, exist_ok=True)
        image.save(out_path)
        print(f"[INFO] Saved: {out_path}")


# loader = SimpleMosaicImage("./samples/akaza.jpg")
# loader.quantize_colors(16).crop_to_grid(2)

# cells = loader.build_adaptive_cells(
#     start_size=64,  # Initial block size
#     min_size=4,    # Minimum block size
#     threshold=5.0  # Lower value = more subdivision
# )

# tiles_10, tile_means_10, tile_labels_10 = build_cifar10_tile_library(max_per_class=1000)
# tiles_100, tile_means_100, tile_labels_100 = build_cifar100_tile_library(max_per_class=400)


# vis1 = loader.draw_cells(cells, outline=(0, 255, 0), width=1)
# vis1.save("./out/cells_outline_akaza.png")

# mosaic_adapt = loader.mosaic_average_color_adaptive(cells)
# loader.save(mosaic_adapt, "./out/mosaic_adaptive_akaza.png")

# mosaic_tiles_adapt_10 = loader.mosaic_with_tiles_adaptive(cells, tiles=tiles_10, tile_means=tile_means_10)
# loader.save(mosaic_tiles_adapt_10, "./out/mosaic_cifar_adaptive_akaza_10.png")

# mosaic_tiles_adapt_100 = loader.mosaic_with_tiles_adaptive(cells, tiles=tiles_100, tile_means=tile_means_100)
# loader.save(mosaic_tiles_adapt_100, "./out/mosaic_cifar_adaptive_akaza_100.png")