wonfy-tools 0.2.0

use std::collections::{HashMap, VecDeque};

use image::RgbaImage;

use crate::error::StitchError;
use crate::util::image::ImageCrop;

use super::gpu::{self, GpuContext, SearchParams};
use super::params::{CheckDirection, MatchMode, Order, OverlapScore, Position};

pub struct ImageStitcher {
    images: Vec<RgbaImage>,
    order: Order,
    direction: CheckDirection,
    window_size: usize,
    match_mode: MatchMode,
    crop: u32,
}

struct StitchOutput {
    image: RgbaImage,
    top_shift: Position,
    bottom_shift: Position,
    bottom_placement: Position,
}

impl ImageStitcher {
    pub fn new(
        images: Vec<RgbaImage>,
        order: Order,
        direction: CheckDirection,
        window_size: usize,
        match_mode: MatchMode,
        crop: u32,
    ) -> Self {
        Self {
            images,
            order,
            direction,
            window_size,
            match_mode,
            crop,
        }
    }

    pub fn validate(&self) -> Result<(), StitchError> {
        if self.images.len() < 2 {
            return Err(StitchError::NotEnoughImages(self.images.len()));
        }

        let window = u32::try_from(self.window_size).map_err(|_| StitchError::InvalidWindowSize)?;
        if window == 0 {
            return Err(StitchError::InvalidWindowSize);
        }

        let sideways = matches!(
            self.direction,
            CheckDirection::Sideways | CheckDirection::SidewaysRight | CheckDirection::SidewaysLeft
        );

        let mut width_bound: u64 = 0;

        for image in &self.images {
            let (width, height) = Self::match_dims(image, self.direction);

            let too_small = width == 0
                || height == 0
                || self.crop + window > height
                || (sideways && self.crop >= width);

            if too_small {
                return Err(StitchError::ImageTooSmall {
                    width: image.width(),
                    height: image.height(),
                    window_size: window,
                    crop: self.crop,
                });
            }

            width_bound = match sideways {
                // sideways composites widen with every stitch, bound by the total
                true => width_bound + width as u64,
                false => width_bound.max(width as u64),
            };
        }

        if 255 * window as u64 * 3 * width_bound >= u32::MAX as u64 {
            return Err(StitchError::ScoreOverflow {
                window_size: window,
                max_width: width_bound as u32,
            });
        }

        Ok(())
    }

