//! Setting the white balance from a point on the photograph.
#[cfg(test)]
use dr_decode::RawImage;
use dr_pipeline::Edit;

use crate::labels;

use super::session::DevelopSession;

impl DevelopSession {
    /// TRACES: FR-DEV-3 | FR-DEV-5
    /// Set the white balance from a point on the photograph.
    ///
    /// `x` and `y` are fractions of the *visible* image — the coordinates a
    /// click on the canvas arrives in — so a photographer inspecting a
    /// highlight at 4× samples the pixel they are actually looking at.
    ///
    /// **Nothing here knows it is white balance.** The colour goes to
    /// [`dr_pipeline::neutral`], which finds the operation that asked to be
    /// driven by a pixel and inverts its declared response; this side supplies
    /// the pixel and the undo step and nothing else. That is FR-DEV-3a's line
    /// in its awkward case: a picker genuinely needs to know how far a hundred
    /// units of temperature move red against blue, and that number is declared
    /// in the node's own file, so the interface must not be the thing that
    /// holds a second copy of it.
    ///
    /// **One step per sample**, and no step at all for a sample that could not
    /// be used — a point in the deep shadows has no balance in it to correct.
    /// `Edit::Action` never coalesces, so two clicks are two decisions
    /// however quickly they follow each other, which is what a photographer
    /// trying a wall and then a cloud expects to be able to undo one at a
    /// time.
    ///
    /// Returns whether the photograph moved.
    pub fn sample_neutral(&mut self, x: f32, y: f32) -> bool {
        let Some(sample) = self.sample_as_shot(x, y) else {
            return false;
        };
        if !dr_pipeline::neutral::neutralise(&mut self.graph, sample) {
            return false;
        }
        self.history
            .record(&self.graph, Edit::Action(labels::step::SAMPLED_NEUTRAL));
        true
    }

