//! Turning a stock into something a shader can run.
//!
//! # The decomposition
//!
//! A physically-honest film simulation looks like it needs a spectral
//! integration per pixel, and vkdt's does exactly that. It does not have to,
//! and the reason is worth writing down because it is what makes this cheap
//! enough to run on a phone:
//!
//! 1. **Exposure is a 3×3 matrix.** A layer's exposure is
//!    `∫ S(λ)·L(λ) dλ`, and the scene spectrum `L` reconstructed from an sRGB
//!    triple is *linear* in that triple — that is what a spectral basis is. So
//!    the whole integral collapses into nine numbers, computed once, exactly.
//!    No approximation is involved.
//!
//! 2. **The characteristic curve is three 1D functions.** Sampled exactly, at
//!    [`crate::profile::CURVE_SAMPLES`].
//!
//! 3. **Everything after that is a function of three densities.** The dye
//!    transmittance, the print exposure through the negative, the paper's own
//!    curves and dyes, the viewing illuminant, the adaptation — all of it takes
//!    three numbers in and gives three numbers out. So it bakes into one small
//!    3D lookup, and the per-pixel cost is a matrix multiply, three curve taps
//!    and one texture fetch.
//!
//! Splitting 2 from 3 rather than baking a single LUT over exposure is
//! deliberate and measured: the curve carries all of the sharp shape and the
//! dye mixing is smooth, so putting the curve in the 3D LUT would force it
//! three times larger for the same error.

use crate::profile::Profile;
use crate::spectrum::{illuminant, Spectrum, Viewing};
use crate::tables::{SPECTRUM, SRGB_BASIS};

/// The mid-grey a photographic exposure is reckoned from.
///
/// 18.4% rather than 18%: it is the value the upstream profiles are calibrated
/// against, and a profile calibrated at one grey and rendered at another is
/// off by a fraction of a stop everywhere.
pub const MID_GREY: f32 = 0.184;

/// The edge length of the baked density lookup.
///
/// 32 holds the worst-case interpolation error to about 0.003 in linear sRGB,
/// which is below one 8-bit code value, in 384 kB. Doubling it buys a factor
/// of four in error for eight times the memory, and there is nothing to spend
/// that on: the error is already under what the output can represent.
pub const LUT_SIZE: usize = 32;

/// What to develop, and how.
pub struct Recipe<'a> {
    /// The stock the picture was taken on.
    pub film: &'a Profile,
    /// The paper it is printed on. `None` views the film directly, which is
    /// what a reversal stock wants and what makes a negative come out orange
    /// and inverted — that being what a negative actually looks like.
    pub print: Option<&'a Profile>,
    /// Camera exposure, in stops.
    pub exposure_ev: f32,
    /// Enlarger exposure, in stops. Ignored without a `print`.
    pub print_exposure_ev: f32,
    /// TRACES: FR-DEV-3f
    /// Development, in stops of push. Positive develops longer.
    ///
    /// Ignored by a stock measured at one process, of which there are many —
    /// see [`crate::profile::Profile::curves_at_push`], which returns the one
    /// measured curve rather than inventing a pushed one.
    pub push_stops: f32,
}

impl<'a> Recipe<'a> {
    /// The straightforward reading of a stock: reversal viewed directly,
    /// negative printed on the paper its datasheet names.
    pub fn new(film: &'a Profile, print: Option<&'a Profile>) -> Self {
        Self {
            film,
            print,
            exposure_ev: 0.0,
            print_exposure_ev: 0.0,
            push_stops: 0.0,
        }
    }
}

/// A recipe reduced to three tables.
///
/// Plain `f32` with a documented layout, and no notion of a texture: what to
/// bind this to is dr-gpu's decision, and keeping it out of here is what lets
/// the whole model be tested on the CPU.
#[derive(Debug, Clone)]
pub struct Baked {
    /// Linear sRGB to the three layers' log₁₀ exposure, before the log — row
    /// `l`, column `c` is layer `l`'s response to sRGB channel `c`.
    pub exposure_matrix: [[f32; 3]; 3],
    /// The characteristic curves, `CURVE_SAMPLES` samples per layer, uniform
    /// over `[curve_log_min, curve_log_max]`.
    pub curves: Vec<[f32; 3]>,
    pub curve_log_min: f32,
    pub curve_log_max: f32,
    /// Density to linear sRGB, `LUT_SIZE³` entries uniform over
    /// `[0, density_max]` on each axis.
    ///
    /// **The red axis varies fastest**, then green, then blue — that is,
    /// `lut[(b * size + g) * size + r]`. Stated because it is not the order
    /// this loop reads most naturally, and it is not arbitrary: it is the
    /// order a 3D texture upload expects, so the consumer can hand the slice
    /// straight to the driver. Filling it the other way round renders a
    /// picture with red and blue transposed, which looks like a plausible
    /// photograph of the wrong colour.
    pub lut: Vec<[f32; 3]>,
    pub density_max: f32,
    pub lut_size: usize,
}

