src/quadtree_engine.py

"""
Quadtree Image Engine - Python implementation of the C quadtree image manipulation.
Mirrors the logic from compress.c, decompress.c, filters.c, rotate.c, union.c
"""
import numpy as np
from dataclasses import dataclass, field
from typing import Optional
from PIL import Image
import io


@dataclass
class QtNode:
    red: int = 0
    green: int = 0
    blue: int = 0
    area: int = 0
    topLeft: Optional['QtNode'] = field(default=None, repr=False)
    topRight: Optional['QtNode'] = field(default=None, repr=False)
    bottomLeft: Optional['QtNode'] = field(default=None, repr=False)
    bottomRight: Optional['QtNode'] = field(default=None, repr=False)

    def is_leaf(self):
        return (self.topLeft is None and self.topRight is None and
                self.bottomLeft is None and self.bottomRight is None)


def get_mean(matrix: np.ndarray, x: int, y: int, size: int):
    """Compute average color and variance score for a region. Mirrors getMean() in suppl.c"""
    region = matrix[y:y+size, x:x+size]
    red = int(np.mean(region[:, :, 0]))
    green = int(np.mean(region[:, :, 1]))
    blue = int(np.mean(region[:, :, 2]))
    variance = (
        np.mean((region[:, :, 0].astype(int) - red)**2) +
        np.mean((region[:, :, 1].astype(int) - green)**2) +
        np.mean((region[:, :, 2].astype(int) - blue)**2)
    ) / 3
    return red, green, blue, int(variance)


def compress_image(matrix: np.ndarray, x: int, y: int, size: int, threshold: int) -> QtNode:
    """Build quadtree. Mirrors compressImage() in compress.c"""
    red, green, blue, score = get_mean(matrix, x, y, size)
    node = QtNode(red=red, green=green, blue=blue, area=size * size)
    if size > 1 and score > threshold:
        half = size // 2
        node.topLeft = compress_image(matrix, x, y, half, threshold)
        node.topRight = compress_image(matrix, x + half, y, half, threshold)
        node.bottomRight = compress_image(matrix, x + half, y + half, half, threshold)
        node.bottomLeft = compress_image(matrix, x, y + half, half, threshold)
    return node


def decompress_image(node: QtNode, matrix: np.ndarray, x: int, y: int, size: int):
    """Reconstruct pixel matrix from quadtree. Mirrors decompressImage() in decompress.c"""
    if node.is_leaf():
        matrix[y:y+size, x:x+size, 0] = node.red
        matrix[y:y+size, x:x+size, 1] = node.green
        matrix[y:y+size, x:x+size, 2] = node.blue
    else:
        half = size // 2
        decompress_image(node.topLeft, matrix, x, y, half)
        decompress_image(node.topRight, matrix, x + half, y, half)
        decompress_image(node.bottomRight, matrix, x + half, y + half, half)
        decompress_image(node.bottomLeft, matrix, x, y + half, half)


# ── Filters (mirrors filters.c) ──────────────────────────────────────────────

def apply_grayscale(node: QtNode) -> QtNode:
    """True luminance-weighted grayscale.
    Compute single luma value and assign to all 3 channels so the
    output is actually grey (not green-tinted).
    luma = 0.299R + 0.587G + 0.114B  (BT.601 standard)
    """
    res = QtNode(area=node.area)
    if not node.is_leaf():
        res.topLeft = apply_grayscale(node.topLeft)
        res.topRight = apply_grayscale(node.topRight)
        res.bottomLeft = apply_grayscale(node.bottomLeft)
        res.bottomRight = apply_grayscale(node.bottomRight)
    luma = int(0.299 * node.red + 0.587 * node.green + 0.114 * node.blue)
    res.red = luma
    res.green = luma
    res.blue = luma
    return res


def apply_negative(node: QtNode) -> QtNode:
    """Invert channels: 255 - value. Mirrors negativeImage() in filters.c"""
    res = QtNode(area=node.area)
    if not node.is_leaf():
        res.topLeft = apply_negative(node.topLeft)
        res.topRight = apply_negative(node.topRight)
        res.bottomLeft = apply_negative(node.bottomLeft)
        res.bottomRight = apply_negative(node.bottomRight)
    res.red = 255 - node.red
    res.green = 255 - node.green
    res.blue = 255 - node.blue
    return res