    /// The colour at a point in the space the white balance gains multiply:
    /// camera RGB with the camera's own balance on, linear, nothing else.
    ///
    /// **Measured where the operation acts, not where the photographer
    /// looks.** The white balance node runs first in the chain, on camera
    /// RGB, before the camera matrix and the view transform; the canvas
    /// shows the pixel after all three. The probe used to be read off a display
    /// render with the adjustments stripped, and the solve then treated an
    /// sRGB triple as if the gains multiplied it directly. On a JPEG the two
    /// spaces coincide, so it worked; on a raw file from any real body the
    /// matrix mixes the channels, and a slightly blue wall on a Canon 6D
    /// came back tint −77 with the whole frame green. This reads the
    /// camera-space tap a merge stitches from — the sensor's numbers after
    /// the lens warp — and puts the as-shot balance on itself, which is
    /// exactly the value the operation's gains are about to multiply.
    ///
    /// That also means nothing has to be stripped and restored: the tap
    /// runs no operations at all, and the display target is untouched, so
    /// a sample that found nothing usable leaves the canvas exactly as it
    /// was.
    ///
    /// The framing is the edit's own, exactly as for [`Self::render_original`]
    /// and for the same reason: `x` and `y` are fractions of what is on
    /// screen, and a probe rendered without the crop and the zoom would be
    /// answering about a different part of the photograph.
    ///
    /// **A patch, not a point.** The shader fetches the source at one
    /// position per output pixel — nearest, or four photosites blended — so
    /// a probe of the whole visible region rendered at 192px was not
    /// "averaging a neighbourhood into each pixel" as its comment claimed;
    /// it was one point sample of a noisy sensor, and two painted-white air
    /// conditioners on the same wall answered +37 and −50. Every eyedropper
    /// averages for exactly this reason: the photographer is pointing at a
    /// grey card, not at a photosite. So the tap is narrowed to the
    /// [`PATCH`] of the canvas around the click — a couple of percent of
    /// its width, square on screen — and rendered at [`PROBE_PX`] square
    /// with interpolation on, which puts a sample on every sensor pixel
    /// under the patch at any ordinary zoom. Those are averaged; a sample
    /// the tap marked void (outside the frame after the lens correction) or
    /// clipped is left out rather than allowed to pull the mean, and if
    /// fewer than half the patch survives there was nothing there to
    /// balance against. One small dispatch and a 64 KB readback on a click.
    pub(super) fn sample_as_shot(&mut self, x: f32, y: f32) -> Option<[f32; 3]> {
        /// Width of the patch as a fraction of what is on the canvas.
        const PATCH: f32 = 0.015;
        /// Side of the probe render, in pixels.
        const PROBE_PX: u32 = 64;

        // Square on screen: the height fraction follows the aspect of the
        // visible region, which is the crop's shape times the view's.
        let (sw, sh) = self.demosaiced.size();
        let (cw, ch) = self.graph.output_size(sw, sh);
        let view = self.graph.framing().view();
        let aspect = (cw as f32 * view.width) / (ch as f32 * view.height).max(f32::EPSILON);
        let (pw, ph) = (PATCH, PATCH * aspect);
        let patch = dr_pipeline::CropRect {
            x: x.clamp(0.0, 1.0) - pw * 0.5,
            y: y.clamp(0.0, 1.0) - ph * 0.5,
            width: pw,
            height: ph,
        };

        let shader = self.graph.compose_camera_probe(patch);
        let rendered = self
            .adjust
            .render_camera_linear(&self.demosaiced, &shader, PROBE_PX, PROBE_PX)
            .map(|_| ());
        let (rgba, _, _) = rendered
            .and_then(|()| self.adjust.read_camera_linear())
            .inspect_err(|e| log::warn!("could not read a neutral off the frame: {e}"))
            .ok()?;

        let mut sum = [0.0f32; 3];
        let mut kept = 0usize;
        let mut seen = 0usize;
        for pixel in rgba.chunks_exact(4) {
            seen += 1;
            // The tap marks a pixel the lens correction pulled in from
            // outside the frame with alpha 0. There is nothing there to
            // balance against.
            if pixel[3] < 0.5 {
                continue;
            }
            // Nor in a clipped one. A blown sky reads as sensor white, and
            // sensor white with the as-shot balance on is strongly magenta —
            // a solve over it drives tint to its stop for a pixel that, on
            // the canvas, the shader has already desaturated to neutral. The
            // same threshold the shader fades from, so what is refused here
            // is what it would have hidden there.
            if pixel[..3].iter().any(|c| *c >= dr_pipeline::CLIP_ONSET) {
                continue;
            }
            for (acc, c) in sum.iter_mut().zip(pixel) {
                *acc += c;
            }
            kept += 1;
        }
        if kept == 0 || kept * 2 < seen {
            return None;
        }

        // The tap is the sensor's numbers with the profile filled neutral;
        // the operation multiplies them *after* the camera's own balance, so
        // that goes on here and the solve sees what the gains will see.
        let wb = self.demosaiced.as_shot_wb();
        let n = kept as f32;
        Some([sum[0] / n * wb[0], sum[1] / n * wb[1], sum[2] / n * wb[2]])
    }
}

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

    /// TRACES: FR-DEV-3 | FR-DEV-5
    /// Sampling something that is already neutral corrects nothing, and says
    /// so by leaving the stack alone.
    ///
    /// The failure this guards is a picker that lands a ten-thousandth off
    /// zero: the photograph would come back marked modified, an undo step
    /// would appear for a correction of nothing, and the sidecar would gain a
    /// temperature the photographer never chose. The solve rounds to the
    /// precision the control is drawn at, which is what makes "no correction"
    /// representable at all.
    #[test]
    fn sampling_a_grey_that_is_already_grey_leaves_the_photograph_alone() {
        let Some(ctx) = headless() else { return };
        let (mut session, _) = grey_session(&ctx);
        let steps = session.history_rows().len();

        assert!(
            session.sample_neutral(0.5, 0.5),
            "a flat grey frame is a usable sample"
        );
        assert!(
            session.is_neutral(),
            "there was nothing to correct, so nothing was corrected"
        );
        assert_eq!(
            session.history_rows().len(),
            steps,
            "and a correction of nothing is not a step"
        );
    }

    /// TRACES: FR-DEV-3
    /// The whole point of the picker, measured where the photographer sees
    /// it: a cast grey on a *raw* frame, sampled, renders grey.
    ///
    /// On a raw frame and not a JPEG, because that is where it was wrong. The
    /// white balance gains multiply camera RGB, before the body's matrix
    /// turns it into sRGB; the probe was read *after* the matrix, and the
    /// solve treated the two as the same space. On a body whose matrix mixes
    /// the channels as much as a Canon's does, a slightly blue wall came back
    /// tint −77 and the whole frame went green. A JPEG carries an identity
    /// matrix, so the same test on one passed while the picker was broken.
    #[test]
    fn sampling_a_cast_grey_on_a_raw_frame_renders_it_grey() {
        let Some(ctx) = headless() else { return };

        // A Canon EOS 6D's D65 matrix (rows summing to one, as
        // `neutral_stays_neutral_through_the_colour_matrix` requires) and a
        // typical as-shot balance for it.
        let cam_to_srgb = [
            1.9125, -1.0587, 0.1461, //
            -0.2249, 1.6466, -0.4217, //
            0.0099, -0.5093, 1.4994,
        ];
        let as_shot = [1.9, 1.0, 1.7];
        // What the wall should look like once the camera's own balance is on:
        // a warm cast, a little over half a stop between red and blue.
        let balanced = [0.30f32, 0.25, 0.20];
        let sensor: Vec<u16> = (0..3)
            .map(|c| (balanced[c] / as_shot[c] * 65535.0).round() as u16)
            .collect();
        let size = 64u32;
        let raw = RawImage {
            width: size,
            height: size,
            data: sensor.repeat((size * size) as usize),
            cfa_pattern: dr_decode::CfaPattern::Rggb,
            black_level: [0; 4],
            white_level: 65535,
            wb_coeffs: [as_shot[0], as_shot[1], as_shot[2], 0.0],
            color_matrix: Some(cam_to_srgb),
            samples_per_pixel: 3,
            profile: None,
            profile_tables: None,
            make: String::new(),
            model: String::new(),
            crop: dr_decode::CropRect {
                x: 0,
                y: 0,
                width: size,
                height: size,
            },
        };
        let mut session =
            DevelopSession::open(&ctx, &raw, dr_types::Orientation::NORMAL).expect("session");

        let at = ((size / 2) * size + size / 2) as usize * 4;
        let before = read_back(&ctx, &session.render(size, size).expect("render"));
        let cast = |px: &[u8]| px.iter().max().unwrap() - px.iter().min().unwrap();
        assert!(
            cast(&before[at..at + 3]) > 20,
            "the premise: the wall renders with a cast, {:?}",
            &before[at..at + 3]
        );

        assert!(
            session.sample_neutral(0.5, 0.5),
            "a mid-grey is a usable sample"
        );

        let after = read_back(&ctx, &session.render(size, size).expect("render"));
        let px = &after[at..at + 3];
        assert!(
            cast(px) <= 3,
            "the sampled point should render neutral, got {px:?} with {:?}",
            session
                .rows()
                .iter()
                .filter(|r| r.value != r.default_value)
                .map(|r| (r.param_label.to_string(), r.value))
                .collect::<Vec<_>>()
        );
    }

    /// TRACES: FR-DEV-3
    /// The picker reads a patch, not a photosite.
    ///
    /// A frame whose pixels alternate warm and cool grey, averaging to a
    /// neutral: a point sample lands on one or the other and swings the
    /// controls hard one way, which is what two white boxes on the same wall
    /// answering +37 and −50 looked like. Averaged, there is nothing to
    /// correct, and the graph says so.
    #[test]
    fn sampling_averages_a_patch_rather_than_reading_one_photosite() {
        let Some(ctx) = headless() else { return };
        // Large enough that the patch — a couple of percent of the frame —
        // holds many sensor pixels; on a 64px frame it would hold one, and
        // the test would be asserting about interpolation instead.
        let size = 1536u32;
        let warm = [0.30f32, 0.25, 0.20];
        let cool = [0.20f32, 0.25, 0.30];
        let mut data = Vec::with_capacity((size * size * 3) as usize);
        for i in 0..(size * size) as usize {
            let p = if i % 2 == 0 { warm } else { cool };
            data.extend(p.iter().map(|c| (c * 65535.0).round() as u16));
        }
        let raw = RawImage {
            width: size,
            height: size,
            data,
            cfa_pattern: dr_decode::CfaPattern::Rggb,
            black_level: [0; 4],
            white_level: 65535,
            wb_coeffs: [1.0, 1.0, 1.0, 0.0],
            color_matrix: None,
            samples_per_pixel: 3,
            profile: None,
            profile_tables: None,
            make: String::new(),
            model: String::new(),
            crop: dr_decode::CropRect {
                x: 0,
                y: 0,
                width: size,
                height: size,
            },
        };
        let mut session =
            DevelopSession::open(&ctx, &raw, dr_types::Orientation::NORMAL).expect("session");

        assert!(
            session.sample_neutral(0.5, 0.5),
            "a mid-grey patch is usable"
        );
        let moved: Vec<_> = session
            .rows()
            .iter()
            .filter(|r| r.value != r.default_value)
            .map(|r| (r.param_label.to_string(), r.value))
            .collect();
        assert!(
            moved.iter().all(|(_, v)| v.abs() <= 2.0),
            "the patch averages neutral, so nothing should move far: {moved:?}"
        );
    }

    /// TRACES: FR-DEV-3
    /// A blown highlight is refused, the way black is.
    ///
    /// Sensor white is not a colour: every channel stopped counting, so the
    /// ratio between them is the as-shot multipliers and nothing about the
    /// scene. Sampling the overcast sky on a Canon 6D frame drove tint to
    /// -100 and temperature to -15 for a patch the canvas showed as pure
    /// white, which is the picker being wrong rather than the point being a
    /// poor choice. Refused, nothing moves and no step is taken.
    #[test]
    fn sampling_a_blown_highlight_moves_nothing() {
        let Some(ctx) = headless() else { return };
        let size = 16u32;
        let raw = RawImage {
            width: size,
            height: size,
            data: vec![65535; (size * size * 3) as usize],
            cfa_pattern: dr_decode::CfaPattern::Rggb,
            black_level: [0; 4],
            white_level: 65535,
            wb_coeffs: [1.9, 1.0, 1.7, 0.0],
            color_matrix: None,
            samples_per_pixel: 3,
            profile: None,
            profile_tables: None,
            make: String::new(),
            model: String::new(),
            crop: dr_decode::CropRect {
                x: 0,
                y: 0,
                width: size,
                height: size,
            },
        };
        let mut session =
            DevelopSession::open(&ctx, &raw, dr_types::Orientation::NORMAL).expect("session");
        let steps = session.history_rows().len();

        assert!(
            !session.sample_neutral(0.5, 0.5),
            "a clipped photosite has no balance in it"
        );
        assert!(session.is_neutral(), "and so nothing was corrected");
        assert_eq!(session.history_rows().len(), steps);
    }
}