PixelVarianceOptimizer.py

"""
PixelVarianceOptimizer.py — Pixel-level text embedding
Métodos avanzados sin tipografía visible:
  1. LSB (Least Significant Bit) — altera solo el bit menos significativo de cada canal
  2. DCT (Discrete Cosine Transform) — embedding en coeficientes de frecuencia media
  3. Adaptive noise — ruido gaussiano orientado por el texto codificado

Ninguno de los tres produce letras visibles: el OCR los lee
a través de patrones de varianza estadística en la matriz de píxeles.
"""

from __future__ import annotations

import math
import struct
import zlib
from dataclasses import dataclass
from enum import Enum
from typing import Literal

import numpy as np
from PIL import Image
from scipy.fft import dctn, idctn


class EmbedMethod(str, Enum):
    LSB      = "lsb"        # bit menos significativo
    DCT      = "dct"        # coeficientes DCT de frecuencia media
    ADAPTIVE = "adaptive"   # ruido gaussiano adaptativo + LSB
    ALL      = "all"        # los tres en cascada (máx densidad)


@dataclass
class PVOConfig:
    method:          EmbedMethod = EmbedMethod.LSB
    lsb_bits:        int   = 1        # cuántos bits LSB usar (1-3); 1 = mínimo artefacto
    dct_strength:    float = 0.8      # amplitud del delta en coeficientes DCT (0.5-3.0)
    dct_block:       int   = 8        # tamaño del bloque DCT (8 estándar)
    noise_sigma:     float = 1.2      # desviación gaussiana para adaptive (0.5-3.0)
    channel:         int   = 0        # canal RGB donde embedir (0=R, 1=G, 2=B, -1=todos)
    output_format:   Literal["PNG"] = "PNG"


