#!/usr/bin/env catnip
# Dithering multi-algorithmes (threshold, Bayer ordonné, Floyd-Steinberg)
#
# Le dithering échange de la résolution spatiale contre de la résolution tonale :
# on ajoute un bruit CONTRÔLÉ avant de quantifier, si bien que l'erreur de
# quantification s'annule en moyenne locale au lieu de former des bandes plates
# (banding) à faible profondeur de couleur. Réf : Wikipedia « Dither ».
#
# Trois stratégies pour réduire une image en niveaux de gris à `levels` paliers :
#   - threshold : quantification au palier le plus proche, sans bruit (banding) ;
#   - Bayer     : seuil ordonné périodique 4x4, bruit déterministe et structuré ;
#   - Floyd-Steinberg : diffusion d'erreur, chaque pixel reverse son résidu de
#     quantification sur ses voisins non encore traités (séquentiel, raster).
#
# pyvips charge, réduit et assemble la planche ; numpy porte la quantification.
#
# DEPS: pyvips[binary], numpy

pyvips = import('pyvips')
numpy = import('numpy')
import('pathlib', 'Path')

script_dir = Path(META.file).parent
input_path = script_dir / 'data/catnip_test.jpg'
output_dir = script_dir / 'output'
output_dir.mkdir(exist_ok=True)

# Types somme : ensemble fermé d'algorithmes, exhaustivité vérifiée par le linter
# dans les match. label() donne la forme d'affichage.
union Algo {
    threshold; ordered; floyd

    label(self): str => {
        match self {
            Algo.threshold => { "threshold" }
            Algo.ordered   => { "bayer-ordered" }
            Algo.floyd     => { "floyd-steinberg" }
        }
    }
}

# levels=2 : palette bilevel (noir/blanc pur), le cas où le dithering est le plus
# visible. scale/step convertissent une valeur 0..255 vers un index de palette et
# inversement : idx = round(v * scale), v = idx * step.
struct DitherConfig {
    levels: int; base_width: int; zoom: int

    scale(self): float => { (self.levels - 1) / 255.0 }
    step(self): float => { 255.0 / (self.levels - 1) }
    max_level(self): int => { self.levels - 1 }
}

cfg = DitherConfig(2, 280, 2)

# Quantification vectorisée : chaque pixel vers le palier le plus proche. Partagée
# par threshold (bruit nul) et Bayer (bruit ajouté en amont dans l'argument).
quantize_array = (arr, config: DitherConfig) => {
    idx = numpy.clip(numpy.round(numpy.clip(arr, 0, 255) * config.scale()), 0, config.max_level())
    idx * config.step()
}

# Floyd-Steinberg : diffusion d'erreur en balayage raster. Le résidu old - newv
# est réparti sur 4 voisins (7/16 droite, 3/16 bas-gauche, 5/16 bas, 1/16 bas-droite),
# tous encore à traiter, donc le passage est intrinsèquement séquentiel.
quantize_floyd = (arr, config: DitherConfig) => {
    work = numpy.array(arr)
    h = work.shape[0]
    w = work.shape[1]
    scale = config.scale()
    step = config.step()
    y = 0
    while y < h {
        row = work[y]
        has_next = y + 1 < h
        nextrow = if has_next { work[y + 1] } else { row }
        x = 0
        while x < w {
            old = row[x]
            clampv = if old < 0.0 { 0.0 } else { if old > 255.0 { 255.0 } else { old } }
            newv = int(clampv * scale + 0.5) * step
            err = old - newv
            row[x] = newv
            if x + 1 < w { row[x + 1] = row[x + 1] + err * 0.4375 }
            if has_next {
                if x > 0 { nextrow[x - 1] = nextrow[x - 1] + err * 0.1875 }
                nextrow[x] = nextrow[x] + err * 0.3125
                if x + 1 < w { nextrow[x + 1] = nextrow[x + 1] + err * 0.0625 }
            }
            x = x + 1
        }
        y = y + 1
    }
    work
}

