ezu-paint 0.8.2

//! `mosaic` — `Raster -> Raster`. Quantize the input into uniform
//! square blocks. Each block is summarized to a single color via
//! `mode`:
//!
//! - `mode: "average"` (default): the mean color of the source pixels
//!   covered by the block. Smooth across colour transitions — the
//!   classic "mosaic" filter.
//! - `mode: "nearest"`: the source pixel at the block's center. No
//!   blending across boundaries — adjacent blocks step crisply between
//!   source colours.
//!
//! With `anchor: "world"` (default) the block grid is anchored to z=0
//! world pixels, so adjacent map tiles see the same block boundaries
//! and seams disappear. This requires upstream padding so blocks
//! straddling the tile edge can be summarized from the same source
//! pixels in both tiles. With `anchor: "tile"` each map tile starts
//! the grid at its own top-left and no extra padding is needed (but
//! tile seams may be visible).

use std::collections::HashMap;
use std::sync::Arc;

use ezu_graph::{
    schema_frag, take_input_ref, BuiltNode, Connection, CoordSpace, EvalCtx, EvalError, FactoryCtx,
    FactoryError, InReader, Node, NodeFactory, PaddingIn, PortKind, PortSpec, PortValue, RasterBuf,
};
use serde_json::Value;
use xxhash_rust::xxh3::Xxh3;

use crate::nodes::common::{read_optional_string, Anchor};

#[derive(Clone, Copy)]
enum Mode {
    Average,
    Nearest,
}

struct MosaicNode {
    block: PaddingIn,
    /// `block`'s bound rounded to whole pixels, for pad propagation.
    block_bound: u32,
    anchor: Anchor,
    mode: Mode,
    ports: Vec<PortSpec>,
    param_refs: Vec<String>,
}

impl Node for MosaicNode {
    fn op_name(&self) -> &'static str {
        "mosaic"
    }
    fn inputs(&self) -> &[PortSpec] {
        &self.ports
    }
    fn output(&self, _input_kinds: &[Option<PortKind>]) -> PortKind {
        PortKind::Raster
    }
    fn coord_space(&self) -> CoordSpace {
        match self.anchor {
            Anchor::World => CoordSpace::World,
            Anchor::Tile => CoordSpace::Tile,
        }
    }
    fn required_pad(&self, downstream: u32) -> u32 {
        // World-anchored blocks may straddle the tile edge. `average`
        // mode needs the full straddling block visible (so both tiles
        // see the same samples and produce the same mean). `nearest`
        // mode only needs the block's center pixel reachable, which is
        // at most `block / 2` pixels outside the tile.
        match self.anchor {
            Anchor::World => match self.mode {
                Mode::Average => downstream.saturating_add(self.block_bound),
                Mode::Nearest => downstream.saturating_add(self.block_bound.div_ceil(2)),
            },
            Anchor::Tile => downstream,
        }
    }
    fn eval(
        &self,
        ctx: &EvalCtx<'_>,
        inputs: &[Option<PortValue>],
    ) -> Result<PortValue, EvalError> {
        let src = inputs[0]
            .as_ref()
            .and_then(PortValue::as_raster)
            .ok_or_else(|| EvalError::MissingInput("input".into()))?;
        let block = (self.block.get(ctx, inputs)?.round() as i64).max(1);
        if block == 1 {
            return Ok(PortValue::Raster(src.clone()));
        }
        let w = src.width as i64;
        let h = src.height as i64;
        let pad = ctx.canvas.pad as i64;
        let tile_size = ctx.canvas.tile_w as i64;

        // World coordinate of source pixel (0, 0).
        let (origin_x, origin_y) = match self.anchor {
            Anchor::World => (
                ctx.tile.x as i64 * tile_size - pad,
                ctx.tile.y as i64 * tile_size - pad,
            ),
            Anchor::Tile => (0, 0),
        };

        // Cache per-block averages keyed by world block index. Each
        // block is averaged once and reused for every output pixel
        // inside it.
        let mut cache: HashMap<(i64, i64), [u8; 4]> = HashMap::new();
        let src_px = &src.pixels;
        let mut out = RasterBuf::new(src.width, src.height);
        let dst = &mut out.pixels;

        for y in 0..h {
            let world_y = origin_y + y;
            let by = world_y.div_euclid(block);
            for x in 0..w {
                let world_x = origin_x + x;
                let bx = world_x.div_euclid(block);
                let color = *cache.entry((bx, by)).or_insert_with(|| match self.mode {
                    Mode::Average => {
                        // World-space block extent.
                        let wx0 = bx * block;
                        let wy0 = by * block;
                        let wx1 = wx0 + block;
                        let wy1 = wy0 + block;
                        // Clip to the available source pixels.
                        let sx0 = (wx0 - origin_x).max(0);
                        let sy0 = (wy0 - origin_y).max(0);
                        let sx1 = (wx1 - origin_x).min(w);
                        let sy1 = (wy1 - origin_y).min(h);
                        if sx1 <= sx0 || sy1 <= sy0 {
                            return [0; 4];
                        }
                        let mut sum = [0u64; 4];
                        let mut count = 0u64;
                        for yy in sy0..sy1 {
                            let row = (yy * w) as usize * 4;
                            for xx in sx0..sx1 {
                                let i = row + (xx as usize) * 4;
                                sum[0] += src_px[i] as u64;
                                sum[1] += src_px[i + 1] as u64;
                                sum[2] += src_px[i + 2] as u64;
                                sum[3] += src_px[i + 3] as u64;
                                count += 1;
                            }
                        }
                        [
                            (sum[0] / count) as u8,
                            (sum[1] / count) as u8,
                            (sum[2] / count) as u8,
                            (sum[3] / count) as u8,
                        ]
                    }
                    Mode::Nearest => {
                        // Sample the source pixel at the block's center
                        // (in world coords). Using a fixed offset means
                        // adjacent tiles pick the same representative
                        // for any straddling block.
                        let half = block / 2;
                        let cx = bx * block + half - origin_x;
                        let cy = by * block + half - origin_y;
                        if cx < 0 || cy < 0 || cx >= w || cy >= h {
                            return [0; 4];
                        }
                        let i = ((cy * w + cx) as usize) * 4;
                        [src_px[i], src_px[i + 1], src_px[i + 2], src_px[i + 3]]
                    }
                });
                let i = ((y * w + x) as usize) * 4;
                dst[i] = color[0];
                dst[i + 1] = color[1];
                dst[i + 2] = color[2];
                dst[i + 3] = color[3];
            }
        }
        Ok(PortValue::Raster(Arc::new(out)))
    }
    fn param_hash(&self, h: &mut Xxh3) {
        h.update(b"mosaic");
        self.block.param_hash(h);
        match self.anchor {
            Anchor::World => h.update(b"w"),
            Anchor::Tile => h.update(b"t"),
        }
        match self.mode {
            Mode::Average => h.update(b"a"),
            Mode::Nearest => h.update(b"n"),
        }
    }
    fn param_refs(&self) -> Vec<String> {
        self.param_refs.clone()
    }
}