class PixelVarianceOptimizer:
    """
    Embide texto en los píxeles de una imagen a nivel matemático
    sin renderizar ningún glifo tipográfico.

    El resultado es indistinguible para el ojo humano pero genera
    patrones estadísticos que los modelos de visión (CLIP, ViT, etc.)
    y los motores OCR de alta sensibilidad pueden extraer.
    """

    _HEADER_MAGIC = b"PVO1"   # 4 bytes de identificación en el stream

    def __init__(self, config: PVOConfig | None = None):
        self.cfg = config or PVOConfig()

    # ------------------------------------------------------------------
    # API pública
    # ------------------------------------------------------------------
    def embed(self, image: Image.Image, text: str) -> Image.Image:
        """
        Embide `text` en `image` según el método configurado.
        Siempre devuelve RGBA para conservar la precisión de píxel.
        """
        arr = self._to_array(image)
        payload = self._encode_payload(text)

        if self.cfg.method == EmbedMethod.LSB:
            arr = self._embed_lsb(arr, payload)
        elif self.cfg.method == EmbedMethod.DCT:
            arr = self._embed_dct(arr, payload)
        elif self.cfg.method == EmbedMethod.ADAPTIVE:
            arr = self._embed_adaptive(arr, payload)
        elif self.cfg.method == EmbedMethod.ALL:
            # cascada: cada método usa canales distintos para evitar colisiones
            arr = self._embed_lsb(arr, payload, channel=0)
            arr = self._embed_dct(arr, payload, channel=1)
            arr = self._embed_adaptive(arr, payload, channel=2)

        return self._from_array(arr)

    # ------------------------------------------------------------------
    # Método 1: LSB
    # ------------------------------------------------------------------
    def _embed_lsb(
        self, arr: np.ndarray, payload: bytes, channel: int | None = None
    ) -> np.ndarray:
        ch     = channel if channel is not None else self.cfg.channel
        bits   = self.cfg.lsb_bits
        mask   = 0xFF ^ ((1 << bits) - 1)   # máscara para limpiar los LSBs

        flat   = arr[:, :, ch].flatten().astype(np.int32)
        stream = self._bytes_to_bits(payload)

        capacity = len(flat) * bits
        if len(stream) > capacity:
            # Truncar payload si excede la capacidad del canal
            stream = stream[:capacity]

        for i, bit_chunk in enumerate(self._chunked(stream, bits)):
            if i >= len(flat):
                break
            # Limpiar LSBs y escribir los bits del payload
            val = int(flat[i]) & mask
            for j, b in enumerate(bit_chunk):
                val |= (b << j)
            flat[i] = np.clip(val, 0, 255)

        result = arr.copy()
        result[:, :, ch] = flat.reshape(arr.shape[:2]).astype(np.uint8)
        return result

    # ------------------------------------------------------------------
    # Método 2: DCT
    # ------------------------------------------------------------------
    def _embed_dct(
        self, arr: np.ndarray, payload: bytes, channel: int | None = None
    ) -> np.ndarray:
        ch      = channel if channel is not None else self.cfg.channel
        bs      = self.cfg.dct_block
        strength= self.cfg.dct_strength

        plane   = arr[:, :, ch].astype(np.float64)
        h, w    = plane.shape
        bits    = self._bytes_to_bits(payload)
        bit_idx = 0

        result  = plane.copy()

        for row in range(0, h - bs + 1, bs):
            for col in range(0, w - bs + 1, bs):
                if bit_idx >= len(bits):
                    break
                block = plane[row:row+bs, col:col+bs]
                C     = dctn(block, norm="ortho")

                # Modificar el coeficiente [4][3] (frecuencia media; balance invisibilidad/robustez)
                bit   = bits[bit_idx]
                coeff = C[4][3]
                # Cuantización binaria del coeficiente
                if bit == 1:
                    C[4][3] = abs(coeff) + strength if coeff >= 0 else -(abs(coeff) + strength)
                else:
                    C[4][3] = -(abs(coeff) + strength) if coeff >= 0 else abs(coeff) + strength

                block_back = idctn(C, norm="ortho")
                result[row:row+bs, col:col+bs] = block_back
                bit_idx += 1

        result = np.clip(result, 0, 255).astype(np.uint8)
        out    = arr.copy()
        out[:, :, ch] = result
        return out

    # ------------------------------------------------------------------
    # Método 3: Adaptive noise
    # ------------------------------------------------------------------
    def _embed_adaptive(
        self, arr: np.ndarray, payload: bytes, channel: int | None = None
    ) -> np.ndarray:
        """
        Añade ruido gaussiano firmado por el payload.
        Zonas de alta varianza natural de la imagen absorben más energía
        → artefactos visualmente mínimos donde hay textura, nulos en zonas lisas.
        """
        ch    = channel if channel is not None else self.cfg.channel
        sigma = self.cfg.noise_sigma

        plane = arr[:, :, ch].astype(np.float64)
        h, w  = plane.shape

        # Mapa de varianza local (ventana 5x5)
        variance_map = self._local_variance(plane, window=5)
        variance_map = variance_map / (variance_map.max() + 1e-9)  # normalizar 0-1

        bits    = self._bytes_to_bits(payload)
        rng     = np.random.default_rng(seed=self._payload_seed(payload))
        noise   = rng.normal(0, sigma, (h, w))

        # Firmar el ruido con los bits del payload (repetir si es necesario)
        sign_pattern = np.array(
            [(1 if bits[i % len(bits)] == 1 else -1) for i in range(h * w)],
            dtype=np.float64
        ).reshape(h, w)

        # Escalar ruido por mapa de varianza (más ruido donde hay textura)
        modulated = noise * sign_pattern * variance_map * sigma

        result = np.clip(plane + modulated, 0, 255).astype(np.uint8)
        out    = arr.copy()
        out[:, :, ch] = result
        return out

    # ------------------------------------------------------------------
    # Helpers internos
    # ------------------------------------------------------------------
    @staticmethod
    def _to_array(image: Image.Image) -> np.ndarray:
        img = image.convert("RGBA")
        return np.array(img, dtype=np.uint8)

    @staticmethod
    def _from_array(arr: np.ndarray) -> Image.Image:
        return Image.fromarray(arr.astype(np.uint8), mode="RGBA")

    def _encode_payload(self, text: str) -> bytes:
        """
        Serializa el texto con header mágico + longitud + CRC32.
        Permite verificar extracción futura.
        """
        raw     = text.encode("utf-8")
        length  = struct.pack(">I", len(raw))
        crc     = struct.pack(">I", zlib.crc32(raw) & 0xFFFFFFFF)
        payload = self._HEADER_MAGIC + length + crc + raw
        return payload

    @staticmethod
    def _bytes_to_bits(data: bytes) -> list[int]:
        bits: list[int] = []
        for byte in data:
            for i in range(8):
                bits.append((byte >> i) & 1)
        return bits

    @staticmethod
    def _chunked(lst: list, n: int):
        for i in range(0, len(lst), n):
            chunk = lst[i:i+n]
            # Rellenar con 0 si el último chunk es incompleto
            while len(chunk) < n:
                chunk.append(0)
            yield chunk

    @staticmethod
    def _local_variance(plane: np.ndarray, window: int = 5) -> np.ndarray:
        from scipy.ndimage import uniform_filter
        mean  = uniform_filter(plane.astype(np.float64), size=window)
        mean2 = uniform_filter(plane.astype(np.float64) ** 2, size=window)
        return np.maximum(mean2 - mean ** 2, 0)

    @staticmethod
    def _payload_seed(payload: bytes) -> int:
        return zlib.crc32(payload) & 0xFFFFFFFF