src/backends/geometric.py

"""Geometric fallback backend — no GPU and no HF token required.

When neither a CUDA GPU (local Qwen) nor an HF token (serverless Inference
Providers) is available, AngleForge would otherwise have no usable engine.
This backend approximates each camera-angle preset with cheap Pillow
geometric transforms (perspective tilt, rotation, zoom, translation) so the
Space is always functional as a bootstrap. It is **not** AI image editing —
results are geometric approximations of the requested viewpoint.
"""

from __future__ import annotations

from typing import Dict, List

import numpy as np
from PIL import Image

from ..config import ANGLE_PRESETS
from .base import ImageEditBackend


def _find_coeffs(dst: List[tuple], src: List[tuple]) -> List[float]:
    matrix = []
    for (dx, dy), (sx, sy) in zip(dst, src):
        matrix.append([dx, dy, 1, 0, 0, 0, -sx * dx, -sx * dy])
        matrix.append([0, 0, 0, dx, dy, 1, -sy * dx, -sy * dy])
    a = np.array(matrix, dtype=float)
    b = np.array(src, dtype=float).reshape(8)
    res, *_ = np.linalg.lstsq(a, b, rcond=None)
    return res.tolist()


def _perspective(img: Image.Image, src_quad: List[tuple]) -> Image.Image:
    w, h = img.size
    dst = [(0, 0), (w, 0), (w, h), (0, h)]
    coeffs = _find_coeffs(dst, src_quad)
    return img.transform((w, h), Image.PERSPECTIVE, coeffs, resample=Image.BICUBIC)


def _tilt(img: Image.Image, top_inset: float, bottom_inset: float) -> Image.Image:
    w, h = img.size
    src = [
        (w * top_inset, 0),
        (w * (1 - top_inset), 0),
        (w * (1 - bottom_inset), h),
        (w * bottom_inset, h),
    ]
    return _perspective(img, src)


def _zoom(img: Image.Image, factor: float) -> Image.Image:
    w, h = img.size
    if factor >= 1.0:  # crop in, then scale back up
        cw, ch = int(w / factor), int(h / factor)
        left, top = (w - cw) // 2, (h - ch) // 2
        return img.crop((left, top, left + cw, top + ch)).resize((w, h), Image.LANCZOS)
    # zoom out: paste shrunk image onto a padded canvas
    sw, sh = int(w * factor), int(h * factor)
    small = img.resize((sw, sh), Image.LANCZOS)
    canvas = Image.new("RGB", (w, h), (20, 20, 20))
    canvas.paste(small, ((w - sw) // 2, (h - sh) // 2))
    return canvas


def _shift(img: Image.Image, dx_frac: float, dy_frac: float) -> Image.Image:
    w, h = img.size
    dx, dy = int(w * dx_frac), int(h * dy_frac)
    return img.transform(
        (w, h), Image.AFFINE, (1, 0, -dx, 0, 1, -dy), resample=Image.BICUBIC
    )


def _transform_for_key(img: Image.Image, key: str) -> Image.Image:
    if key in ("top_down", "birds_eye"):
        return _tilt(img, top_inset=0.0, bottom_inset=0.20 if key == "top_down" else 0.12)
    if key == "worms_eye":
        return _tilt(img, top_inset=0.16, bottom_inset=0.0)
    if key == "rotate_left_45":
        return img.rotate(45, resample=Image.BICUBIC, expand=False)
    if key == "rotate_right_45":
        return img.rotate(-45, resample=Image.BICUBIC, expand=False)
    if key == "rotate_left_90":
        return img.rotate(90, resample=Image.BICUBIC, expand=False)
    if key == "rotate_right_90":
        return img.rotate(-90, resample=Image.BICUBIC, expand=False)
    if key == "close_up":
        return _zoom(img, 1.45)
    if key == "wide_angle":
        return _zoom(img, 0.7)
    if key == "move_left":
        return _shift(img, dx_frac=0.15, dy_frac=0.0)
    if key == "move_right":
        return _shift(img, dx_frac=-0.15, dy_frac=0.0)
    if key == "move_forward":
        return _zoom(img, 1.2)
    if key == "move_down":
        return _shift(img, dx_frac=0.0, dy_frac=-0.15)
    return img  # original / unknown


class GeometricBackend(ImageEditBackend):
    """Token-free, CPU-only viewpoint approximation using Pillow transforms."""

    source = "geometric_fallback"

    def __init__(self, image_size: int = 512) -> None:
        self.image_size = image_size
        # Reverse map: bilingual prompt -> preset key.
        self._prompt_to_key: Dict[str, str] = {
            prompt: key for key, (_label, prompt) in ANGLE_PRESETS.items()
        }

    def prepare(self) -> None:
        return None

    def edit(
        self,
        image: Image.Image,
        prompt: str,
        seed: int,
        num_inference_steps: int,
        true_guidance_scale: float,
    ) -> Image.Image:
        img = image.convert("RGB")
        if not prompt or not prompt.strip():
            return img
        key = self._prompt_to_key.get(prompt.strip(), "original")
        return _transform_for_key(img, key)