splitbit_llm/vision/image_gen.py

"""Image Generation — local, NumPy-based procedural image generation.

No external dependencies. Generates images using:
- Procedural patterns (gradients, noise, fractals)
- Text-to-image via prompt interpretation (maps keywords to visual properties)
- SVG generation for vector graphics
- ASCII art generation from text prompts

All generation is 100% local — no API calls to DALL-E, Stable Diffusion, etc.
For production-quality image generation, connect an external service via
the connectors module.
"""

from __future__ import annotations

import base64
import hashlib
import io
import logging
import math
import os
import struct
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)


class ImageGenerator:
    """Local image generator using NumPy.

    Generates images from text prompts using procedural techniques:
    - Color palette extraction from prompt keywords
    - Pattern selection based on prompt themes
    - Composition using gradients, noise, geometric shapes
    - Output as BMP (no external deps) or SVG (vector)

    Prompt interpretation:
    - "sunset" → warm orange/pink gradient
    - "ocean" → blue gradient with wave pattern
    - "forest" → green gradient with noise texture
    - "fire" → red/orange with flicker pattern
    - "abstract" → random colorful shapes
    - "geometric" → structured geometric patterns
    """

    PROMPT_PALETTES = {
        "sunset": [(255, 140, 50), (255, 80, 120), (100, 50, 150)],
        "sunrise": [(255, 180, 80), (255, 120, 100), (150, 100, 200)],
        "ocean": [(20, 80, 180), (40, 120, 200), (80, 180, 220)],
        "sea": [(20, 80, 180), (40, 120, 200), (80, 180, 220)],
        "water": [(40, 100, 180), (60, 140, 200), (100, 180, 220)],
        "forest": [(20, 80, 30), (40, 120, 50), (80, 160, 70)],
        "tree": [(20, 80, 30), (60, 100, 40), (100, 70, 40)],
        "fire": [(255, 50, 0), (255, 120, 0), (255, 200, 50)],
        "flame": [(255, 50, 0), (255, 120, 0), (255, 200, 50)],
        "ice": [(180, 220, 255), (200, 240, 255), (220, 250, 255)],
        "snow": [(200, 220, 240), (220, 240, 250), (240, 250, 255)],
        "night": [(10, 10, 40), (20, 20, 60), (40, 40, 80)],
        "space": [(0, 0, 20), (20, 10, 40), (60, 40, 100)],
        "star": [(0, 0, 20), (40, 40, 80), (255, 255, 200)],
        "desert": [(200, 170, 100), (220, 190, 130), (240, 210, 160)],
        "sky": [(100, 150, 220), (130, 180, 240), (180, 210, 250)],
        "grass": [(40, 120, 30), (60, 160, 40), (100, 200, 60)],
        "mountain": [(80, 70, 60), (120, 110, 100), (160, 150, 140)],
        "abstract": [(255, 50, 100), (50, 200, 255), (255, 200, 50)],
        "geometric": [(50, 50, 150), (150, 50, 100), (50, 150, 200)],
        "rainbow": [(255, 0, 0), (255, 128, 0), (255, 255, 0), (0, 255, 0), (0, 128, 255), (128, 0, 255)],
        "metal": [(100, 100, 110), (140, 140, 150), (180, 180, 190)],
        "gold": [(180, 140, 40), (220, 180, 60), (255, 220, 100)],
        "neon": [(255, 0, 255), (0, 255, 255), (255, 255, 0)],
        "dark": [(10, 10, 15), (20, 20, 30), (40, 40, 50)],
        "light": [(240, 240, 250), (220, 220, 240), (200, 200, 230)],
    }

    PATTERN_TYPES = ["gradient", "radial", "noise", "fractal", "geometric", "waves"]

    def __init__(self, default_size: tuple[int, int] = (256, 256)) -> None:
        self.default_size = default_size
        self._stats = {
            "images_generated": 0,
            "total_pixels_generated": 0,
            "avg_generation_time_s": 0.0,
        }

    def generate(self, prompt: str, width: int = 0, height: int = 0,
                 pattern: str = "") -> dict[str, Any]:
        """Generate an image from a text prompt.

        Args:
            prompt: text description of the image
            width: image width (0 = default)
            height: image height (0 = default)
            pattern: force a specific pattern type
        Returns:
            dict with image data (BMP bytes, base64), metadata
        """
        import time
        t0 = time.time()

        w, h = (width, height) if width and height else self.default_size
        palette = self._extract_palette(prompt)
        pat = pattern or self._select_pattern(prompt)

        # Generate pixel array
        pixels = self._render(pat, w, h, palette)

        # Encode as BMP
        bmp_bytes = self._encode_bmp(pixels)

        # Encode as base64 for embedding in HTML
        b64 = base64.b64encode(bmp_bytes).decode()

        elapsed = time.time() - t0
        self._stats["images_generated"] += 1
        self._stats["total_pixels_generated"] += w * h
        self._stats["avg_generation_time_s"] = (
            (self._stats["avg_generation_time_s"] * (self._stats["images_generated"] - 1) + elapsed)
            / self._stats["images_generated"]
        )

        return {
            "prompt": prompt,
            "width": w,
            "height": h,
            "pattern": pat,
            "palette": palette,
            "format": "bmp",
            "size_bytes": len(bmp_bytes),
            "base64": b64,
            "elapsed_s": round(elapsed, 4),
        }

    def generate_svg(self, prompt: str, width: int = 256, height: int = 256) -> str:
        """Generate an SVG image from a prompt."""
        palette = self._extract_palette(prompt)
        pat = self._select_pattern(prompt)

        svg_parts = [
            f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}">',
        ]

        if pat == "gradient":
            stops = "".join(
                f'<stop offset="{i/(len(palette)-1):.2f}" '
                f'style="stop-color:rgb({r},{g},{b})" />'
                for i, (r, g, b) in enumerate(palette)
            )
            svg_parts.append(f'<defs><linearGradient id="g1">{stops}</linearGradient></defs>')
            svg_parts.append(f'<rect width="{width}" height="{height}" fill="url(#g1)" />')

        elif pat == "radial":
            cx, cy = width // 2, height // 2
            r = min(width, height) // 2
            stops = "".join(
                f'<stop offset="{i/(len(palette)-1):.2f}" '
                f'style="stop-color:rgb({r_},{g_},{b_})" />'
                for i, (r_, g_, b_) in enumerate(palette)
            )
            svg_parts.append(f'<defs><radialGradient id="g2">{stops}</radialGradient></defs>')
            svg_parts.append(f'<rect width="{width}" height="{height}" fill="url(#g2)" />')

        elif pat == "geometric":
            svg_parts.append(f'<rect width="{width}" height="{height}" fill="rgb({palette[0][0]},{palette[0][1]},{palette[0][2]})" />')
            for i, (r, g, b) in enumerate(palette[1:]):
                cx = (width * (i + 1)) // (len(palette) - 1)
                cy = (height * (i + 1)) // (len(palette) - 1)
                radius = min(width, height) // 6
                svg_parts.append(f'<circle cx="{cx}" cy="{cy}" r="{radius}" '
                                 f'fill="rgb({r},{g},{b})" opacity="0.7" />')

        else:
            # Fallback: solid color
            r, g, b = palette[0]
            svg_parts.append(f'<rect width="{width}" height="{height}" fill="rgb({r},{g},{b})" />')

        svg_parts.append('</svg>')
        return "\n".join(svg_parts)

    def generate_ascii(self, prompt: str, width: int = 60, height: int = 20) -> str:
        """Generate ASCII art from a prompt."""
        palette = self._extract_palette(prompt)
        chars = " .:-=+*#%@"
        pixels = self._render("noise", width, height, palette)

        lines: list[str] = []
        for y in range(height):
            line = ""
            for x in range(width):
                brightness = int(np.mean(pixels[y, x])) // 26
                line += chars[min(brightness, len(chars) - 1)]
            lines.append(line)

        return "\n".join(lines)

    def _extract_palette(self, prompt: str) -> list[tuple[int, int, int]]:
        """Extract color palette from prompt keywords."""
        prompt_lower = prompt.lower()
        for keyword, palette in self.PROMPT_PALETTES.items():
            if keyword in prompt_lower:
                return palette
        # Default: blue-purple gradient
        return [(30, 30, 80), (60, 50, 120), (100, 80, 160)]

    def _select_pattern(self, prompt: str) -> str:
        """Select a pattern type based on prompt keywords."""
        prompt_lower = prompt.lower()
        if any(kw in prompt_lower for kw in ["gradient", "sky", "sunset", "sunrise", "dawn", "dusk"]):
            return "gradient"
        if any(kw in prompt_lower for kw in ["radial", "burst", "explosion", "sun", "star"]):
            return "radial"
        if any(kw in prompt_lower for kw in ["noise", "texture", "rough", "chaos", "random"]):
            return "noise"
        if any(kw in prompt_lower for kw in ["fractal", "recursive", "mandelbrot", "julia"]):
            return "fractal"
        if any(kw in prompt_lower for kw in ["geometric", "shape", "circle", "square", "triangle"]):
            return "geometric"
        if any(kw in prompt_lower for kw in ["wave", "ocean", "sea", "water", "ripple"]):
            return "waves"
        # Hash-based selection for variety
        h = int(hashlib.md5(prompt.encode()).hexdigest(), 16) % len(self.PATTERN_TYPES)
        return self.PATTERN_TYPES[h]

    def _render(self, pattern: str, w: int, h: int,
                palette: list[tuple[int, int, int]]) -> np.ndarray:
        """Render a pattern to a pixel array."""
        pixels = np.zeros((h, w, 3), dtype=np.uint8)

        if pattern == "gradient":
            pixels = self._render_gradient(w, h, palette)
        elif pattern == "radial":
            pixels = self._render_radial(w, h, palette)
        elif pattern == "noise":
            pixels = self._render_noise(w, h, palette)
        elif pattern == "fractal":
            pixels = self._render_fractal(w, h, palette)
        elif pattern == "geometric":
            pixels = self._render_geometric(w, h, palette)
        elif pattern == "waves":
            pixels = self._render_waves(w, h, palette)
        else:
            pixels = self._render_gradient(w, h, palette)

        return pixels

    def _render_gradient(self, w: int, h: int,
                         palette: list[tuple[int, int, int]]) -> np.ndarray:
        """Render a vertical gradient."""
        pixels = np.zeros((h, w, 3), dtype=np.uint8)
        for y in range(h):
            t = y / max(h - 1, 1)
            idx = t * (len(palette) - 1)
            i0 = int(idx)
            i1 = min(i0 + 1, len(palette) - 1)
            frac = idx - i0
            r = int(palette[i0][0] * (1 - frac) + palette[i1][0] * frac)
            g = int(palette[i0][1] * (1 - frac) + palette[i1][1] * frac)
            b = int(palette[i0][2] * (1 - frac) + palette[i1][2] * frac)
            pixels[y, :] = [r, g, b]
        return pixels

    def _render_radial(self, w: int, h: int,
                       palette: list[tuple[int, int, int]]) -> np.ndarray:
        """Render a radial gradient."""
        pixels = np.zeros((h, w, 3), dtype=np.float32)
        cx, cy = w / 2, h / 2
        max_dist = math.sqrt(cx**2 + cy**2)
        yy, xx = np.ogrid[:h, :w]
        dist = np.sqrt((xx - cx)**2 + (yy - cy)**2) / max_dist

        for c in range(3):
            channel = np.zeros((h, w), dtype=np.float32)
            for i in range(len(palette) - 1):
                t0 = i / (len(palette) - 1)
                t1 = (i + 1) / (len(palette) - 1)
                mask = (dist >= t0) & (dist <= t1)
                frac = (dist[mask] - t0) / max(t1 - t0, 1e-6)
                channel[mask] = palette[i][c] * (1 - frac) + palette[i + 1][c] * frac
            pixels[:, :, c] = channel

        return np.clip(pixels, 0, 255).astype(np.uint8)

    def _render_noise(self, w: int, h: int,
                      palette: list[tuple[int, int, int]]) -> np.ndarray:
        """Render a noise-based texture."""
        noise = np.random.rand(h, w)
        pixels = np.zeros((h, w, 3), dtype=np.uint8)

        for c in range(3):
            channel = np.zeros((h, w), dtype=np.float32)
            for i in range(len(palette) - 1):
                t0 = i / (len(palette) - 1)
                t1 = (i + 1) / (len(palette) - 1)
                mask = (noise >= t0) & (noise <= t1)
                frac = (noise[mask] - t0) / max(t1 - t0, 1e-6)
                channel[mask] = palette[i][c] * (1 - frac) + palette[i + 1][c] * frac
            pixels[:, :, c] = np.clip(channel, 0, 255)

        return pixels

    def _render_fractal(self, w: int, h: int,
                        palette: list[tuple[int, int, int]]) -> np.ndarray:
        """Render a simple fractal (Mandelbrot-like)."""
        pixels = np.zeros((h, w, 3), dtype=np.uint8)
        max_iter = 50

        # Create coordinate arrays
        x_vals = np.linspace(-2.0, 1.0, w, dtype=np.float32)
        y_vals = np.linspace(-1.5, 1.5, h, dtype=np.float32)
        cx, cy = np.meshgrid(x_vals, y_vals)

        zx = np.zeros((h, w), dtype=np.float32)
        zy = np.zeros((h, w), dtype=np.float32)
        iterations = np.zeros((h, w), dtype=np.float32)

        for i in range(max_iter):
            mask = zx**2 + zy**2 < 4
            zx_new = zx[mask]**2 - zy[mask]**2 + cx[mask]
            zy[mask] = 2 * zx[mask] * zy[mask] + cy[mask]
            zx[mask] = zx_new
            iterations[mask] = i

        norm = iterations / max_iter

        for c in range(3):
            channel = np.zeros((h, w), dtype=np.float32)
            for i in range(len(palette) - 1):
                t0 = i / (len(palette) - 1)
                t1 = (i + 1) / (len(palette) - 1)
                mask = (norm >= t0) & (norm <= t1)
                frac = (norm[mask] - t0) / max(t1 - t0, 1e-6)
                channel[mask] = palette[i][c] * (1 - frac) + palette[i + 1][c] * frac
            pixels[:, :, c] = np.clip(channel, 0, 255)

        return pixels

    def _render_geometric(self, w: int, h: int,
                          palette: list[tuple[int, int, int]]) -> np.ndarray:
        """Render geometric shapes."""
        pixels = np.zeros((h, w, 3), dtype=np.uint8)
        # Background
        r, g, b = palette[0]
        pixels[:, :] = [r, g, b]

        # Draw circles
        for i, (r, g, b) in enumerate(palette[1:]):
            cx = (w * (i + 1)) // (len(palette) - 1)
            cy = (h * (i + 1)) // (len(palette) - 1)
            radius = min(w, h) // 6
            yy, xx = np.ogrid[:h, :w]
            mask = (xx - cx)**2 + (yy - cy)**2 <= radius**2
            pixels[mask] = [r, g, b]

        return pixels

    def _render_waves(self, w: int, h: int,
                      palette: list[tuple[int, int, int]]) -> np.ndarray:
        """Render a wave pattern."""
        pixels = np.zeros((h, w, 3), dtype=np.float32)
        yy, xx = np.ogrid[:h, :w]

        wave = (np.sin(xx * 0.05) + np.sin(yy * 0.03) + np.sin((xx + yy) * 0.02)) / 3
        wave = (wave + 1) / 2  # normalize to 0-1

        for c in range(3):
            channel = np.zeros((h, w), dtype=np.float32)
            for i in range(len(palette) - 1):
                t0 = i / (len(palette) - 1)
                t1 = (i + 1) / (len(palette) - 1)
                mask = (wave >= t0) & (wave <= t1)
                frac = (wave[mask] - t0) / max(t1 - t0, 1e-6)
                channel[mask] = palette[i][c] * (1 - frac) + palette[i + 1][c] * frac
            pixels[:, :, c] = channel

        return np.clip(pixels, 0, 255).astype(np.uint8)

    def _encode_bmp(self, pixels: np.ndarray) -> bytes:
        """Encode a pixel array as BMP format (no external deps)."""
        h, w = pixels.shape[:2]
        row_size = (w * 3 + 3) & ~3  # BMP rows are padded to 4 bytes
        pixel_data_size = row_size * h
        file_size = 54 + pixel_data_size

        bmp = io.BytesIO()
        # BMP header
        bmp.write(b"BM")
        bmp.write(struct.pack("<I", file_size))
        bmp.write(struct.pack("<HH", 0, 0))
        bmp.write(struct.pack("<I", 54))
        # DIB header
        bmp.write(struct.pack("<I", 40))
        bmp.write(struct.pack("<i", w))
        bmp.write(struct.pack("<i", h))
        bmp.write(struct.pack("<HH", 1, 24))
        bmp.write(struct.pack("<I", 0))
        bmp.write(struct.pack("<I", pixel_data_size))
        bmp.write(struct.pack("<i", 2835))
        bmp.write(struct.pack("<i", 2835))
        bmp.write(struct.pack("<I", 0))
        bmp.write(struct.pack("<I", 0))

        # Pixel data (BMP is bottom-up, BGR)
        for y in range(h - 1, -1, -1):
            row = pixels[y]
            for x in range(w):
                r, g, b = int(row[x, 0]), int(row[x, 1]), int(row[x, 2])
                bmp.write(struct.pack("BBB", b, g, r))
            # Pad row to 4-byte boundary
            padding = row_size - w * 3
            bmp.write(b"\x00" * padding)

        return bmp.getvalue()

    def save(self, image_data: dict[str, Any], path: str) -> None:
        """Save generated image data to a file."""
        if image_data.get("format") == "bmp":
            b64 = image_data["base64"]
            with open(path, "wb") as f:
                f.write(base64.b64decode(b64))
        elif image_data.get("format") == "svg":
            with open(path, "w") as f:
                f.write(image_data["svg"])

    def get_stats(self) -> dict[str, Any]:
        return {**self._stats}