//! Quality of the Remove Tool's fill (`photocraft_algo::remove::remove_with`) on real photos.
//! `cargo run --release -p photocraft-engine --example remove_quality -- <images dir> [out dir]`
//!
//! The images are `01.jpg`, `02.jpg`, `03.jpg`, `04.jpg` and `06.jpg`: public-domain photos from
//! Wikimedia Commons, not committed (the list is in `docs/development.md` › Remove Tool quality).
//! Each gets holes over background (discs and thin lines, placed in fractions of the image size)
//! that are filled through a window like the command's. Measured in the hole against the original:
//!
//! * `psnr`: fidelity in dB (a plausible fill need not match, so read it as a trend);
//! * `texture`: mean gradient magnitude of the fill over the original's (1 is ideal; below 1 the
//!   fill is smoother than what it replaces);
//! * `seam`: the same on the band along the hole's edge (above 1 the edge shows).
//!
//! With an out dir, each case is also written as `<image>-<case>.png`: original | hole | fill.
use std::time::Instant;

use photocraft_codecs::{ChannelLayout, EncodeOptions, Format, Image, SampleType};
use photocraft_raster::Interrupt;

/// A hole: a disc (centre in fractions of the width and height, radius in fractions of the width)
/// or a line (ends in fractions, width in pixels).
#[derive(Clone, Copy)]
enum Shape {
    Disc { c: (f64, f64), r: f64 },
    Line { a: (f64, f64), b: (f64, f64), width: f64 },
}

fn cases(image: &str) -> Vec<(&'static str, Shape)> {
    use Shape::{Disc, Line};
    match image {
        // The Tetons and the Snake River: forest, river bend, ridge, clouds.
        "01" => vec![
            ("forest", Disc { c: (0.75, 0.62), r: 0.03 }),
            ("river-bend", Disc { c: (0.45, 0.72), r: 0.045 }),
            ("ridge", Disc { c: (0.5, 0.28), r: 0.035 }),
            ("clouds", Disc { c: (0.25, 0.12), r: 0.06 }),
            ("wire", Line { a: (0.3, 0.55), b: (0.7, 0.85), width: 5.0 }),
        ],
        // A dam in a canyon: rock, water, the dam's wall.
        "02" => vec![
            ("rock", Disc { c: (0.85, 0.45), r: 0.04 }),
            ("water", Disc { c: (0.6, 0.7), r: 0.035 }),
            ("dam-wall", Disc { c: (0.35, 0.25), r: 0.025 }),
            ("wire", Line { a: (0.05, 0.6), b: (0.95, 0.5), width: 4.0 }),
        ],
        // Oak trees by a gravel path: grass, gravel, a trunk, branches.
        "03" => vec![
            ("grass", Disc { c: (0.75, 0.85), r: 0.045 }),
            ("gravel", Disc { c: (0.25, 0.85), r: 0.035 }),
            ("trunk", Disc { c: (0.48, 0.45), r: 0.02 }),
            ("branches", Disc { c: (0.8, 0.2), r: 0.04 }),
        ],
        // A meadow anemone in grass, shallow depth of field.
        "04" => vec![
            ("bokeh", Disc { c: (0.2, 0.35), r: 0.05 }),
            ("stems", Disc { c: (0.85, 0.6), r: 0.04 }),
            ("wire", Line { a: (0.1, 0.1), b: (0.9, 0.3), width: 6.0 }),
        ],
        // A coast from above: surf, sand, mountains.
        "06" => vec![
            ("surf", Disc { c: (0.75, 0.85), r: 0.045 }),
            ("sand", Disc { c: (0.3, 0.75), r: 0.04 }),
            ("mountains", Disc { c: (0.5, 0.12), r: 0.03 }),
            ("wire", Line { a: (0.0, 0.55), b: (1.0, 0.65), width: 5.0 }),
        ],
        _ => Vec::new(),
    }
}

