pbrt/src/image/ops.rs

// use rayon::prelude::*;
use crate::core::pbrt::Float;
use crate::geometry::{Bounds2i, Point2i};
use crate::image::pixel::PixelStorage;
use crate::image::{Image, PixelData, PixelFormat, WrapMode, WrapMode2D};
use crate::utils::math::windowed_sinc;
use rayon::prelude::*;
use std::sync::{Arc, Mutex};

#[derive(Debug, Clone, Copy)]
pub struct ResampleWeight {
    pub first_pixel: i32,
    pub weight: [Float; 4],
}

impl Image {
    pub fn flip_y(&mut self) {
        let res = self.resolution;
        let nc = self.n_channels();

        match &mut self.pixels {
            PixelData::U8(d) => flip_y_kernel(d, res, nc),
            PixelData::F16(d) => flip_y_kernel(d, res, nc),
            PixelData::F32(d) => flip_y_kernel(d, res, nc),
        }
    }

    pub fn crop(&self, bounds: Bounds2i) -> Image {
        let new_res = Point2i::new(
            bounds.p_max.x() - bounds.p_min.x(),
            bounds.p_max.y() - bounds.p_min.y(),
        );

        let mut new_image = Image::from_vector(
            self.format,
            new_res,
            self.channel_names.clone(),
            self.encoding,
        );

        match (&self.pixels, &mut new_image.pixels) {
            (PixelData::U8(src), PixelData::U8(dst)) => {
                crop_kernel(src, dst, self.resolution, bounds, self.n_channels())
            }
            (PixelData::F16(src), PixelData::F16(dst)) => {
                crop_kernel(src, dst, self.resolution, bounds, self.n_channels())
            }
            (PixelData::F32(src), PixelData::F32(dst)) => {
                crop_kernel(src, dst, self.resolution, bounds, self.n_channels())
            }
            _ => panic!("Format mismatch in crop"),
        }

        new_image
    }

    pub fn copy_rect_out(&self, extent: Bounds2i, buf: &mut [Float], wrap: WrapMode2D) {
        match &self.pixels {
            PixelData::U8(d) => copy_rect_out_kernel(d, self, extent, buf, wrap),
            PixelData::F16(d) => copy_rect_out_kernel(d, self, extent, buf, wrap),
            PixelData::F32(d) => copy_rect_out_kernel(d, self, extent, buf, wrap),
        }
    }

    pub fn copy_rect_in(&mut self, extent: Bounds2i, buf: &[Float]) {
        let resolution = self.resolution;
        let n_channels = self.n_channels();
        let encoding = self.encoding;

        match &mut self.pixels {
            PixelData::U8(d) => {
                copy_rect_in_kernel(d, resolution, n_channels, encoding, extent, buf)
            }
            PixelData::F16(d) => {
                copy_rect_in_kernel(d, resolution, n_channels, encoding, extent, buf)
            }
            PixelData::F32(d) => {
                copy_rect_in_kernel(d, resolution, n_channels, encoding, extent, buf)
            }
        }
    }

    pub fn float_resize_up(&self, new_res: Point2i, wrap_mode: WrapMode2D) -> Image {
        assert!(new_res.x() >= self.resolution.x() && new_res.y() >= self.resolution.y());
        assert!(
            matches!(self.format, PixelFormat::F32),
            "ResizeUp requires Float format"
        );

        let resampled_image = Arc::new(Mutex::new(Image::from_vector(
            PixelFormat::F32, // Force float output
            new_res,
            self.channel_names.clone(),
            self.encoding,
        )));

        let x_weights = resample_weights(self.resolution.x() as usize, new_res.x() as usize);
        let y_weights = resample_weights(self.resolution.y() as usize, new_res.y() as usize);
        let n_channels = self.n_channels();

        let tile_size = 16;
        let tiles = generate_tiles(new_res, tile_size);

        tiles.par_iter().for_each(|out_extent| {
            let x_start = x_weights[out_extent.p_min.x() as usize].first_pixel;
            let x_end = x_weights[(out_extent.p_max.x() - 1) as usize].first_pixel + 4;
            let y_start = y_weights[out_extent.p_min.y() as usize].first_pixel;
            let y_end = y_weights[(out_extent.p_max.y() - 1) as usize].first_pixel + 4;

            let in_extent =
                Bounds2i::from_points(Point2i::new(x_start, y_start), Point2i::new(x_end, y_end));

            let mut in_buf = vec![0.0; in_extent.area() as usize * n_channels];
            self.copy_rect_out(in_extent, &mut in_buf, wrap_mode);

            let out_buf = compute_resize_tile(
                &in_buf,
                in_extent,
                *out_extent,
                n_channels,
                &x_weights,
                &y_weights,
            );

            let mut guard = resampled_image.lock().unwrap();
            guard.copy_rect_in(*out_extent, &out_buf);
        });

        Arc::try_unwrap(resampled_image)
            .unwrap()
            .into_inner()
            .unwrap()
    }