def apply_sepia(node: QtNode) -> QtNode:
    """Cinematic warm sepia tone. Mirrors sepia() in filters.c"""
    res = QtNode(area=node.area)
    if not node.is_leaf():
        res.topLeft = apply_sepia(node.topLeft)
        res.topRight = apply_sepia(node.topRight)
        res.bottomLeft = apply_sepia(node.bottomLeft)
        res.bottomRight = apply_sepia(node.bottomRight)
    r, g, b = node.red, node.green, node.blue
    res.red = min(255, int(0.393 * r + 0.769 * g + 0.189 * b))
    res.green = min(255, int(0.272 * r + 0.534 * g + 0.131 * b))
    res.blue = min(255, int(0.349 * r + 0.686 * g + 0.168 * b))
    return res


def apply_brighten(node: QtNode, factor: float = 1.3) -> QtNode:
    """Brighten by scaling channels. Mirrors brighten() in filters.c"""
    res = QtNode(area=node.area)
    if not node.is_leaf():
        res.topLeft = apply_brighten(node.topLeft, factor)
        res.topRight = apply_brighten(node.topRight, factor)
        res.bottomLeft = apply_brighten(node.bottomLeft, factor)
        res.bottomRight = apply_brighten(node.bottomRight, factor)
    res.red = min(255, int(node.red * factor))
    res.green = min(255, int(node.green * factor))
    res.blue = min(255, int(node.blue * factor))
    return res


# ── Spatial Transforms (mirrors rotate.c) ────────────────────────────────────

def get_water_image(node: QtNode):
    """Vertical flip (top-bottom). Mirrors getWaterImage() in rotate.c"""
    if not node.is_leaf():
        get_water_image(node.topLeft)
        get_water_image(node.topRight)
        get_water_image(node.bottomLeft)
        get_water_image(node.bottomRight)
        node.topLeft, node.bottomLeft = node.bottomLeft, node.topLeft
        node.topRight, node.bottomRight = node.bottomRight, node.topRight


def get_mirror_image(node: QtNode):
    """Horizontal mirror (left-right). Mirrors getMirrorImage() in rotate.c"""
    if not node.is_leaf():
        get_mirror_image(node.topLeft)
        get_mirror_image(node.topRight)
        get_mirror_image(node.bottomLeft)
        get_mirror_image(node.bottomRight)
        node.topLeft, node.topRight = node.topRight, node.topLeft
        node.bottomLeft, node.bottomRight = node.bottomRight, node.bottomLeft


def rotate_left(node: QtNode):
    """90° counter-clockwise. Mirrors rotateLeft() in rotate.c"""
    if not node.is_leaf():
        rotate_left(node.topLeft)
        rotate_left(node.topRight)
        rotate_left(node.bottomLeft)
        rotate_left(node.bottomRight)
        tl, tr, bl, br = node.topLeft, node.topRight, node.bottomLeft, node.bottomRight
        node.topLeft = tr
        node.topRight = br
        node.bottomLeft = tl
        node.bottomRight = bl


def rotate_right(node: QtNode):
    """90° clockwise. Mirrors rotateRight() in rotate.c"""
    if not node.is_leaf():
        rotate_right(node.topLeft)
        rotate_right(node.topRight)
        rotate_right(node.bottomLeft)
        rotate_right(node.bottomRight)
        tl, tr, bl, br = node.topLeft, node.topRight, node.bottomLeft, node.bottomRight
        node.topLeft = bl
        node.topRight = tl
        node.bottomLeft = br
        node.bottomRight = tr


# ── Union / Blend (mirrors union.c) ──────────────────────────────────────────

def union_of_images(t1: QtNode, t2: QtNode) -> QtNode:
    """Average-blend two quadtrees. Mirrors unionOfImages() in union.c"""
    res = QtNode(area=min(t1.area, t2.area))
    res.red = (t1.red + t2.red) // 2
    res.green = (t1.green + t2.green) // 2
    res.blue = (t1.blue + t2.blue) // 2
    if not t1.is_leaf() and not t2.is_leaf():
        res.topLeft = union_of_images(t1.topLeft, t2.topLeft)
        res.topRight = union_of_images(t1.topRight, t2.topRight)
        res.bottomLeft = union_of_images(t1.bottomLeft, t2.bottomLeft)
        res.bottomRight = union_of_images(t1.bottomRight, t2.bottomRight)
    elif t1.is_leaf() and not t2.is_leaf():
        res.topLeft = union_of_images(t1, t2.topLeft)
        res.topRight = union_of_images(t1, t2.topRight)
        res.bottomLeft = union_of_images(t1, t2.bottomLeft)
        res.bottomRight = union_of_images(t1, t2.bottomRight)
    elif not t1.is_leaf() and t2.is_leaf():
        res.topLeft = union_of_images(t1.topLeft, t2)
        res.topRight = union_of_images(t1.topRight, t2)
        res.bottomLeft = union_of_images(t1.bottomLeft, t2)
        res.bottomRight = union_of_images(t1.bottomRight, t2)
    return res