    pub async fn stitch(self) -> Result<(RgbaImage, VecDeque<Position>), StitchError> {
        self.validate()?;

        let shared = gpu::shared_context().await?;
        #[cfg(not(target_arch = "wasm32"))]
        let context: &GpuContext = shared;
        #[cfg(target_arch = "wasm32")]
        let context: &GpuContext = &shared;

        match self.order {
            Order::Ordered => self.stitch_ordered(context).await,
            Order::Unordered => self.stitch_unordered(context).await,
        }
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn stitch_blocking(self) -> Result<(RgbaImage, VecDeque<Position>), StitchError> {
        pollster::block_on(self.stitch())
    }

    async fn stitch_ordered(
        self,
        context: &GpuContext,
    ) -> Result<(RgbaImage, VecDeque<Position>), StitchError> {
        let window = self.window_size as u32;
        let mut images = self.images.into_iter();
        let mut composite = images.next().expect("validated");
        let mut positions: VecDeque<Position> = VecDeque::new();
        let mut skip = 0u32;

        for image in images {
            let region = Self::search_pair(
                context,
                &composite,
                &image,
                self.direction,
                window,
                &self.match_mode,
                self.crop,
                skip,
            )
            .await?;

            let output = Self::stitch_pair(
                &composite,
                &image,
                &region.position,
                self.crop,
                self.direction,
            );

            for position in positions.iter_mut() {
                *position += &output.top_shift;
            }

            skip = match self.direction {
                CheckDirection::Horizontal => output.bottom_placement.x.max(0) as u32,
                _ => output.bottom_placement.y.max(0) as u32,
            };

            positions.push_back(output.bottom_placement);
            composite = output.image;
        }

        Ok((composite, positions))
    }

    async fn stitch_unordered(
        self,
        context: &GpuContext,
    ) -> Result<(RgbaImage, VecDeque<Position>), StitchError> {
        let window = self.window_size as u32;
        let mut next_id = 0u64;
        let mut id = || {
            next_id += 1;
            next_id
        };

        let mut pool: Vec<(u64, RgbaImage)> =
            self.images.into_iter().map(|image| (id(), image)).collect();
        let mut scores: HashMap<(u64, u64), (f64, OverlapScore)> = HashMap::new();
        let mut positions: Vec<(u64, Position)> = Vec::new();

        while pool.len() > 1 {
            for first in 0..pool.len() {
                for second in first + 1..pool.len() {
                    let (id1, image1) = &pool[first];
                    let (id2, image2) = &pool[second];

                    if scores.contains_key(&(*id1, *id2)) {
                        continue;
                    }

                    let forward = Self::search_pair(
                        context,
                        image1,
                        image2,
                        self.direction,
                        window,
                        &self.match_mode,
                        self.crop,
                        0,
                    )
                    .await?;

                    let mut backward = Self::search_pair(
                        context,
                        image2,
                        image1,
                        self.direction,
                        window,
                        &self.match_mode,
                        self.crop,
                        0,
                    )
                    .await?;
                    backward.flipped = true;

                    let best = match forward.score <= backward.score {
                        true => forward,
                        false => backward,
                    };

                    let (width1, _) = Self::match_dims(image1, self.direction);
                    let (width2, _) = Self::match_dims(image2, self.direction);
                    let normalized =
                        best.score as f64 / (window as f64 * width1.min(width2) as f64);

                    scores.insert((*id1, *id2), (normalized, best));
                }
            }

            let mut best_pair: Option<((u64, u64), f64, OverlapScore)> = None;

            for (first, (id1, _)) in pool.iter().enumerate() {
                for (id2, _) in &pool[first + 1..] {
                    let key = (*id1, *id2);
                    let (normalized, region) = &scores[&key];

                    let is_better = match &best_pair {
                        None => true,
                        Some((best_key, best_norm, _)) => {
                            *normalized < *best_norm
                                || (*normalized == *best_norm && key < *best_key)
                        }
                    };

                    if is_better {
                        best_pair = Some((key, *normalized, region.clone()));
                    }
                }
            }

            let ((id1, id2), _, region) = best_pair.expect("pool has at least one pair");

            let index1 = pool.iter().position(|(id, _)| *id == id1).unwrap();
            let index2 = pool.iter().position(|(id, _)| *id == id2).unwrap();
            let (later, earlier) = (index1.max(index2), index1.min(index2));
            let second = pool.remove(later);
            let first = pool.remove(earlier);

            let ((top_id, top), (bottom_id, bottom)) = match region.flipped {
                false => (first, second),
                true => (second, first),
            };

            let output =
                Self::stitch_pair(&top, &bottom, &region.position, self.crop, self.direction);

            let new_id = id();

            for (owner, position) in positions.iter_mut() {
                if *owner == top_id {
                    *position += &output.top_shift;
                } else if *owner == bottom_id {
                    *position += &output.bottom_shift;
                } else {
                    continue;
                }

                *owner = new_id;
            }

            positions.push((new_id, output.bottom_placement));
            scores.retain(|(key1, key2), _| ![id1, id2].contains(key1) && ![id1, id2].contains(key2));
            pool.push((new_id, output.image));
        }

        let (_, image) = pool.pop().expect("validated");
        Ok((image, positions.into_iter().map(|(_, p)| p).collect()))
    }

    /// (width, height) in match space: Horizontal matches along columns, so it swaps.
    fn match_dims(image: &RgbaImage, direction: CheckDirection) -> (u32, u32) {
        match direction {
            CheckDirection::Horizontal => (image.height(), image.width()),
            _ => (image.width(), image.height()),
        }
    }

    #[allow(clippy::too_many_arguments)]
    async fn search_pair(
        context: &GpuContext,
        part1: &RgbaImage,
        part2: &RgbaImage,
        direction: CheckDirection,
        window: u32,
        match_mode: &MatchMode,
        crop: u32,
        skip: u32,
    ) -> Result<OverlapScore, StitchError> {
        use CheckDirection as CD;

        let (offset_start, offset_end) = match direction {
            CD::Sideways | CD::SidewaysRight | CD::SidewaysLeft => {
                let width = part1.width().max(part2.width());
                let start = -((width - 1 - crop) as i32);
                let end = (width - crop) as i32;

                match direction {
                    CD::SidewaysLeft => (start, 0),
                    CD::SidewaysRight => (0, end),
                    _ => (start, end),
                }
            }
            _ => (0, 0),
        };

        let best = gpu::find_best_overlap(
            context,
            part1,
            part2,
            &SearchParams {
                transpose: direction == CD::Horizontal,
                edges: matches!(match_mode, MatchMode::Edges),
                window,
                crop,
                skip,
                offset_start,
                offset_end,
            },
        )
        .await?;

        let position = match direction {
            CD::Horizontal => Position {
                x: best.anchor as i32,
                y: 0,
            },
            _ => Position {
                x: best.offset,
                y: best.anchor as i32,
            },
        };

        Ok(OverlapScore {
            score: best.score as u64,
            position,
            flipped: false,
        })
    }

    /// `position` places `bottom`'s matched content relative to `top` in uncropped
    /// match coordinates; crops are applied to the seam sides before stacking.
    fn stitch_pair(
        top: &RgbaImage,
        bottom: &RgbaImage,
        position: &Position,
        crop: u32,
        direction: CheckDirection,
    ) -> StitchOutput {
        use CheckDirection::*;

        let top_crop = match direction {
            Vertical => ImageCrop {
                bottom: crop,
                ..Default::default()
            },
            Horizontal => ImageCrop {
                right: crop,
                ..Default::default()
            },
            Sideways | SidewaysRight | SidewaysLeft => {
                let mut top_crop = ImageCrop {
                    bottom: crop,
                    ..Default::default()
                };

                if position.x < 0 {
                    top_crop.right = crop;
                } else {
                    top_crop.left = crop;
                }

                top_crop
            }
        };
        let bottom_crop = top_crop.clone().reverse();

        // matching ran on uncropped columns, so side crops shift the anchor
        let offset = match direction {
            Sideways | SidewaysRight | SidewaysLeft => Position {
                x: match position.x >= 0 {
                    true => position.x - crop as i32,
                    false => position.x + crop as i32,
                },
                y: position.y,
            },
            _ => position.clone(),
        };

        let top_image = top_crop.crop_image(top);
        let bottom_image = bottom_crop.crop_image(bottom);

        let top_at = Position {
            x: (-offset.x).max(0),
            y: (-offset.y).max(0),
        };
        let bottom_at = Position {
            x: offset.x.max(0),
            y: offset.y.max(0),
        };

        let out_width = (top_at.x as u32 + top_image.width())
            .max(bottom_at.x as u32 + bottom_image.width());
        let out_height = (top_at.y as u32 + top_image.height())
            .max(bottom_at.y as u32 + bottom_image.height());

        let mut output = RgbaImage::new(out_width, out_height);
        Self::blit(&mut output, &top_image, top_at.x as u32, top_at.y as u32);
        Self::blit(
            &mut output,
            &bottom_image,
            bottom_at.x as u32,
            bottom_at.y as u32,
        );

        StitchOutput {
            image: output,
            top_shift: Position {
                x: top_at.x - top_crop.left as i32,
                y: top_at.y - top_crop.top as i32,
            },
            bottom_shift: Position {
                x: bottom_at.x - bottom_crop.left as i32,
                y: bottom_at.y - bottom_crop.top as i32,
            },
            bottom_placement: bottom_at,
        }
    }

    fn blit(destination: &mut RgbaImage, source: &RgbaImage, x: u32, y: u32) {
        let destination_width = destination.width() as usize;
        let source_width = source.width() as usize;
        let destination_raw: &mut [u8] = destination;
        let source_raw: &[u8] = source;

        for row in 0..source.height() as usize {
            let source_offset = row * source_width * 4;
            let destination_offset = ((y as usize + row) * destination_width + x as usize) * 4;

            destination_raw[destination_offset..destination_offset + source_width * 4]
                .copy_from_slice(&source_raw[source_offset..source_offset + source_width * 4]);
        }
    }
}