//! TRACES: FR-MRG-10
//! Where each frame gives way to the next.
//!
//! The first merges averaged every overlap: each frame weighted by its
//! distance from its own edge, so that across two hundred pixels one frame
//! faded into the other. That hides an exposure step and does not hide
//! anything that differs between the frames — parallax on a near slope, a
//! walker, a branch in the wind — which the average draws twice, half as
//! bright, a soft double edge at 1:1.
//!
//! A seam answers it the way every stitcher does: in an overlap, each output
//! pixel is taken from *one* frame, and the line where the choice changes is
//! put where the frames agree and the picture is smooth — through sky,
//! along a shadow, round the walker rather than through him — and away from
//! either frame's edge, where vignetting and the lens correction's fringe
//! live. The blend is then narrow and only across that line.
//!
//! # How
//!
//! At proxy resolution, on the output surface, which fits (panorama.md §5:
//! "it is a mask, not an image"):
//!
//! 1. Frames are laid down one at a time, each next to one already placed.
//!    The composite so far is a label per texel and the value its owner saw.
//! 2. Where a new frame overlaps the composite, a cost per texel: the
//!    difference between the two (after the gains), how much detail either
//!    has there, and how near either frame's edge it is — smoothed over a
//!    few texels, because "agree" means locally, not at one pixel.
//! 3. The cut is a path across the overlap, perpendicular to the line from
//!    the composite's frames to the new one, found by dynamic programming
//!    one row at a time: the per-column seam panorama.md §4 chose over a
//!    graph cut because it is the GPU-friendly shape. Texels on the new
//!    frame's side of the path become its own.
//!
//! What the merge reads is [`SeamMap::share`]: the fraction of a small
//! window about a point that is labelled with a frame, tent-weighted, which
//! is a narrow blend that follows the seam. `merge.wgsl` computes the same
//! thing on the GPU from the same labels.

use crate::bundle::Cameras;
use crate::image::Gray;
use crate::projection::{self, Projection};

/// No frame owns this texel.
pub const NONE: u8 = 255;

/// The most frames a map can label: one less than [`NONE`].
pub const MAX_FRAMES: usize = NONE as usize;

/// Which frame each texel of the output takes its pixels from.
#[derive(Debug, Clone, PartialEq)]
pub struct SeamMap {
    pub width: usize,
    pub height: usize,
    /// The projection scale the map was laid out at: the proxies' focal
    /// length. Output coordinates at any other scale are this times the
    /// ratio of the scales.
    pub scale: f64,
    /// Centred output coordinates, at `scale`, of texel (0, 0)'s top-left
    /// corner.
    pub origin: (f64, f64),
    /// Output units per texel, at `scale`.
    pub px: f64,
    /// Row-major, one per texel: the frame's index, or [`NONE`].
    pub labels: Vec<u8>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SeamOptions {
    /// The widest the map is laid out, in texels. Wider than the proxies'
    /// own resolution buys nothing.
    pub max_width: usize,
    /// How much detail costs against disagreement: a seam through texture
    /// shows even where the frames agree, because the blend across it
    /// softens it.
    pub detail: f32,
    /// How much a frame's edge costs, and how far in from it the cost
    /// reaches, in proxy pixels. Frame edges are where vignetting is
    /// darkest and the lens correction ran out of sensor.
    pub edge: f32,
    pub edge_margin: f32,
    /// The radius, in texels, a texel's cost looks about it for the worst
    /// of its neighbours: at least the radius the merge blends across.
    pub smoothing: usize,
}

impl Default for SeamOptions {
    fn default() -> Self {
        SeamOptions {
            max_width: 2048,
            detail: 0.5,
            edge: 0.5,
            edge_margin: 24.0,
            smoothing: 4,
        }
    }
}

/// The most texels a blend reaches either side of a seam. The merge's
/// shader loads the square of twice this per pixel per frame near a seam.
pub const MAX_BLEND_RADIUS: f64 = 4.0;

/// Cost of a texel outside the overlap: high enough that the path keeps to
/// the overlap wherever there is one, finite so that a row with a gap in it
/// still has an answer.
const OUTSIDE: f32 = 1.0e3;

impl SeamMap {
    /// The map's origin and texel size in the coordinates of an output
    /// laid out at `scale` (the full-resolution focal length, or a fraction
    /// of it).
    pub fn at_scale(&self, scale: f64) -> ((f64, f64), f64) {
        let r = scale / self.scale;
        ((self.origin.0 * r, self.origin.1 * r), self.px * r)
    }