    pub fn generate_pyramid(base: Image, _wrap: WrapMode) -> Vec<Image> {
        let mut levels = vec![base];
        let internal_wrap = WrapMode2D {
            uv: [WrapMode::Clamp; 2],
        };

        loop {
            let prev = levels.last().unwrap();
            let old = prev.resolution;
            if old.x() == 1 && old.y() == 1 {
                break;
            }

            let new_res = Point2i::new((old.x() / 2).max(1), (old.y() / 2).max(1));
            let mut next = Image::from_vector(
                prev.format,
                new_res,
                prev.channel_names.clone(),
                prev.encoding,
            );

            match &mut next.pixels {
                PixelData::U8(d) => downsample_kernel(d, new_res, prev, internal_wrap),
                PixelData::F16(d) => downsample_kernel(d, new_res, prev, internal_wrap),
                PixelData::F32(d) => downsample_kernel(d, new_res, prev, internal_wrap),
            }
            levels.push(next);
        }
        levels
    }
}

fn flip_y_kernel<T: PixelStorage>(pixels: &mut [T], res: Point2i, channels: usize) {
    let w = res.x() as usize;
    let h = res.y() as usize;
    let stride = w * channels;

    for y in 0..(h / 2) {
        let bot = h - 1 - y;
        for i in 0..stride {
            pixels.swap(y * stride + i, bot * stride + i);
        }
    }
}

fn crop_kernel<T: PixelStorage>(
    src: &[T],
    dst: &mut [T],
    src_res: Point2i,
    bounds: Bounds2i,
    channels: usize,
) {
    let dst_w = (bounds.p_max.x() - bounds.p_min.x()) as usize;
    // let dst_h = (bounds.p_max.y() - bounds.p_min.y()) as usize;

    dst.par_chunks_mut(dst_w * channels)
        .enumerate()
        .for_each(|(dy, dst_row)| {
            let sy = bounds.p_min.y() as usize + dy;
            let sx_start = bounds.p_min.x() as usize;

            let src_offset = (sy * src_res.x() as usize + sx_start) * channels;
            let count = dst_w * channels;

            dst_row.copy_from_slice(&src[src_offset..src_offset + count]);
        });
}

fn copy_rect_out_kernel<T: PixelStorage>(
    src: &[T],
    image: &Image,
    extent: Bounds2i,
    buf: &mut [Float],
    wrap: WrapMode2D,
) {
    let w = (extent.p_max.x() - extent.p_min.x()) as usize;
    let channels = image.n_channels();
    let enc = image.encoding;
    let res = image.resolution;

    buf.par_chunks_mut(w * channels)
        .enumerate()
        .for_each(|(y_rel, row_buf)| {
            let y = extent.p_min.y() + y_rel as i32;
            for x_rel in 0..w {
                let x = extent.p_min.x() + x_rel as i32;

                // This allows us to use 'src' directly (Fast Path).
                if x >= 0 && x < res.x() && y >= 0 && y < res.y() {
                    let offset = (y as usize * res.x() as usize + x as usize) * channels;

                    for c in 0..channels {
                        row_buf[x_rel * channels + c] = T::to_linear(src[offset + c], enc);
                    }
                } else {
                    // Slow path: Out of bounds, requires Wrap Mode logic.
                    // We fall back to get_channel which handles the wrapping math.
                    let p = Point2i::new(x, y);
                    for c in 0..channels {
                        row_buf[x_rel * channels + c] = image.get_channel_with_wrap(p, c, wrap);
                    }
                }
            }
        });
}

fn copy_rect_in_kernel<T: PixelStorage>(
    dst: &mut [T],
    res: Point2i,
    channels: usize,
    enc: crate::utils::color::ColorEncoding,
    extent: Bounds2i,
    buf: &[Float],
) {
    let w = (extent.p_max.x() - extent.p_min.x()) as usize;
    let res_x = res.x() as usize;

    let rows = buf.chunks(w * channels);
    for (y_rel, row) in rows.enumerate() {
        let y = extent.p_min.y() + y_rel as i32;
        if y < 0 || y >= res.y() {
            continue;
        }

        let dst_row_start = (y as usize * res_x) * channels;

        for (x_rel, &val) in row.iter().enumerate() {
            let c = x_rel % channels;
            let x_pixel = x_rel / channels;
            let x = extent.p_min.x() + x_pixel as i32;

            if x >= 0 && x < res.x() {
                let idx = dst_row_start + (x as usize * channels) + c;
                dst[idx] = T::from_linear(val, enc);
            }
        }
    }
}

