//! Every way of editing a mask, as frames you can watch.
//!
//! The mask tools are hard to review from a still: what makes them right is
//! how the mask *moves* as a stroke is painted, as a correction is subtracted,
//! as an edge is shaped. This renders that — a synthetic photograph and one
//! frame per step of each mode — so the pipeline's behaviour can be watched
//! before any of it is wired to a finger.
//!
//! ```sh
//! cargo run -p dr-gpu --example mask_modes --release -- out
//! ffmpeg -y -framerate 12 -i out/paint-%03d.ppm out/paint.gif
//! ```
//!
//! PPM for the reason every other example here writes it: no encoder
//! dependency, and ffmpeg, ImageMagick and every viewer read it.
//!
//! # What it is really showing
//!
//! The fold that builds a layer's mask (`MaskPass::render`), through the
//! composed shader that samples it. A part drawn in the wrong order, a
//! subtraction that took the base with it, an erase stroke that punched
//! through the selection underneath — each of those is a frame here that looks
//! wrong, and none of them is visible in a single rendered still.

use dr_gpu::{AdjustPass, DemosaicedImage, GpuContext, MaskPass};
use dr_pipeline::descriptor::ParamId;
use dr_pipeline::mask::{Join, MaskLayer, MaskPart, MaskSource, MaskStack, Morphology};
use dr_pipeline::operation::compose_full;
use dr_pipeline::spot::SpotSet;
use dr_pipeline::{ops, Framing};
use dr_types::ColourSpace;

const W: u32 = 480;
const H: u32 = 320;

fn main() {
    env_logger::init();
    let dir = std::env::args().nth(1).unwrap_or_else(|| "out".into());
    std::fs::create_dir_all(&dir).expect("output directory");

    let Some(ctx) = pollster::block_on(GpuContext::new_headless()).ok() else {
        eprintln!("no adapter; nothing to render");
        return;
    };

    let source = DemosaicedImage::from_rgba8(&ctx, &scene(), W, H).expect("upload");
    let mut masks = MaskPass::new(&ctx).expect("mask pass");
    let mut adjust = AdjustPass::new(&ctx);
    let mut shot = Shot {
        source: &source,
        masks: &mut masks,
        adjust: &mut adjust,
        dir: &dir,
    };

    write_ppm(&format!("{dir}/original.ppm"), &to_rgb(&scene()), W, H);

    paint(&mut shot);
    erase(&mut shot);
    subtract(&mut shot);
    invert(&mut shot);
    shape(&mut shot);

    println!("frames in {dir}/");
}

/// One layer, brightened hard, so the mask is legible rather than tasteful.
fn lifted(source: MaskSource) -> MaskLayer {
    let mut layer = MaskLayer::new("m1", source);
    layer.set_param("exposure", ParamId("exposure"), 1.6);
    layer.set_param("saturation", ParamId("vibrance"), 0.6);
    layer
}

/// A stroke painted from left to right across the subject, one frame per dab.
///
/// Each frame is the whole mask rasterised again, which is what the
/// application does today on every shape change — so the frames are also a
/// crude answer to "is a stroke's cost growing as it is painted".
fn paint(shot: &mut Shot) {
    let mut layer = lifted(MaskSource::brush());
    layer.begin_stroke(0, false, 0.13, 0.5, 1.0);
    for (i, x) in steps(0.18, 0.82, 28).enumerate() {
        layer.extend_stroke(0, x, 0.52 + 0.06 * (x * 9.0).sin());
        shot.frame("paint", i, &layer);
    }
    layer.end_stroke(0);
}

/// The same layer, with an erase stroke taken back through the middle of it.
fn erase(shot: &mut Shot) {
    let mut layer = lifted(MaskSource::brush());
    layer.begin_stroke(0, false, 0.16, 0.5, 1.0);
    for x in steps(0.18, 0.82, 20) {
        layer.extend_stroke(0, x, 0.5);
    }
    layer.end_stroke(0);

    for i in 0..8 {
        shot.frame("erase", i, &layer);
    }
    layer.begin_stroke(0, true, 0.09, 0.7, 1.0);
    for (i, x) in steps(0.25, 0.75, 20).enumerate() {
        layer.extend_stroke(0, x, 0.5);
        shot.frame("erase", 8 + i, &layer);
    }
    layer.end_stroke(0);
}

/// A correction joined to a radial selection and then taken out of it: the
/// part appears, is painted, and the mask loses exactly what it covers.
fn subtract(shot: &mut Shot) {
    let mut layer = lifted(MaskSource::Radial {
        centre: (0.5, 0.5),
        radii: (0.42, 0.34),
        angle: 0.0,
        feather: 0.35,
    });

    for i in 0..8 {
        shot.frame("subtract", i, &layer);
    }

    layer.push_part(MaskPart::painted("p2", Join::Subtract));
    for (i, x) in steps(0.3, 0.72, 22).enumerate() {
        if i == 0 {
            layer.begin_stroke(1, false, 0.1, 0.6, 1.0);
        }
        layer.extend_stroke(1, x, 0.46);
        shot.frame("subtract", 8 + i, &layer);
    }
    layer.end_stroke(1);

    // And off again, which is the half a stroke cannot do: a part is a thing
    // that can be switched off after the fact.
    for i in 0..8 {
        let mut without = layer.clone();
        without.remove_part(1);
        shot.frame("subtract", 30 + i, &without);
    }
}

