//! Optical flow by block matching, and motion-compensated frame interpolation (the Pixel Motion
//! mode of frame blending).
//!
//! The method is the textbook hierarchical block matcher used in video coding and frame-rate
//! conversion: build a luma pyramid of both frames; at the coarsest level, find for each block of
//! frame A the integer displacement into frame B with the smallest sum of absolute differences
//! (SAD) in a search window; at each finer level, start from the doubled coarse vectors (and the
//! neighbours' vectors) and refine in a small window; finish with a sub-pixel parabola fit of the
//! SAD around the best vector and a 3×3 vector median to remove outliers.
//!
//! Interpolation at a fraction `w` between the frames warps both towards the in-between time —
//! A forward by `w·F`, B back by `(1 − w)·F` — and cross-fades them, so moving detail stays sharp
//! instead of double-exposing as with a plain frame mix.

use rayon::prelude::*;

use crate::Image;

/// A block motion field: one vector per `block × block` tile of frame A (in pixels, A → B).
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Flow {
    pub width: u32,
    pub height: u32,
    pub block: u32,
    /// Tiles across / down.
    pub cols: usize,
    pub rows: usize,
    pub v: Vec<[f32; 2]>,
}

impl Flow {
    /// The motion at pixel position (x, y), bilinearly interpolated between tile centres.
    pub fn at(&self, x: f64, y: f64) -> [f32; 2] {
        if self.v.is_empty() {
            return [0.0; 2];
        }
        let b = self.block as f64;
        let fx = (x / b - 0.5).clamp(0.0, (self.cols - 1) as f64);
        let fy = (y / b - 0.5).clamp(0.0, (self.rows - 1) as f64);
        let (x0, y0) = (fx.floor() as usize, fy.floor() as usize);
        let (x1, y1) = ((x0 + 1).min(self.cols - 1), (y0 + 1).min(self.rows - 1));
        let (tx, ty) = ((fx - x0 as f64) as f32, (fy - y0 as f64) as f32);
        let g = |c: usize, r: usize| self.v[r * self.cols + c];
        let mut o = [0.0; 2];
        for i in 0..2 {
            let top = g(x0, y0)[i] + (g(x1, y0)[i] - g(x0, y0)[i]) * tx;
            let bot = g(x0, y1)[i] + (g(x1, y1)[i] - g(x0, y1)[i]) * tx;
            o[i] = top + (bot - top) * ty;
        }
        o
    }
}

/// A single-channel plane.
struct Plane {
    w: usize,
    h: usize,
    d: Vec<f32>,
}

impl Plane {
    fn luma(img: &Image) -> Plane {
        let d = img.data.par_iter().map(|p| 0.2126 * p[0] + 0.7152 * p[1] + 0.0722 * p[2] + 0.25 * p[3]).collect();
        Plane { w: img.width as usize, h: img.height as usize, d }
    }
    fn half(&self) -> Plane {
        let (w, h) = ((self.w / 2).max(1), (self.h / 2).max(1));
        let d = (0..w * h)
            .into_par_iter()
            .map(|i| {
                let (x, y) = (i % w * 2, i / w * 2);
                let g = |dx: usize, dy: usize| self.d[(y + dy).min(self.h - 1) * self.w + (x + dx).min(self.w - 1)];
                (g(0, 0) + g(1, 0) + g(0, 1) + g(1, 1)) * 0.25
            })
            .collect();
        Plane { w, h, d }
    }
    #[inline]
    fn get(&self, x: i64, y: i64) -> f32 {
        self.d[y.clamp(0, self.h as i64 - 1) as usize * self.w + x.clamp(0, self.w as i64 - 1) as usize]
    }
}

/// SAD of the `bs` block at (bx, by) in `a` against `b` displaced by (dx, dy).
fn sad(a: &Plane, b: &Plane, bx: i64, by: i64, bs: i64, dx: i64, dy: i64) -> f32 {
    let mut s = 0.0;
    for y in by..by + bs {
        for x in bx..bx + bs {
            s += (a.get(x, y) - b.get(x + dx, y + dy)).abs();
        }
    }
    s
}

/// Block-matching flow from `a` to `b` (same size). `block` is the tile size at full resolution;
/// `radius` the search range at the coarsest level (motion up to about `radius · 2^levels`).
pub fn block_flow(a: &Image, b: &Image, block: u32, radius: i32) -> Flow {
    block_flow_window(a, b, block, block, radius)
}