impl Baked {
    /// Look a colour up the way the shader will, for tests and for previews.
    pub fn apply(&self, rgb: [f32; 3]) -> [f32; 3] {
        let mut log_exposure = [0.0f32; 3];
        for (l, slot) in log_exposure.iter_mut().enumerate() {
            let m = self.exposure_matrix[l];
            let e = m[0] * rgb[0] + m[1] * rgb[1] + m[2] * rgb[2];
            *slot = (e.max(0.0) + 1e-10).log10();
        }
        self.sample_lut(self.sample_curves(log_exposure))
    }

    fn sample_curves(&self, log_exposure: [f32; 3]) -> [f32; 3] {
        let last = self.curves.len() - 1;
        let span = self.curve_log_max - self.curve_log_min;
        let mut out = [0.0f32; 3];
        for (c, slot) in out.iter_mut().enumerate() {
            let t = ((log_exposure[c] - self.curve_log_min) / span).clamp(0.0, 1.0) * last as f32;
            let i = (t.floor() as usize).min(last - 1);
            let f = t - i as f32;
            *slot = self.curves[i][c] * (1.0 - f) + self.curves[i + 1][c] * f;
        }
        out
    }

    fn sample_lut(&self, density: [f32; 3]) -> [f32; 3] {
        let n = self.lut_size;
        let mut base = [0usize; 3];
        let mut frac = [0f32; 3];
        for c in 0..3 {
            let t = (density[c] / self.density_max).clamp(0.0, 1.0) * (n - 1) as f32;
            base[c] = (t.floor() as usize).min(n - 2);
            frac[c] = t - base[c] as f32;
        }
        let mut out = [0.0f32; 3];
        for dx in 0..2 {
            for dy in 0..2 {
                for dz in 0..2 {
                    let w = if dx == 0 { 1.0 - frac[0] } else { frac[0] }
                        * if dy == 0 { 1.0 - frac[1] } else { frac[1] }
                        * if dz == 0 { 1.0 - frac[2] } else { frac[2] };
                    let e = self.lut[((base[2] + dz) * n + base[1] + dy) * n + base[0] + dx];
                    for c in 0..3 {
                        out[c] += w * e[c];
                    }
                }
            }
        }
        out
    }
}

/// Linear sRGB to the three layers' exposure, mid-grey normalised.
///
/// Normalised on the *green* layer alone, one shared scalar for all three.
/// Doing it per layer is tempting and wrong: it would silently flatten the
/// film's own channel balance, which is a large part of what distinguishes one
/// stock from another. Where the balance genuinely has to come out — printing a
/// negative — it is [`print_balance`]'s job, which is also where it belongs
/// physically.
pub fn exposure_matrix(film: &Profile) -> [[f32; 3]; 3] {
    let reference = illuminant(&film.reference_illuminant);
    let sensitivity = film.sensitivity();

    let mut m = [[0.0f32; 3]; 3];
    let mut mid_grey = [0.0f32; 3];
    for i in 0..SPECTRUM {
        for layer in 0..3 {
            let s = sensitivity[i][layer] * reference[i];
            mid_grey[layer] += s * MID_GREY;
            for channel in 0..3 {
                m[layer][channel] += s * SRGB_BASIS[i][channel];
            }
        }
    }

    let scale = 1.0 / mid_grey[1];
    for row in &mut m {
        for v in row.iter_mut() {
            *v *= scale;
        }
    }
    m
}