fn downsample_kernel<T: PixelStorage>(
    dst: &mut [T],
    dst_res: Point2i,
    prev: &Image,
    wrap: WrapMode2D,
) {
    let w = dst_res.x() as usize;
    let channels = prev.n_channels();
    let enc = prev.encoding;
    let old_res = prev.resolution;

    dst.par_chunks_mut(w * channels)
        .enumerate()
        .for_each(|(y, row)| {
            let src_y = y * 2;
            for x in 0..w {
                let src_x = x * 2;
                for c in 0..channels {
                    let mut sum = 0.0;
                    let mut count = 0.0;

                    for dy in 0..2 {
                        for dx in 0..2 {
                            let sx = src_x as i32 + dx;
                            let sy = src_y as i32 + dy;
                            if sx < old_res.x() && sy < old_res.y() {
                                sum += prev.get_channel_with_wrap(Point2i::new(sx, sy), c, wrap);
                                count += 1.0;
                            }
                        }
                    }

                    let avg = if count > 0.0 { sum / count } else { 0.0 };
                    row[x * channels + c] = T::from_linear(avg, enc);
                }
            }
        });
}

fn resample_weights(old_res: usize, new_res: usize) -> Vec<ResampleWeight> {
    let filter_radius = 2.0;
    let tau = 2.0;
    (0..new_res)
        .map(|i| {
            let center = (i as Float + 0.5) * old_res as Float / new_res as Float;
            let first_pixel = ((center - filter_radius) + 0.5).floor() as i32;
            let mut weights = [0.0; 4];
            let mut sum = 0.0;
            for j in 0..4 {
                let pos = (first_pixel + j) as Float + 0.5;
                weights[j as usize] = windowed_sinc(pos - center, filter_radius, tau);
                sum += weights[j as usize];
            }
            let inv_sum = 1.0 / sum;
            for w in &mut weights {
                *w *= inv_sum;
            }
            ResampleWeight {
                first_pixel,
                weight: weights,
            }
        })
        .collect()
}

fn generate_tiles(res: Point2i, tile_size: i32) -> Vec<Bounds2i> {
    let nx = (res.x() + tile_size - 1) / tile_size;
    let ny = (res.y() + tile_size - 1) / tile_size;
    let mut tiles = Vec::new();
    for y in 0..ny {
        for x in 0..nx {
            let p_min = Point2i::new(x * tile_size, y * tile_size);
            let p_max = Point2i::new(
                (x * tile_size + tile_size).min(res.x()),
                (y * tile_size + tile_size).min(res.y()),
            );
            tiles.push(Bounds2i::from_points(p_min, p_max));
        }
    }
    tiles
}

fn compute_resize_tile(
    in_buf: &[Float],
    in_extent: Bounds2i,
    out_extent: Bounds2i,
    n_channels: usize,
    x_w: &[ResampleWeight],
    y_w: &[ResampleWeight],
) -> Vec<Float> {
    let nx_out = (out_extent.p_max.x() - out_extent.p_min.x()) as usize;
    let ny_out = (out_extent.p_max.y() - out_extent.p_min.y()) as usize;
    let nx_in = (in_extent.p_max.x() - in_extent.p_min.x()) as usize;
    let ny_in = (in_extent.p_max.y() - in_extent.p_min.y()) as usize;

    let mut x_buf = vec![0.0; n_channels * ny_in * nx_out];

    // Resize X
    for y in 0..ny_in {
        for x in 0..nx_out {
            let x_global = out_extent.p_min.x() + x as i32;
            let w = &x_w[x_global as usize];
            let x_in_start = (w.first_pixel - in_extent.p_min.x()) as usize;

            let in_idx_base = (y * nx_in + x_in_start) * n_channels;
            let out_idx_base = (y * nx_out + x) * n_channels;

            for c in 0..n_channels {
                let mut val = 0.0;
                for k in 0..4 {
                    val += w.weight[k] * in_buf[in_idx_base + k * n_channels + c];
                }
                x_buf[out_idx_base + c] = val;
            }
        }
    }

    let mut out_buf = vec![0.0; n_channels * nx_out * ny_out];

    // Resize Y
    for x in 0..nx_out {
        for y in 0..ny_out {
            let y_global = out_extent.p_min.y() + y as i32;
            let w = &y_w[y_global as usize];
            let y_in_start = (w.first_pixel - in_extent.p_min.y()) as usize;

            let in_idx_base = (y_in_start * nx_out + x) * n_channels;
            let out_idx_base = (y * nx_out + x) * n_channels;
            let stride = nx_out * n_channels;

            for c in 0..n_channels {
                let mut val = 0.0;
                for k in 0..4 {
                    val += w.weight[k] * x_buf[in_idx_base + k * stride + c];
                }
                out_buf[out_idx_base + c] = val.max(0.0);
            }
        }
    }
    out_buf
}