# ── I/O Helpers ───────────────────────────────────────────────────────────────

def read_ppm_bytes(data: bytes):
    """
    Parse P6 binary PPM data robustly.
    Returns numpy array (H, W, 3) uint8.
    """
    # First try Pillow (handles PNG/JPG/PPM automatically)
    try:
        img = Image.open(io.BytesIO(data)).convert("RGB")
        return np.array(img, dtype=np.uint8)
    except Exception:
        pass

    # Manual P6 binary PPM parser — reads header byte-by-byte
    pos = 0
    def read_token():
        nonlocal pos
        # Skip whitespace and comments
        while pos < len(data):
            c = data[pos:pos+1]
            if c == b'#':
                while pos < len(data) and data[pos:pos+1] != b'\n':
                    pos += 1
            elif c in (b' ', b'\t', b'\n', b'\r'):
                pos += 1
            else:
                break
        start = pos
        while pos < len(data) and data[pos:pos+1] not in (b' ', b'\t', b'\n', b'\r'):
            pos += 1
        return data[start:pos].decode('ascii')

    magic = read_token()
    if magic != 'P6':
        raise ValueError(f"Not a P6 PPM file (got: {magic!r})")
    w = int(read_token())
    h = int(read_token())
    _maxval = int(read_token())
    # skip exactly one whitespace byte after maxval
    pos += 1
    raw = data[pos:pos + h * w * 3]
    arr = np.frombuffer(raw, dtype=np.uint8).reshape(h, w, 3)
    return arr.copy()  # make writable


def arr_to_pil(arr: np.ndarray) -> Image.Image:
    return Image.fromarray(arr.astype(np.uint8), 'RGB')


def next_power_of_two(n: int) -> int:
    p = 1
    while p < n:
        p <<= 1
    return p


def pad_to_square_pow2(arr: np.ndarray):
    """Pad image to square power-of-2 as required by the quadtree."""
    h, w = arr.shape[:2]
    size = next_power_of_two(max(h, w))
    padded = np.zeros((size, size, 3), dtype=np.uint8)
    padded[:h, :w] = arr
    return padded, h, w


def process_image(matrix: np.ndarray, operation: str, threshold: int,
                  matrix2: Optional[np.ndarray] = None, return_tree: bool = False) -> np.ndarray:
    """
    Full pipeline: pad → compress → transform → decompress → crop.
    Returns the result as a numpy RGB array, and optionally the tree.
    """
    padded, oh, ow = pad_to_square_pow2(matrix)
    size = padded.shape[0]

    tree = compress_image(padded, 0, 0, size, threshold)

    if operation == "grayscale":
        tree = apply_grayscale(tree)
    elif operation == "negative":
        tree = apply_negative(tree)
    elif operation == "sepia":
        tree = apply_sepia(tree)
    elif operation == "brighten":
        tree = apply_brighten(tree)
    elif operation == "mirror":
        get_mirror_image(tree)
    elif operation == "water":
        get_water_image(tree)
    elif operation == "rotate_left":
        rotate_left(tree)
    elif operation == "rotate_right":
        rotate_right(tree)
    elif operation == "union" and matrix2 is not None:
        padded2, oh2, ow2 = pad_to_square_pow2(matrix2)
        s2 = padded2.shape[0]
        s_common = max(size, s2)
        p1 = np.zeros((s_common, s_common, 3), dtype=np.uint8)
        p1[:padded.shape[0], :padded.shape[1]] = padded
        p2 = np.zeros((s_common, s_common, 3), dtype=np.uint8)
        p2[:padded2.shape[0], :padded2.shape[1]] = padded2
        t1 = compress_image(p1, 0, 0, s_common, threshold)
        t2 = compress_image(p2, 0, 0, s_common, threshold)
        tree = union_of_images(t1, t2)
        size = s_common
        # For union, we want the bounding box that encompasses both original images
        oh = max(oh, oh2)
        ow = max(ow, ow2)

    out = np.zeros((size, size, 3), dtype=np.uint8)
    decompress_image(tree, out, 0, 0, size)
    
    if operation == "mirror":
        res = out[:oh, size - ow : size]
    elif operation == "water":
        res = out[size - oh : size, :ow]
    elif operation == "rotate_left":
        res = out[size - ow : size, :oh]
    elif operation == "rotate_right":
        res = out[:ow, size - oh : size]
    else:
        res = out[:oh, :ow]
    
    if return_tree:
        return res, tree
    return res