/// The enlarger head's filtration, solved rather than dialled.
///
/// Returns the per-layer log exposure offsets that make a mid-grey scene print
/// as a neutral mid-grey. This is also where a colour negative's orange mask
/// goes: the mask is a fixed density, so balancing mid-grey to neutral cancels
/// it — which is why a printed negative looks like a photograph while a scanned
/// one looks orange.
fn print_balance(
    film: &Profile,
    paper: &Profile,
    exposure_ev: f32,
    print_exposure_ev: f32,
) -> [f32; 3] {
    let matrix = exposure_matrix(film);
    let scene = MID_GREY * 2f32.powf(exposure_ev);
    let mut log_exposure = [0.0f32; 3];
    for (l, slot) in log_exposure.iter_mut().enumerate() {
        let m = matrix[l];
        *slot = ((m[0] + m[1] + m[2]) * scene + 1e-10).log10();
    }
    let mid_raw = paper_exposure(film, paper, film.density_at(log_exposure));

    // Where on the paper's curve mid-grey belongs: the density that reflects
    // 18%, read off the average of the three curves. Averaged because the
    // point of the balance is that the three end up at the same place.
    let target_density = -MID_GREY.log10() - mean(&paper.base_density);
    let target = invert_mean_curve(paper, target_density);

    let mut offsets = [0.0f32; 3];
    for (l, slot) in offsets.iter_mut().enumerate() {
        *slot = target - (mid_raw[l] + 1e-10).log10() + print_exposure_ev * 2f32.log10();
    }
    offsets
}

/// The paper's three layer exposures, printing through a negative at these
/// densities.
///
/// The one genuinely spectral step left in the chain — the negative's
/// transmittance is `10^-D`, so this is not a matrix and cannot be made into
/// one. It takes exactly three numbers in, which is what lets the whole
/// negative-and-print chain still bake into a 3D lookup.
fn paper_exposure(film: &Profile, paper: &Profile, density: [f32; 3]) -> [f32; 3] {
    let enlarger = illuminant(&paper.reference_illuminant);
    let sensitivity = paper.sensitivity();
    let transmittance = film.transmittance(density);

    let mut raw = [0.0f32; 3];
    for i in 0..SPECTRUM {
        let light = transmittance[i] * enlarger[i];
        for layer in 0..3 {
            raw[layer] += light * sensitivity[i][layer];
        }
    }
    raw
}

/// Bake a recipe into the tables a shader runs.
pub fn bake(recipe: &Recipe) -> Baked {
    let film = recipe.film;
    let mut matrix = exposure_matrix(film);
    // Camera exposure rides in the matrix rather than in the shader: it is a
    // scalar on a linear quantity, and folding it in here costs nothing and
    // keeps the per-pixel work identical whether or not it has been moved.
    let gain = 2f32.powf(recipe.exposure_ev);
    for row in &mut matrix {
        for v in row.iter_mut() {
            *v *= gain;
        }
    }

    // TRACES: FR-DEV-3f
    // Developed to the requested push before anything else reads the curves:
    // the density ceiling, the print balance and the grain all depend on how
    // far this film was taken, and a push that only reached one of them would
    // be a contrast change wearing a push's name.
    let curves = film.curves_at_push(recipe.push_stops);
    let density_max = curves
        .iter()
        .flat_map(|row| row.iter())
        .fold(0.0f32, |a, &b| a.max(b))
        .max(1e-3);

    let viewing = match recipe.print {
        Some(paper) => Viewing::new(&paper.viewing_illuminant),
        None => Viewing::new(&film.viewing_illuminant),
    };
    let balance = recipe
        .print
        .map(|paper| print_balance(film, paper, recipe.exposure_ev, recipe.print_exposure_ev));

    let n = LUT_SIZE;
    let mut lut = Vec::with_capacity(n * n * n);
    // Blue outermost and red innermost, so the red axis varies fastest. See
    // `Baked::lut`: this is the layout a 3D texture upload wants, and getting
    // it backwards transposes red and blue in the finished picture.
    for b in 0..n {
        for g in 0..n {
            for r in 0..n {
                let density = [
                    density_max * r as f32 / (n - 1) as f32,
                    density_max * g as f32 / (n - 1) as f32,
                    density_max * b as f32 / (n - 1) as f32,
                ];
                lut.push(match recipe.print.zip(balance) {
                    Some((paper, offsets)) => {
                        let raw = paper_exposure(film, paper, density);
                        let mut log_exposure = [0.0f32; 3];
                        for (l, slot) in log_exposure.iter_mut().enumerate() {
                            *slot = (raw[l] + 1e-10).log10() + offsets[l];
                        }
                        let paper_density = paper.density_at(log_exposure);
                        viewing.to_srgb(&paper.transmittance(paper_density))
                    }
                    None => viewing.to_srgb(&film.transmittance(density)),
                });
            }
        }
    }

    Baked {
        exposure_matrix: matrix,
        curves,
        curve_log_min: film.log_exposure_min,
        curve_log_max: film.log_exposure_max,
        lut,
        density_max,
        lut_size: n,
    }
}