/// The hole of `shape` on a `w × h` image.
fn mask(shape: Shape, w: usize, h: usize) -> Vec<bool> {
    let (wf, hf) = (w as f64, h as f64);
    (0..w * h)
        .map(|i| {
            let (x, y) = ((i % w) as f64 + 0.5, (i / w) as f64 + 0.5);
            match shape {
                Shape::Disc { c, r } => (x - c.0 * wf).hypot(y - c.1 * hf) < r * wf,
                Shape::Line { a, b, width } => {
                    let (ax, ay, bx, by) = (a.0 * wf, a.1 * hf, b.0 * wf, b.1 * hf);
                    let (dx, dy) = (bx - ax, by - ay);
                    let t = (((x - ax) * dx + (y - ay) * dy) / (dx * dx + dy * dy)).clamp(0.0, 1.0);
                    (x - ax - t * dx).hypot(y - ay - t * dy) < width / 2.0
                }
            }
        })
        .collect()
}

/// Bounding box `(x0, y0, x1, y1)` of the set cells.
fn bbox(w: usize, m: &[bool]) -> Option<(usize, usize, usize, usize)> {
    let mut b: Option<(usize, usize, usize, usize)> = None;
    for (i, _) in m.iter().enumerate().filter(|(_, v)| **v) {
        let (x, y) = (i % w, i / w);
        b = Some(b.map_or((x, y, x + 1, y + 1), |(x0, y0, x1, y1)| (x0.min(x), y0.min(y), x1.max(x + 1), y1.max(y + 1))));
    }
    b
}

/// A `(x0, y0, x1, y1)` crop of an interleaved buffer.
fn crop(src: &[f32], w: usize, ch: usize, r: (usize, usize, usize, usize)) -> Vec<f32> {
    (r.1..r.3).flat_map(|y| src[(y * w + r.0) * ch..(y * w + r.2) * ch].iter().copied()).collect()
}

/// Gradient magnitude (forward differences, summed over channels) at every cell.
fn gradient(img: &[f32], w: usize, h: usize, ch: usize) -> Vec<f32> {
    (0..w * h)
        .map(|i| {
            let (x, y) = (i % w, i / w);
            let (r, d) = ((y * w + (x + 1).min(w - 1)) * ch, ((y + 1).min(h - 1) * w + x) * ch);
            (0..ch).map(|c| (img[r + c] - img[i * ch + c]).abs() + (img[d + c] - img[i * ch + c]).abs()).sum()
        })
        .collect()
}

fn mean_over(v: &[f32], m: &[bool]) -> f64 {
    let (s, n) = v.iter().zip(m).filter(|(_, m)| **m).fold((0.0f64, 0usize), |(s, n), (v, _)| (s + f64::from(*v), n + 1));
    if n == 0 { 0.0 } else { s / n as f64 }
}

struct Score {
    psnr: f64,
    texture: f64,
    seam: f64,
}

fn score(orig: &[f32], out: &[f32], hole: &[bool], w: usize, h: usize, ch: usize) -> Score {
    let (mut se, mut n) = (0.0f64, 0usize);
    for (i, _) in hole.iter().enumerate().filter(|(_, m)| **m) {
        for c in 0..ch {
            se += f64::from(out[i * ch + c] - orig[i * ch + c]).powi(2);
            n += 1;
        }
    }
    let psnr = if se == 0.0 { 99.0 } else { 10.0 * (n as f64 / se).log10() };
    let (go, gf) = (gradient(orig, w, h, ch), gradient(out, w, h, ch));
    // The edge band: cells within 2 px of the hole's boundary, on both sides.
    let grown = photocraft_algo::remove::dilate(w, h, hole, 2);
    let inner: Vec<bool> = {
        let outside: Vec<bool> = hole.iter().map(|m| !*m).collect();
        photocraft_algo::remove::dilate(w, h, &outside, 2).iter().map(|o| !*o).collect()
    };
    let band: Vec<bool> = grown.iter().zip(&inner).map(|(g, i)| *g && !*i).collect();
    let ratio = |m: &[bool]| {
        let o = mean_over(&go, m);
        if o > 0.0 { mean_over(&gf, m) / o } else { 1.0 }
    };
    Score { psnr, texture: ratio(hole), seam: ratio(&band) }
}