    /// The radius, in texels, of a blend `blend_px` output pixels wide in an
    /// output laid out at `scale`: what [`Self::share`] and the shader are
    /// given, so that the preview and the merge blend alike.
    pub fn blend_radius(&self, scale: f64, blend_px: f64) -> f64 {
        let (_, px) = self.at_scale(scale);
        (blend_px / 2.0 / px).clamp(1.0, MAX_BLEND_RADIUS)
    }

    /// The share frame `k` has of output point `(u, v)` given at `scale`:
    /// the tent-weighted fraction of the texels within `radius` (in texels)
    /// that it owns. `None` where no texel in reach is owned at all — the
    /// map has nothing to say there, and the caller falls back to its
    /// feather.
    ///
    /// This is the function `merge.wgsl`'s `seam_share` repeats; the two
    /// must agree.
    pub fn share(&self, k: usize, u: f64, v: f64, scale: f64, radius: f64) -> Option<f32> {
        let ((ou, ov), px) = self.at_scale(scale);
        let x = (u - ou) / px - 0.5;
        let y = (v - ov) / px - 0.5;
        let r = radius.max(1.0);
        let (x0, x1) = ((x - r).ceil() as i64, (x + r).floor() as i64);
        let (y0, y1) = ((y - r).ceil() as i64, (y + r).floor() as i64);
        let (mut mine, mut all) = (0.0f64, 0.0f64);
        for j in y0.max(0)..=y1.min(self.height as i64 - 1) {
            let wy = 1.0 - (y - j as f64).abs() / r;
            if wy <= 0.0 {
                continue;
            }
            for i in x0.max(0)..=x1.min(self.width as i64 - 1) {
                let wx = 1.0 - (x - i as f64).abs() / r;
                if wx <= 0.0 {
                    continue;
                }
                let l = self.labels[j as usize * self.width + i as usize];
                if l == NONE {
                    continue;
                }
                all += wx * wy;
                if usize::from(l) == k {
                    mine += wx * wy;
                }
            }
        }
        (all > 0.0).then(|| (mine / all) as f32)
    }
}

/// One frame warped onto the map: its gain-corrected value and its distance
/// from its own edge (in proxy pixels) per texel, NaN where it does not
/// reach.
struct Warped {
    value: Vec<f32>,
    edge: Vec<f32>,
}

/// Lay seams across the overlaps of `proxies`, aligned by `cameras` (at the
/// proxies' scale), with `gains` the linear multipliers the merge will
/// apply. `None` if the frames project nowhere or there are more than
/// [`MAX_FRAMES`].
pub fn find(
    proxies: &[&Gray],
    cameras: &Cameras,
    gains: &[f32],
    projection: Projection,
    opts: &SeamOptions,
) -> Option<SeamMap> {
    let n = proxies.len();
    if n == 0 || n > MAX_FRAMES || cameras.rotations.len() != n || gains.len() != n {
        return None;
    }
    let (fw, fh) = (proxies[0].width as f64, proxies[0].height as f64);
    let scale = cameras.focal;
    let bounds = projection::bounds(projection, scale, cameras, (fw, fh))?;
    let width = opts.max_width.min(bounds.width().ceil() as usize).max(1);
    let px = bounds.width() / width as f64;
    let height = ((bounds.height() / px).ceil() as usize).max(1);
    let mut map = SeamMap {
        width,
        height,
        scale,
        origin: (bounds.min_u, bounds.min_v),
        px,
        labels: vec![NONE; width * height],
    };

    // Where each frame's centre lands, in texels: what orders the frames
    // and orients each cut.
    let centres: Vec<(f64, f64)> = (0..n)
        .map(|k| {
            let d = cameras.bearing(k, (0.0, 0.0));
            projection
                .from_direction(scale, d)
                .map(|(u, v)| ((u - bounds.min_u) / px, (v - bounds.min_v) / px))
                .unwrap_or((width as f64 / 2.0, height as f64 / 2.0))
        })
        .collect();

    // The composite so far: what its owner saw, and how far from the
    // owner's edge.
    let mut value = vec![f32::NAN; width * height];
    let mut edge = vec![f32::NAN; width * height];

    for k in order(&centres, (width as f64 / 2.0, height as f64 / 2.0)) {
        let w = warp(&map, proxies[k], cameras, k, gains[k], projection);
        let overlap: Vec<usize> = (0..width * height)
            .filter(|&i| map.labels[i] != NONE && !w.value[i].is_nan())
            .collect();
        // Texels nobody owns yet are the new frame's without a cut.
        let mut take: Vec<bool> = map
            .labels
            .iter()
            .zip(&w.value)
            .map(|(&l, v)| l == NONE && !v.is_nan())
            .collect();
        if !overlap.is_empty() {
            cut(
                &map, &value, &edge, &w, &overlap, &centres, k, opts, &mut take,
            );
        }
        for i in 0..width * height {
            if take[i] {
                map.labels[i] = k as u8;
                value[i] = w.value[i];
                edge[i] = w.edge[i];
            }
        }
    }
    Some(map)
}

/// The order frames are laid down in: the one nearest the middle first,
/// then always the unplaced frame nearest any placed one, so that each new
/// frame meets the composite along an overlap rather than across a gap.
fn order(centres: &[(f64, f64)], middle: (f64, f64)) -> Vec<usize> {
    let d2 = |a: (f64, f64), b: (f64, f64)| (a.0 - b.0).powi(2) + (a.1 - b.1).powi(2);
    let n = centres.len();
    let mut placed = vec![false; n];
    let mut out = Vec::with_capacity(n);
    let first = (0..n)
        .min_by(|&a, &b| d2(centres[a], middle).total_cmp(&d2(centres[b], middle)))
        .expect("at least one frame");
    placed[first] = true;
    out.push(first);
    while out.len() < n {
        let next = (0..n)
            .filter(|&k| !placed[k])
            .min_by(|&a, &b| {
                let near = |k: usize| {
                    out.iter()
                        .map(|&p| d2(centres[k], centres[p]))
                        .fold(f64::MAX, f64::min)
                };
                near(a).total_cmp(&near(b))
            })
            .expect("an unplaced frame");
        placed[next] = true;
        out.push(next);
    }
    out
}

/// Frame `k` sampled at every texel's centre, bilinearly. The proxy is
/// gamma-encoded grey, so the gain (linear) becomes `gain^(1/2.2)` on it.
fn warp(
    map: &SeamMap,
    g: &Gray,
    cameras: &Cameras,
    k: usize,
    gain: f32,
    projection: Projection,
) -> Warped {
    let (fw, fh) = (g.width as f64, g.height as f64);
    let gain = gain.max(1e-6).powf(1.0 / 2.2);
    let mut value = vec![f32::NAN; map.width * map.height];
    let mut edge = vec![f32::NAN; map.width * map.height];
    for ty in 0..map.height {
        let v = map.origin.1 + (ty as f64 + 0.5) * map.px;
        for tx in 0..map.width {
            let u = map.origin.0 + (tx as f64 + 0.5) * map.px;
            let d = projection.to_direction(map.scale, u, v);
            let Some((x, y)) = cameras.project(k, d) else {
                continue;
            };
            let (x, y) = (x + fw / 2.0 - 0.5, y + fh / 2.0 - 0.5);
            let e = x.min(fw - 1.0 - x).min(y).min(fh - 1.0 - y);
            if e < 0.0 {
                continue;
            }
            let (x0, y0) = (x.floor() as usize, y.floor() as usize);
            let (x1, y1) = ((x0 + 1).min(g.width - 1), (y0 + 1).min(g.height - 1));
            let (ax, ay) = ((x - x0 as f64) as f32, (y - y0 as f64) as f32);
            let at = |xx: usize, yy: usize| g.data[yy * g.width + xx];
            let top = at(x0, y0) * (1.0 - ax) + at(x1, y0) * ax;
            let bot = at(x0, y1) * (1.0 - ax) + at(x1, y1) * ax;
            let i = ty * map.width + tx;
            value[i] = (top * (1.0 - ay) + bot * ay) * gain;
            edge[i] = e as f32;
        }
    }
    Warped { value, edge }
}

/// Central-difference gradient magnitude of `plane` at texel `i`, from the
/// neighbours that exist.
fn detail(plane: &[f32], width: usize, height: usize, i: usize) -> f32 {
    let (x, y) = (i % width, i / width);
    let c = plane[i];
    let mut g = 0.0f32;
    let mut diff = |j: usize| {
        let n = plane[j];
        if !n.is_nan() {
            g = g.max((n - c).abs());
        }
    };
    if x > 0 {
        diff(i - 1);
    }
    if x + 1 < width {
        diff(i + 1);
    }
    if y > 0 {
        diff(i - width);
    }
    if y + 1 < height {
        diff(i + width);
    }
    g
}

/// Cut the overlap between the composite and frame `k`, marking in `take`
/// the overlap texels that go to `k`.
#[allow(clippy::too_many_arguments)]
fn cut(
    map: &SeamMap,
    value: &[f32],
    edge: &[f32],
    new: &Warped,
    overlap: &[usize],
    centres: &[(f64, f64)],
    k: usize,
    opts: &SeamOptions,
    take: &mut [bool],
) {
    let (w, h) = (map.width, map.height);

    // The raw cost per overlap texel.
    let mut raw = vec![f32::NAN; w * h];
    let margin = opts.edge_margin.max(1.0);
    for &i in overlap {
        let differ = (value[i] - new.value[i]).abs();
        let detail = detail(value, w, h, i).max(detail(&new.value, w, h, i));
        let near = (1.0 - edge[i].min(new.edge[i]) / margin).max(0.0);
        raw[i] = differ + opts.detail * detail + opts.edge * near * near + 1e-3;
    }
    // The worst over a small window: a texel is only cheap if its whole
    // neighbourhood agrees, so the path keeps at least the blend's radius
    // clear of a difference rather than threading the one lucky texel
    // beside it — the blend straddles the path by that much and would
    // otherwise reach the difference anyway.
    let r = opts.smoothing as isize;
    let mut cost = vec![OUTSIDE; w * h];
    for &i in overlap {
        let (x, y) = ((i % w) as isize, (i / w) as isize);
        let mut worst = 0.0f32;
        for dy in -r..=r {
            for dx in -r..=r {
                let (xx, yy) = (x + dx, y + dy);
                if xx < 0 || yy < 0 || xx >= w as isize || yy >= h as isize {
                    continue;
                }
                let c = raw[yy as usize * w + xx as usize];
                if !c.is_nan() {
                    worst = worst.max(c);
                }
            }
        }
        cost[i] = worst;
    }

    // The axis the cut crosses: from the composite's frames, weighted by how
    // much of the overlap each owns, to the new frame.
    let mut from = (0.0f64, 0.0f64);
    for &i in overlap {
        let c = centres[usize::from(map.labels[i])];
        from = (from.0 + c.0, from.1 + c.1);
    }
    let m = overlap.len() as f64;
    from = (from.0 / m, from.1 / m);
    let to = centres[k];
    let (mut ax, mut ay) = (to.0 - from.0, to.1 - from.1);
    let len = (ax * ax + ay * ay).sqrt();
    if len < 1e-6 {
        (ax, ay) = (1.0, 0.0);
    } else {
        (ax, ay) = (ax / len, ay / len);
    }
    // Along the cut: perpendicular to the axis.
    let (bx, by) = (-ay, ax);

    // The overlap's extent in (s along the cut, t across it).
    let st = |i: usize| {
        let (x, y) = ((i % w) as f64 + 0.5, (i / w) as f64 + 0.5);
        (x * bx + y * by, x * ax + y * ay)
    };
    let (mut s0, mut s1, mut t0, mut t1) = (f64::MAX, f64::MIN, f64::MAX, f64::MIN);
    for &i in overlap {
        let (s, t) = st(i);
        s0 = s0.min(s);
        s1 = s1.max(s);
        t0 = t0.min(t);
        t1 = t1.max(t);
    }
    let rows = (s1 - s0).round() as usize + 1;
    let cols = (t1 - t0).round() as usize + 1;

    // The grid in (s, t), each cell sampled from the texel it falls in, so
    // that a rotated overlap has no holes.
    let mut grid = vec![OUTSIDE; rows * cols];
    let mut any = vec![false; rows];
    for si in 0..rows {
        for ti in 0..cols {
            let (s, t) = (s0 + si as f64, t0 + ti as f64);
            let x = s * bx + t * ax;
            let y = s * by + t * ay;
            if x < 0.0 || y < 0.0 {
                continue;
            }
            let (x, y) = (x as usize, y as usize);
            if x >= w || y >= h {
                continue;
            }
            let c = cost[y * w + x];
            if c < OUTSIDE {
                grid[si * cols + ti] = c;
                any[si] = true;
            }
        }
    }

    // Dynamic programming down the rows: the path moves at most one column
    // per row, and starts afresh after a row with no overlap in it.
    let mut acc = grid.clone();
    let mut from_col = vec![0u32; rows * cols];
    for si in 1..rows {
        if !any[si] {
            continue;
        }
        let prev = &acc[(si - 1) * cols..si * cols].to_vec();
        if !any[si - 1] {
            continue;
        }
        for ti in 0..cols {
            let mut best = (prev[ti], ti);
            if ti > 0 && prev[ti - 1] < best.0 {
                best = (prev[ti - 1], ti - 1);
            }
            if ti + 1 < cols && prev[ti + 1] < best.0 {
                best = (prev[ti + 1], ti + 1);
            }
            acc[si * cols + ti] += best.0;
            from_col[si * cols + ti] = best.1 as u32;
        }
    }
    // Back up from the end of each run of rows with overlap.
    let mut seam = vec![usize::MAX; rows];
    let mut si = rows;
    while si > 0 {
        si -= 1;
        if !any[si] {
            continue;
        }
        let row = &acc[si * cols..(si + 1) * cols];
        let mut t = (0..cols)
            .min_by(|&a, &b| row[a].total_cmp(&row[b]))
            .unwrap_or(0);
        loop {
            seam[si] = t;
            if si == 0 || !any[si - 1] {
                break;
            }
            t = from_col[si * cols + t] as usize;
            si -= 1;
        }
    }

    // The new frame takes the side of the path its centre is on.
    for &i in overlap {
        let (s, t) = st(i);
        let si = ((s - s0).round() as usize).min(rows - 1);
        let ti = (t - t0).round();
        if seam[si] != usize::MAX && ti >= seam[si] as f64 {
            take[i] = true;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::linalg::{Mat3, Vec3};

    /// A scene as a function of direction, and frames of it rendered by the
    /// same cameras the seam reads.
    fn render(
        cameras: &Cameras,
        k: usize,
        size: (usize, usize),
        scene: impl Fn(Vec3) -> f32,
    ) -> Gray {
        let (w, h) = size;
        let mut data = vec![0.0; w * h];
        for y in 0..h {
            for x in 0..w {
                let p = (
                    x as f64 + 0.5 - w as f64 / 2.0,
                    y as f64 + 0.5 - h as f64 / 2.0,
                );
                data[y * w + x] = scene(cameras.bearing(k, p));
            }
        }
        Gray {
            width: w,
            height: h,
            data,
        }
    }

    fn yaw(a: f64) -> Mat3 {
        let (s, c) = a.sin_cos();
        Mat3([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]])
    }

    /// Smooth, with a little texture: what a sky over a slope looks like to
    /// the cost.
    fn landscape(d: Vec3) -> f32 {
        let (x, y) = (d.x() / d.z(), d.y() / d.z());
        let texture = if y > 0.1 { 0.1 * (y * 40.0).sin() } else { 0.0 };
        (0.5 + 0.2 * (x * 3.0).sin() + texture).clamp(0.0, 1.0) as f32
    }

    fn pair() -> Cameras {
        Cameras {
            rotations: vec![Mat3::IDENTITY, yaw(0.35)],
            focal: 300.0,
        }
    }

    #[test]
    fn one_frame_owns_everything_it_reaches() {
        let cameras = Cameras {
            rotations: vec![Mat3::IDENTITY],
            focal: 300.0,
        };
        let g = render(&cameras, 0, (320, 240), landscape);
        let map = find(
            &[&g],
            &cameras,
            &[1.0],
            Projection::Perspective,
            &Default::default(),
        )
        .unwrap();
        let owned = map.labels.iter().filter(|&&l| l == 0).count();
        assert!(owned as f64 > 0.95 * (map.width * map.height) as f64);
    }

    #[test]
    fn each_frame_keeps_its_own_side() {
        let cameras = pair();
        let frames: Vec<Gray> = (0..2)
            .map(|k| render(&cameras, k, (320, 240), landscape))
            .collect();
        let refs: Vec<&Gray> = frames.iter().collect();
        let map = find(
            &refs,
            &cameras,
            &[1.0, 1.0],
            Projection::Cylindrical,
            &Default::default(),
        )
        .unwrap();
        let mid = map.height / 2 * map.width;
        assert_eq!(map.labels[mid + 2], 0, "the left edge is frame 0's alone");
        assert_eq!(
            map.labels[mid + map.width - 3],
            1,
            "the right edge is frame 1's"
        );
        // One change of owner along every row that both frames cross.
        for y in 0..map.height {
            let row = &map.labels[y * map.width..(y + 1) * map.width];
            let owned: Vec<u8> = row.iter().copied().filter(|&l| l != NONE).collect();
            let changes = owned.windows(2).filter(|p| p[0] != p[1]).count();
            assert!(changes <= 1, "row {y} changes owner {changes} times");
        }
    }

    #[test]
    fn the_seam_goes_round_what_only_one_frame_saw() {
        // Frame 1 saw something frame 0 did not — a figure that walked into
        // the overlap — in the middle of where the two meet.
        let cameras = pair();
        let figure = Vec3::new(0.175f64.sin(), 0.0, 0.175f64.cos());
        let walker = |d: Vec3| {
            let near = (d.x() - figure.x()).abs() < 0.04 && (d.y() - figure.y()).abs() < 0.15;
            if near {
                0.95
            } else {
                landscape(d)
            }
        };
        let frames = [
            render(&cameras, 0, (320, 240), landscape),
            render(&cameras, 1, (320, 240), walker),
        ];
        let refs: Vec<&Gray> = frames.iter().collect();
        let map = find(
            &refs,
            &cameras,
            &[1.0, 1.0],
            Projection::Cylindrical,
            &Default::default(),
        )
        .unwrap();
        // Every texel of the figure is taken from the same frame, with a
        // blend radius of room to spare, so it is either all there or not at
        // all — never half.
        let (u, v) = Projection::Cylindrical
            .from_direction(map.scale, figure)
            .unwrap();
        let mut owners = std::collections::HashSet::new();
        // The figure's extent on the surface, plus the blend's radius.
        let radius = 3.0;
        let reach = |half: f64| half * map.scale + radius * map.px;
        let (ru, rv) = (reach(0.04), reach(0.15));
        let mut dv = -rv;
        while dv <= rv {
            let mut du = -ru;
            while du <= ru {
                let s = map.share(1, u + du, v + dv, map.scale, radius);
                owners.insert((s.unwrap() * 100.0).round() as i32);
                du += map.px;
            }
            dv += map.px;
        }
        assert_eq!(owners.len(), 1, "the figure is split: shares {owners:?}");
    }

    #[test]
    fn share_is_a_blend_across_the_seam_and_whole_away_from_it() {
        let map = SeamMap {
            width: 8,
            height: 1,
            scale: 1.0,
            origin: (0.0, 0.0),
            px: 1.0,
            labels: vec![0, 0, 0, 0, 1, 1, 1, 1],
        };
        assert_eq!(map.share(0, 1.5, 0.5, 1.0, 2.0), Some(1.0));
        assert_eq!(map.share(1, 6.5, 0.5, 1.0, 2.0), Some(1.0));
        let at_seam = map.share(0, 4.0, 0.5, 1.0, 2.0).unwrap();
        assert!((at_seam - 0.5).abs() < 1e-6, "{at_seam}");
        // And at twice the scale, the same point is twice as far out.
        assert_eq!(
            map.share(0, 8.0, 1.0, 2.0, 2.0),
            map.share(0, 4.0, 0.5, 1.0, 2.0)
        );
        let empty = SeamMap {
            labels: vec![NONE; 8],
            ..map
        };
        assert_eq!(empty.share(0, 4.0, 0.5, 1.0, 2.0), None);
    }
}