/// The layer turned over, and back, holding each state long enough to read.
fn invert(shot: &mut Shot) {
    let mut layer = lifted(MaskSource::Radial {
        centre: (0.42, 0.52),
        radii: (0.3, 0.32),
        angle: 0.0,
        feather: 0.3,
    });

    for i in 0..24 {
        layer.invert = (i / 8) % 2 == 1;
        shot.frame("invert", i, &layer);
    }
}

/// The edge controls, swept: a feather opening up, then a dilation pushing the
/// boundary out and an erosion pulling it back.
fn shape(shot: &mut Shot) {
    let mut layer = lifted(MaskSource::Radial {
        centre: (0.5, 0.5),
        radii: (0.3, 0.3),
        angle: 0.0,
        feather: 0.02,
    });
    layer.base_mut().falloff = dr_pipeline::mask::Falloff::Smooth;

    for (i, f) in steps(0.0, 0.09, 18).enumerate() {
        layer.base_mut().feather = f;
        shot.frame("shape", i, &layer);
    }
    layer.base_mut().morphology = Morphology::Dilate;
    for (i, r) in steps(0.0, 0.06, 12).enumerate() {
        layer.base_mut().morph_radius = r;
        shot.frame("shape", 18 + i, &layer);
    }
    layer.base_mut().morphology = Morphology::Erode;
    for (i, r) in steps(0.0, 0.06, 12).enumerate() {
        layer.base_mut().morph_radius = r;
        shot.frame("shape", 30 + i, &layer);
    }
}

/// Everything one frame needs, so the mode functions read as what they do.
struct Shot<'a> {
    source: &'a DemosaicedImage,
    masks: &'a mut MaskPass,
    adjust: &'a mut AdjustPass,
    dir: &'a str,
}

impl Shot<'_> {
    fn frame(&mut self, mode: &str, index: usize, layer: &MaskLayer) {
        let mut stack = MaskStack::new();
        stack.push(layer.clone());

        let shader = compose_full(
            &ops::chain(),
            &Framing::new(),
            ColourSpace::Srgb,
            &stack,
            &SpotSet::new(),
            &[],
        );

        let array = self
            .masks
            .render(&stack, None, None, Some(self.source), W, H)
            .expect("rasterise");
        self.adjust
            .render_masked(self.source, &shader, W, H, Some(array))
            .expect("render");
        let rgba = self.adjust.export_pixels().expect("readback").0;

        write_ppm(
            &format!("{}/{mode}-{index:03}.ppm", self.dir),
            &to_rgb(&rgba),
            W,
            H,
        );
    }
}

/// `count` values from `from` to `to`, inclusive.
fn steps(from: f32, to: f32, count: usize) -> impl Iterator<Item = f32> {
    (0..count).map(move |i| from + (to - from) * i as f32 / (count.max(2) - 1) as f32)
}

/// A picture with somewhere obvious to put a mask: a graded sky, a ground
/// band, and a warm subject sitting on the join.
fn scene() -> Vec<u8> {
    let mut px = vec![0u8; (W * H * 4) as usize];
    for y in 0..H {
        for x in 0..W {
            let (fx, fy) = (x as f32 / W as f32, y as f32 / H as f32);
            let sky = [
                (60.0 + 90.0 * fy) as u8,
                (110.0 + 90.0 * fy) as u8,
                (190.0 + 50.0 * fy) as u8,
            ];
            let ground = [
                (70.0 + 40.0 * fx) as u8,
                (85.0 + 30.0 * fx) as u8,
                (60.0 + 20.0 * fx) as u8,
            ];
            let mut c = if fy > 0.62 { ground } else { sky };

            // The subject: an ellipse, warm, with a little internal structure
            // so a feathered edge has something to be soft against.
            let (dx, dy) = ((fx - 0.5) / 0.22, (fy - 0.52) / 0.3);
            if dx * dx + dy * dy < 1.0 {
                let shade = 0.75 + 0.25 * (fx * 40.0).sin() * (fy * 30.0).cos();
                c = [
                    (205.0 * shade) as u8,
                    (170.0 * shade) as u8,
                    (140.0 * shade) as u8,
                ];
            }

            let i = ((y * W + x) * 4) as usize;
            px[i..i + 4].copy_from_slice(&[c[0], c[1], c[2], 255]);
        }
    }
    px
}

fn to_rgb(rgba: &[u8]) -> Vec<u8> {
    rgba.chunks_exact(4)
        .flat_map(|p| [p[0], p[1], p[2]])
        .collect()
}

fn write_ppm(path: &str, rgb: &[u8], w: u32, h: u32) {
    use std::io::Write as _;
    let mut f = std::io::BufWriter::new(std::fs::File::create(path).expect("create"));
    write!(f, "P6\n{w} {h}\n255\n").expect("header");
    f.write_all(rgb).expect("body");
}