/// [`block_flow`] with separate vector spacing and matching window: one vector per
/// `tile × tile` tile, each matched with the `window × window` block centred on its tile
/// (Timewarp's Vector Detail and Block Size).
pub fn block_flow_window(a: &Image, b: &Image, tile: u32, window: u32, radius: i32) -> Flow {
    let bs = tile.max(2) as usize;
    let win = window.max(4) as i64;
    // Window origin relative to the tile origin.
    let wo = bs as i64 / 2 - win / 2;
    let (w, h) = (a.width as usize, a.height as usize);
    let cols = w.div_ceil(bs).max(1);
    let rows = h.div_ceil(bs).max(1);
    let mut flow = Flow { width: a.width, height: a.height, block: bs as u32, cols, rows, v: vec![[0.0; 2]; cols * rows] };
    if a.width != b.width || a.height != b.height || w < bs || h < bs || (w as i64) < win || (h as i64) < win {
        return flow;
    }
    // Pyramids (level 0 = full resolution); stop while blocks still fit.
    let mut pa = vec![Plane::luma(a)];
    let mut pb = vec![Plane::luma(b)];
    while pa.len() < 4 && pa.last().is_some_and(|p| p.w / 2 >= (bs as i64).max(win) as usize * 2 && p.h / 2 >= (bs as i64).max(win) as usize * 2) {
        let (na, nb) = (pa.last().map(Plane::half), pb.last().map(Plane::half));
        pa.extend(na);
        pb.extend(nb);
    }
    let top = pa.len() - 1;
    // Integer vectors at the current level, per full-resolution tile.
    let mut cur: Vec<[i64; 2]> = vec![[0, 0]; cols * rows];
    for level in (0..=top).rev() {
        let (la, lb) = (&pa[level], &pb[level]);
        let k = 1i64 << level;
        // The block covers the tile's footprint at this level (at least 4 pixels).
        let lbs = (win / k).max(4);
        let r = if level == top { radius as i64 } else { 1 };
        let prev = cur.clone();
        cur = (0..cols * rows)
            .into_par_iter()
            .map(|i| {
                let (c, rr) = ((i % cols) as i64, (i / cols) as i64);
                let bx = (c * bs as i64 + wo) / k;
                let by = (rr * bs as i64 + wo) / k;
                // Candidates: this tile's and its neighbours' vectors from the coarser level.
                let mut seeds: Vec<[i64; 2]> = vec![[0, 0]];
                if level != top {
                    for (dc, dr) in [(0, 0), (-1, 0), (1, 0), (0, -1), (0, 1)] {
                        let (nc, nr) = (c + dc, rr + dr);
                        if nc >= 0 && nr >= 0 && (nc as usize) < cols && (nr as usize) < rows {
                            let v = prev[nr as usize * cols + nc as usize];
                            seeds.push([v[0] * 2, v[1] * 2]);
                        }
                    }
                }
                seeds.sort_unstable();
                seeds.dedup();
                let mut best = ([0i64, 0i64], f32::INFINITY);
                for s in seeds {
                    for dy in -r..=r {
                        for dx in -r..=r {
                            let d = [s[0] + dx, s[1] + dy];
                            // A tiny bias towards short vectors breaks ties in flat areas.
                            let e = sad(la, lb, bx, by, lbs, d[0], d[1]) + 1e-4 * (d[0].abs() + d[1].abs()) as f32;
                            if e < best.1 {
                                best = (d, e);
                            }
                        }
                    }
                }
                best.0
            })
            .collect();
    }
    // Sub-pixel refinement (parabola through the SAD of the neighbours) at full resolution.
    let (la, lb) = (&pa[0], &pb[0]);
    let bsi = bs as i64;
    let wsi = win;
    let fine: Vec<[f32; 2]> = (0..cols * rows)
        .into_par_iter()
        .map(|i| {
            let (bx, by) = ((i % cols) as i64 * bsi + wo, (i / cols) as i64 * bsi + wo);
            let d = cur[i];
            let e0 = sad(la, lb, bx, by, wsi, d[0], d[1]);
            let fit = |em: f32, ep: f32| {
                let den = em - 2.0 * e0 + ep;
                if den > 1e-6 { (0.5 * (em - ep) / den).clamp(-0.5, 0.5) } else { 0.0 }
            };
            let sx = fit(sad(la, lb, bx, by, wsi, d[0] - 1, d[1]), sad(la, lb, bx, by, wsi, d[0] + 1, d[1]));
            let sy = fit(sad(la, lb, bx, by, wsi, d[0], d[1] - 1), sad(la, lb, bx, by, wsi, d[0], d[1] + 1));
            [d[0] as f32 + sx, d[1] as f32 + sy]
        })
        .collect();
    // 3×3 vector median (component-wise) against outliers.
    flow.v = (0..cols * rows)
        .into_par_iter()
        .map(|i| {
            let (c, r) = ((i % cols) as i64, (i / cols) as i64);
            let mut xs = Vec::with_capacity(9);
            let mut ys = Vec::with_capacity(9);
            for dr in -1..=1 {
                for dc in -1..=1 {
                    let (nc, nr) = (c + dc, r + dr);
                    if nc >= 0 && nr >= 0 && (nc as usize) < cols && (nr as usize) < rows {
                        let v = fine[nr as usize * cols + nc as usize];
                        xs.push(v[0]);
                        ys.push(v[1]);
                    }
                }
            }
            xs.sort_by(f32::total_cmp);
            ys.sort_by(f32::total_cmp);
            [xs[xs.len() / 2], ys[ys.len() / 2]]
        })
        .collect();
    flow
}