fn mean(s: &Spectrum) -> f32 {
    s.iter().sum::<f32>() / SPECTRUM as f32
}

/// The log exposure at which the paper's average curve reaches `density`.
fn invert_mean_curve(paper: &Profile, density: f32) -> f32 {
    let curves = &paper.density_curves;
    let last = curves.len() - 1;
    let span = paper.log_exposure_max - paper.log_exposure_min;
    let at = |i: usize| (curves[i][0] + curves[i][1] + curves[i][2]) / 3.0;
    let log_at = |i: usize| paper.log_exposure_min + span * i as f32 / last as f32;

    // A paper is negative-working, so its mean curve rises. Walk it rather
    // than binary-search: 256 samples is nothing, and a linear scan is correct
    // even where the curve is flat, which a bisection is not.
    let ascending = at(last) >= at(0);
    for i in 0..last {
        let (lo, hi) = (at(i), at(i + 1));
        let brackets = if ascending {
            lo <= density && density <= hi
        } else {
            hi <= density && density <= lo
        };
        if brackets && (hi - lo).abs() > f32::EPSILON {
            let f = (density - lo) / (hi - lo);
            return log_at(i) + (log_at(i + 1) - log_at(i)) * f;
        }
    }
    // Off the end of the curve: the nearest end is the honest answer, and it
    // keeps a badly-scaled contributed profile from producing a NaN that would
    // propagate silently through the whole LUT.
    if (density <= at(0)) == ascending {
        paper.log_exposure_min
    } else {
        paper.log_exposure_max
    }
}

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

    fn profile(yaml: &str) -> Profile {
        Profile::parse(yaml).unwrap()
    }

    fn portra() -> Profile {
        profile(include_str!("../profiles/kodak_portra_400.yaml"))
    }

    fn endura() -> Profile {
        profile(include_str!("../profiles/kodak_portra_endura.yaml"))
    }

    fn kodachrome() -> Profile {
        profile(include_str!("../profiles/kodak_kodachrome_64.yaml"))
    }

    fn spread(rgb: [f32; 3]) -> f32 {
        rgb.iter().cloned().fold(f32::MIN, f32::max) - rgb.iter().cloned().fold(f32::MAX, f32::min)
    }

    #[test]
    fn the_exposure_matrix_is_diagonally_dominant() {
        // Each layer must respond most strongly to its own primary. A matrix
        // that failed this would mean the sensitivity table had been pasted in
        // the wrong channel order, which produces a picture that renders
        // perfectly and has its colours swapped.
        for film in [portra(), kodachrome()] {
            let m = exposure_matrix(&film);
            for layer in 0..3 {
                for channel in 0..3 {
                    if channel != layer {
                        assert!(
                            m[layer][layer] > m[layer][channel] * 3.0,
                            "{}: layer {layer} responds to channel {channel} too strongly: {m:?}",
                            film.stock
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn a_reversal_stock_renders_a_positive() {
        let film = kodachrome();
        let baked = bake(&Recipe::new(&film, None));
        let shadow = baked.apply([0.02; 3]);
        let mid = baked.apply([MID_GREY; 3]);
        let highlight = baked.apply([0.8; 3]);
        assert!(
            shadow[1] < mid[1] && mid[1] < highlight[1],
            "{shadow:?} {mid:?} {highlight:?}"
        );
    }

    #[test]
    fn a_scanned_negative_is_inverted_and_orange() {
        // Not a defect: it is what a negative looks like, and rendering it any
        // other way would mean the print stage was silently applied.
        let film = portra();
        let baked = bake(&Recipe::new(&film, None));
        let shadow = baked.apply([0.02; 3]);
        let highlight = baked.apply([0.8; 3]);
        assert!(
            shadow[1] > highlight[1],
            "not inverted: {shadow:?} -> {highlight:?}"
        );
        let mid = baked.apply([MID_GREY; 3]);
        assert!(mid[0] > mid[2] * 4.0, "no orange mask: {mid:?}");
    }

    #[test]
    fn printing_a_negative_restores_the_picture() {
        // The property the whole print stage exists for: through the paper,
        // the same negative is the right way up and neutral again.
        let film = portra();
        let paper = endura();
        let baked = bake(&Recipe::new(&film, Some(&paper)));

        let shadow = baked.apply([0.02; 3]);
        let mid = baked.apply([MID_GREY; 3]);
        let highlight = baked.apply([0.8; 3]);
        assert!(
            shadow[1] < mid[1] && mid[1] < highlight[1],
            "print is not a positive: {shadow:?} {mid:?} {highlight:?}"
        );
        for grey in [shadow, mid, highlight] {
            assert!(
                spread(grey) < 0.06,
                "print of a neutral is not neutral: {grey:?}"
            );
        }
    }

    #[test]
    fn exposure_moves_the_print_the_way_it_moves_a_photograph() {
        let film = portra();
        let paper = endura();
        let brighter = bake(&Recipe {
            exposure_ev: 1.0,
            ..Recipe::new(&film, Some(&paper))
        });
        let base = bake(&Recipe::new(&film, Some(&paper)));
        assert!(brighter.apply([MID_GREY; 3])[1] > base.apply([MID_GREY; 3])[1]);
    }

    #[test]
    fn the_lut_holds_its_error_under_a_code_value() {
        // The claim LUT_SIZE is chosen on. Compared against the same chain
        // evaluated exactly, so it measures interpolation error and nothing
        // else.
        let film = kodachrome();
        let baked = bake(&Recipe::new(&film, None));
        let viewing = Viewing::new(&film.viewing_illuminant);

        let mut worst = 0.0f32;
        for i in 0..40 {
            for j in 0..40 {
                let rgb = [
                    i as f32 / 39.0,
                    j as f32 / 39.0,
                    ((i + j) % 40) as f32 / 39.0,
                ];
                let mut log_exposure = [0.0f32; 3];
                for (l, slot) in log_exposure.iter_mut().enumerate() {
                    let m = baked.exposure_matrix[l];
                    *slot =
                        ((m[0] * rgb[0] + m[1] * rgb[1] + m[2] * rgb[2]).max(0.0) + 1e-10).log10();
                }
                let exact = viewing.to_srgb(&film.transmittance(film.density_at(log_exposure)));
                let approx = baked.apply(rgb);
                for c in 0..3 {
                    worst = worst.max((exact[c] - approx[c]).abs());
                }
            }
        }
        assert!(
            worst < 1.0 / 255.0,
            "worst LUT error {worst} exceeds one code value"
        );
    }

    #[test]
    fn the_lut_stores_red_along_its_fastest_axis() {
        // The layout a 3D texture upload expects, and the one bug this whole
        // decomposition is most exposed to: fill it the other way round and
        // the picture comes back with red and blue transposed -- entirely
        // plausible-looking, and wrong. Asserted here rather than only in the
        // end-to-end GPU test, because that one needs a device and this one
        // does not.
        let film = kodachrome();
        let baked = bake(&Recipe::new(&film, None));
        let n = baked.lut_size;

        // Step one along each axis from the origin, and check that the entry
        // found is the one the *density* moved along that axis should give.
        let viewing = Viewing::new(&film.viewing_illuminant);
        let step = baked.density_max / (n - 1) as f32;
        for (axis, offset) in [(0usize, 1usize), (1, n), (2, n * n)] {
            let mut density = [0.0f32; 3];
            density[axis] = step;
            let expected = viewing.to_srgb(&film.transmittance(density));
            let stored = baked.lut[offset];
            for c in 0..3 {
                assert!(
                    (stored[c] - expected[c]).abs() < 1e-4,
                    "axis {axis} is not at stride {offset}: stored {stored:?}, \
                     the density one step along that axis gives {expected:?}"
                );
            }
        }
    }

    #[test]
    fn the_lut_is_the_size_it_says_it_is() {
        let film = kodachrome();
        let baked = bake(&Recipe::new(&film, None));
        assert_eq!(baked.lut.len(), LUT_SIZE * LUT_SIZE * LUT_SIZE);
        assert_eq!(baked.curves.len(), CURVE_SAMPLES);
    }
}