print("⇒ Chargement et passage en niveaux de gris")
source = pyvips.Image.new_from_file(str(input_path)).colourspace('b-w')
small = source.thumbnail_image(cfg.base_width)
gray = numpy.asarray(small.numpy(), dtype='float64')
print(f"  {small.width}×{small.height}, palette à {cfg.levels} paliers")

# Matrice de Bayer 4x4 (indices 0..15) : seuils ordonnés répartissant l'erreur
# spatialement de façon périodique. Normalisée dans [0,1), recentrée sur 0 puis
# pavée aux dimensions de l'image via numpy.take (pas de sous-indexation 2D).
bayer_norm = (numpy.array(list(
        list(0, 8, 2, 10),
        list(12, 4, 14, 6),
        list(3, 11, 1, 9),
        list(15, 7, 13, 5),
    )) +
    0.5) /
    16.0
h = gray.shape[0]
w = gray.shape[1]
big = numpy.tile(bayer_norm, tuple(int((h + 3) / 4), int((w + 3) / 4)))
bayer_map = numpy.take(numpy.take(big, numpy.arange(h), axis=0), numpy.arange(w), axis=1)

# Un bruit d'amplitude un palier (step) est ajouté avant quantification : c'est ce
# décalage ordonné qui casse le banding.
render = (algo: Algo, image, config: DitherConfig) => {
    match algo {
        Algo.threshold => { quantize_array(image, config) }
        Algo.ordered   => { quantize_array(image + config.step() * (bayer_map - 0.5), config) }
        Algo.floyd     => { quantize_floyd(image, config) }
    }
}

# Les trois algorithmes sont diffusés sur l'image : une seule expression, la
# sélection reste dans le match de render.
print()
print("⇒ Application des trois algorithmes")
algos = list(Algo.threshold, Algo.ordered, Algo.floyd)
rendered = algos.[(a) => { render(a, gray, cfg) }]

# Contrôle numérique : le dithering conserve la luminance MOYENNE (l'erreur se
# compense localement), là où un seuil dur la déforme selon la répartition des
# pixels clairs/sombres. On mesure l'écart de moyenne à la source.
print()
print("⇒ Conservation de la luminance moyenne (écart à la source)")
orig_mean = float(numpy.mean(gray))
print(f"  source : moyenne = {round(orig_mean, 2)}")
i = 0
while i < len(algos) {
    dev = float(numpy.abs(numpy.mean(rendered[i]) - orig_mean))
    print(f"  {algos[i].label():>16} : écart = {round(dev, 3)}")
    i = i + 1
}

# Planche comparative : source à gauche, puis les trois rendus. Zoom au plus
# proche voisin pour rendre la trame de points visible sans lissage.
print()
print("⇒ Assemblage de la planche comparative")
panels = list(small.resize(cfg.zoom, kernel='nearest'))
i = 0
while i < len(rendered) {
    panel = pyvips.Image.new_from_array(numpy.asarray(rendered[i], dtype='uint8'))
    panels.append(panel.resize(cfg.zoom, kernel='nearest'))
    i = i + 1
}
board = pyvips.Image.arrayjoin(panels, across=len(panels), shim=12, background=200)

output_path = output_dir / 'pyvips_dithering.png'
board.write_to_file(str(output_path))
print(f"  {len(panels)} panneaux (source + 3 rendus) → {output_path.name}")
print()
print(f"⇒ Sortie disponible dans {output_dir}")

# Aperçu navigateur : sert l'image une fois puis rend la main.
# --no-browser garde un chemin headless (le PNG reste écrit ci-dessus).
if '--no-browser' not in import('sys').argv {
    http = import('http')
    b64 = import('base64').b64encode(output_path.read_bytes()).decode('ascii')
    http.serve(f'<!doctype html><meta charset="utf-8"><title>pyvips dithering</title><body style="margin:0;background:#0d1117"><img style="max-width:100%;display:block;margin:0 auto" src="data:image/png;base64,{b64}">',
        0, 'text/html; charset=utf-8', True)
}