/// Cross-fade of two frames: `a·(1 − w) + b·w` (Frame Mix).
pub fn mix(a: &Image, b: &Image, w: f32) -> Image {
    if a.width != b.width || a.height != b.height {
        return if w < 0.5 { a.clone() } else { b.clone() };
    }
    let mut out = a.clone();
    out.data.par_iter_mut().zip(b.data.par_iter()).for_each(|(p, q)| {
        for c in 0..4 {
            p[c] += (q[c] - p[c]) * w;
        }
    });
    out
}

/// Motion-compensated frame at fraction `w` (0 = `a`, 1 = `b`) between two frames (Pixel Motion).
pub fn interpolate(a: &Image, b: &Image, w: f32) -> Image {
    if a.width != b.width || a.height != b.height {
        return mix(a, b, w);
    }
    let flow = block_flow(a, b, 8, 4);
    let mut out = Image::new(a.width, a.height);
    let wd = w as f64;
    out.rows_mut().for_each(|(y, row)| {
        let cy = y as f64 + 0.5;
        for (x, px) in row.iter_mut().enumerate() {
            let cx = x as f64 + 0.5;
            let f = flow.at(cx, cy);
            let (fx, fy) = (f[0] as f64, f[1] as f64);
            let pa = a.sample_bilinear_clamped(cx - wd * fx, cy - wd * fy);
            let pb = b.sample_bilinear_clamped(cx + (1.0 - wd) * fx, cy + (1.0 - wd) * fy);
            for c in 0..4 {
                px[c] = pa[c] + (pb[c] - pa[c]) * w;
            }
        }
    });
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A textured image shifted right by `dx` pixels.
    fn texture(w: u32, h: u32, dx: f64) -> Image {
        let mut img = Image::new(w, h);
        img.rows_mut().for_each(|(y, row)| {
            for (x, p) in row.iter_mut().enumerate() {
                let (u, v) = (x as f64 - dx, y as f64);
                let s = 0.5 + 0.2 * (u * 0.21 + v * 0.07).sin() + 0.15 * (u * 0.05 - v * 0.19).cos() + 0.1 * (u * 0.013 * v * 0.011).sin();
                let s = s as f32;
                *p = [s, s * 0.8, 1.0 - s, 1.0];
            }
        });
        img
    }

    fn mae(a: &Image, b: &Image, margin: u32) -> f32 {
        let mut s = 0.0;
        let mut n = 0;
        for y in margin..a.height - margin {
            for x in margin..a.width - margin {
                let (p, q) = (a.get(x as i64, y as i64), b.get(x as i64, y as i64));
                s += (0..3).map(|c| (p[c] - q[c]).abs()).sum::<f32>();
                n += 3;
            }
        }
        s / n as f32
    }

    #[test]
    fn flow_finds_translation() {
        let a = texture(128, 96, 0.0);
        let b = texture(128, 96, 6.0);
        let f = block_flow(&a, &b, 8, 4);
        let m = f.at(64.0, 48.0);
        assert!((m[0] - 6.0).abs() < 0.3 && m[1].abs() < 0.3, "{m:?}");
    }

    #[test]
    fn pixel_motion_beats_frame_mix_on_translation() {
        let a = texture(128, 96, 0.0);
        let b = texture(128, 96, 6.0);
        let truth = texture(128, 96, 3.0);
        let pm = mae(&interpolate(&a, &b, 0.5), &truth, 12);
        let fm = mae(&mix(&a, &b, 0.5), &truth, 12);
        assert!(pm < fm * 0.25, "pixel motion {pm} vs frame mix {fm}");
    }
}