fn write_png(path: &std::path::Path, data: &[f32], w: usize, h: usize, ch: usize) {
    let layout = if ch == 1 { ChannelLayout::Gray } else { ChannelLayout::Rgb };
    let img = Image::from_normalized(w as u32, h as u32, layout, SampleType::U8, data).expect("image");
    std::fs::write(path, photocraft_codecs::encode(&img, Format::Png, &EncodeOptions::default()).expect("encode")).expect("write");
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let dir = std::path::PathBuf::from(args.get(1).expect("usage: remove_quality <images dir> [out dir]"));
    let out_dir = args.get(2).map(std::path::PathBuf::from);
    if let Some(d) = &out_dir {
        std::fs::create_dir_all(d).expect("out dir");
    }
    println!("{:<4} {:<11} {:>7} {:>8} {:>6} {:>9}", "img", "case", "psnr", "texture", "seam", "ms");
    let (mut sum_psnr, mut sum_tex, mut sum_seam, mut sum_ms, mut count) = (0.0, 0.0, 0.0, 0.0, 0);
    for name in ["01", "02", "03", "04", "06"] {
        let Ok(bytes) = std::fs::read(dir.join(format!("{name}.jpg"))) else {
            println!("{name}: missing, skipped");
            continue;
        };
        let img = photocraft_codecs::decode(&bytes).expect("decode");
        let (w, h) = (img.width() as usize, img.height() as usize);
        let ch = img.layout().channels();
        let px = img.to_normalized();
        for (case, shape) in cases(name) {
            let full = mask(shape, w, h);
            let Some(b) = bbox(w, &full) else { continue };
            // The command's window: the hole plus three quarters of its extent (32..512 px).
            let ext = (b.2 - b.0).max(b.3 - b.1);
            let m = (ext * 3 / 4).clamp(32, 512);
            let win = (b.0.saturating_sub(m), b.1.saturating_sub(m), (b.2 + m).min(w), (b.3 + m).min(h));
            let (ww, wh) = (win.2 - win.0, win.3 - win.1);
            let orig = crop(&px, w, ch, win);
            let hole: Vec<bool> = (win.1..win.3).flat_map(|y| (win.0..win.2).map(move |x| (x, y))).map(|(x, y)| full[y * w + x]).collect();
            let mut holed = orig.clone();
            for (i, _) in hole.iter().enumerate().filter(|(_, m)| **m) {
                holed[i * ch..(i + 1) * ch].fill(1.0);
            }
            let t = Instant::now();
            let out = photocraft_algo::remove::remove_with(ww, wh, ch, &holed, &hole, 0x0052_e30e, &Interrupt::NONE).expect("never cancelled");
            let ms = t.elapsed().as_secs_f64() * 1000.0;
            let s = score(&orig, &out, &hole, ww, wh, ch);
            println!("{name:<4} {case:<11} {:>7.2} {:>8.3} {:>6.3} {ms:>9.1}", s.psnr, s.texture, s.seam);
            (sum_psnr, sum_tex, sum_seam, sum_ms, count) = (sum_psnr + s.psnr, sum_tex + s.texture, sum_seam + s.seam, sum_ms + ms, count + 1);
            if let Some(d) = &out_dir {
                // Original | hole | fill, cropped to the hole plus half its extent.
                let pad = ext / 2 + 8;
                let lb = b.0 - win.0;
                let tb = b.1 - win.1;
                let r = (lb.saturating_sub(pad), tb.saturating_sub(pad), (lb + (b.2 - b.0) + pad).min(ww), (tb + (b.3 - b.1) + pad).min(wh));
                let (cw, chh) = (r.2 - r.0, r.3 - r.1);
                let parts = [crop(&orig, ww, ch, r), crop(&holed, ww, ch, r), crop(&out, ww, ch, r)];
                let mut sheet = vec![0.0f32; cw * 3 * chh * ch];
                for (k, p) in parts.iter().enumerate() {
                    for y in 0..chh {
                        let dst = (y * cw * 3 + k * cw) * ch;
                        sheet[dst..dst + cw * ch].copy_from_slice(&p[y * cw * ch..(y + 1) * cw * ch]);
                    }
                }
                write_png(&d.join(format!("{name}-{case}.png")), &sheet, cw * 3, chh, ch);
            }
        }
    }
    if count > 0 {
        let n = f64::from(count);
        println!("{:<16} {:>7.2} {:>8.3} {:>6.3} {:>9.1}", "mean", sum_psnr / n, sum_tex / n, sum_seam / n, sum_ms / n);
    }
}