pub(super) struct MosaicFactory;
impl NodeFactory for MosaicFactory {
    fn op_name(&self) -> &'static str {
        "mosaic"
    }
    fn build(
        &self,
        fields: &serde_json::Map<String, Value>,
        ctx: &FactoryCtx<'_>,
    ) -> Result<BuiltNode, FactoryError> {
        let input = take_input_ref(fields, "input")?;
        let mut r = InReader::new(fields, ctx, 1);
        let block = PaddingIn::read(&mut r, fields, "block")?;
        let parts = r.finish();
        let block_bound_raw = block.bound();
        if !(block_bound_raw.is_finite() && block_bound_raw >= 1.0) {
            return Err(FactoryError::BadField {
                field: "block".into(),
                msg: "expected integer >= 1".into(),
            });
        }
        let block_bound = block_bound_raw.round() as u32;
        let anchor = match read_optional_string(fields, "anchor")?.as_deref() {
            None | Some("world") => Anchor::World,
            Some("tile") => Anchor::Tile,
            Some(other) => {
                return Err(FactoryError::BadField {
                    field: "anchor".into(),
                    msg: format!("expected `world` or `tile`, got `{other}`"),
                });
            }
        };
        let mode = match read_optional_string(fields, "mode")?.as_deref() {
            None | Some("average") => Mode::Average,
            Some("nearest") => Mode::Nearest,
            Some(other) => {
                return Err(FactoryError::BadField {
                    field: "mode".into(),
                    msg: format!("expected `average` or `nearest`, got `{other}`"),
                });
            }
        };
        let mut ports = vec![PortSpec {
            name: "input",
            accepts: &[PortKind::Raster],
            optional: false,
        }];
        ports.extend(parts.ports);
        let mut connections = vec![Connection {
            port: "input".into(),
            src: input,
        }];
        connections.extend(parts.connections);

        Ok(BuiltNode {
            node: Box::new(MosaicNode {
                block,
                block_bound,
                anchor,
                mode,
                ports,
                param_refs: parts.param_refs,
            }),
            connections,
        })
    }
    fn schema(&self) -> Value {
        serde_json::json!({
            "description": "Quantize a raster into uniform square blocks. `mode` selects how each block is summarized: `average` (default) takes the mean of covered pixels (smooth, classic mosaic); `nearest` takes the block's centre pixel (crisp, no inter-colour blending). World-anchored by default so adjacent map tiles share the same block grid.",
            "properties": {
                "input": schema_frag::node_ref(),
                "block": schema_frag::in_number(serde_json::json!({
                    "type": "integer", "minimum": 1,
                    "description": "Block edge length in canvas pixels. With `anchor: world` the block size sets the canvas padding, which is fixed before anything renders — so this needs an upper bound the build can see: a literal, a `$param` with `max`, or `block-max` beside an `@node` port."
                })),
                "block-max": { "type": "number", "minimum": 0.0, "description": "Upper bound on `block` for padding, required when `block` is an `@node` port. Values above it are clamped." },
                "anchor": { "type": "string", "enum": ["world", "tile"], "default": "world",
                            "description": "`world` (default) makes the block grid seamless across map tiles by growing the upstream pad. `tile` restarts the grid at every map tile's top-left and requires no extra padding." },
                "mode": { "type": "string", "enum": ["average", "nearest"], "default": "average",
                          "description": "`average` (default) blends covered pixels into a mean colour. `nearest` samples the block's centre pixel verbatim — produces hard block edges without inter-colour averaging." },
            },
            "required": ["input", "block"],
        })
    }
}

ezu_graph::submit_node!(MosaicFactory);