//! TRACES: FR-DEV-3g
//! How noisy each photosite is: the network is told, not left to guess
//! (denoise.md §3.3).
//!
//! The model is `σ² = S·x + O + row² + col²` per photosite, `x` the signal
//! normalised black-to-white the way the demosaic normalises it. Three
//! sources, best first:
//!
//! 1. **A measured table** for the body ([`Source::Table`]) — the Canon EOS 6D
//!    today, from the library's own frames.
//! 2. **The DNG's `NoiseProfile`** ([`Source::DngProfile`]) — what Adobe's
//!    converter measured for the body at that ISO.
//! 3. **The frame itself** ([`Source::Measured`]) — read, row and column
//!    noise from its masked border, which is a dark frame taken in the same
//!    instant, and only the shot gain estimated, from the quietest flat
//!    patches. Checked against the 6D's table on 130 frames: within ±10 % at
//!    ISO 1000 and above, scattered below; the network loses under 0.3 dB for
//!    a σ off by 15–20 %, and over-estimating costs half what
//!    under-estimating does, so the estimate leans high.
//!
//! Row and column noise come from the masked border whenever the frame has
//! one, whatever the source of the rest.

use dr_decode::{CfaPattern, RawImage};
use serde::Deserialize;

/// Where a frame's noise figures came from, for develop to say.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Source {
    Table,
    DngProfile,
    Measured,
}

impl Source {
    pub fn label(self) -> &'static str {
        match self {
            Source::Table => "measured for this camera",
            Source::DngProfile => "from the DNG's noise profile",
            Source::Measured => "estimated from this photograph",
        }
    }
}

/// Per-photosite noise in the frame's own normalisation (black 0, white 1).
#[derive(Clone, Debug, PartialEq)]
pub struct NoiseModel {
    /// Shot gain per colour, R G B.
    pub s: [f32; 3],
    /// Read variance per colour, R G B.
    pub o: [f32; 3],
    /// Standard deviation shared by a whole row, and by a whole column.
    pub row: f32,
    pub col: f32,
    pub source: Source,
}

impl NoiseModel {
    /// σ for a photosite of colour `c` (0 R, 1 G, 2 B) reading `x`.
    #[inline]
    pub fn sigma(&self, c: usize, x: f32) -> f32 {
        (self.s[c] * x.max(0.0) + self.o[c] + self.row * self.row + self.col * self.col).sqrt()
    }

    /// The same figures scaled for the Amount the spec describes (§3.3):
    /// above 1 tells the network there is more noise than there is.
    pub fn scaled(&self, amount: f32) -> NoiseModel {
        let a2 = amount * amount;
        NoiseModel {
            s: self.s.map(|v| v * a2),
            o: self.o.map(|v| v * a2),
            row: self.row * amount,
            col: self.col * amount,
            source: self.source,
        }
    }
}

/// The frame's noise, from the best source it has.
///
/// `bytes` is the file (for a DNG's `NoiseProfile`), `iso` its EXIF ISO.
/// `None` only for a frame with no masked border, no profile and no table
/// that is also too dark or too busy to measure.
pub fn for_frame(raw: &RawImage, bytes: &[u8], iso: Option<u32>) -> Option<NoiseModel> {
    for_frame_with(raw, dr_decode::noise_profile(bytes).as_deref(), iso)
}

/// [`for_frame`], given the file's `NoiseProfile` already read
/// ([`dr_decode::noise_profile`]) rather than the file, for a caller that
/// keeps the header's answer and not the bytes.
pub fn for_frame_with(
    raw: &RawImage,
    profile: Option<&[(f32, f32)]>,
    iso: Option<u32>,
) -> Option<NoiseModel> {
    let dark = dark_border(raw);
    let mut model = iso
        .and_then(|iso| from_table(raw, iso))
        .or_else(|| profile.and_then(|p| from_dng_profile(raw, p)))
        .or_else(|| measured(raw, dark.as_ref()))?;
    if let Some(d) = dark {
        // The border saw this exposure's row and column noise directly.
        if model.source != Source::Table {
            model.row = d.row;
            model.col = d.col;
        }
    }
    Some(model)
}

#[derive(Deserialize)]
struct Table {
    make: String,
    model: String,
    rows: Vec<TableRow>,
}

#[derive(Deserialize)]
struct TableRow {
    iso: u32,
    s_dn: [f32; 4],
    o_dn: [f32; 4],
    row_dn: f32,
    col_dn: f32,
}

const TABLES: &[&str] = &[include_str!("../tables/canon-eos-6d.yaml")];

/// The body's measured table at the nearest ISO it holds, converted from DN
/// to this frame's normalisation.
pub fn from_table(raw: &RawImage, iso: u32) -> Option<NoiseModel> {
    let table = TABLES.iter().find_map(|t| {
        let t: Table = serde_norway::from_str(t).ok()?;
        (t.make.eq_ignore_ascii_case(&raw.make) && t.model.eq_ignore_ascii_case(&raw.model))
            .then_some(t)
    })?;
    let row = table.rows.iter().min_by(|a, b| {
        let d = |r: &TableRow| ((r.iso as f32).ln() - (iso as f32).ln()).abs();
        d(a).total_cmp(&d(b))
    })?;
    let span = span(raw);
    // RGGB positions → colours: the greens share.
    let s = [row.s_dn[0], 0.5 * (row.s_dn[1] + row.s_dn[2]), row.s_dn[3]].map(|v| v / span);
    let o =
        [row.o_dn[0], 0.5 * (row.o_dn[1] + row.o_dn[2]), row.o_dn[3]].map(|v| v / (span * span));
    Some(NoiseModel {
        s,
        o,
        row: row.row_dn / span,
        col: row.col_dn / span,
        source: Source::Table,
    })
}

/// A DNG's `NoiseProfile`: one pair for every plane, or one per colour plane
/// (R, G, B for a Bayer DNG), already in the file's black-to-white units —
/// which are the units `dr-decode` normalises by.
pub fn from_dng_profile(raw: &RawImage, pairs: &[(f32, f32)]) -> Option<NoiseModel> {
    if raw.cfa_pattern.is_xtrans() || raw.samples_per_pixel != 1 {
        return None;
    }
    let (s, o) = match pairs {
        [(s, o)] => ([*s; 3], [*o; 3]),
        [r, g, b, ..] => ([r.0, g.0, b.0], [r.1, g.1, b.1]),
        _ => return None,
    };
    Some(NoiseModel {
        s,
        o,
        row: 0.0,
        col: 0.0,
        source: Source::DngProfile,
    })
}

/// Read, row and column noise measured on the masked border, normalised.
#[derive(Clone, Copy, Debug)]
pub struct Dark {
    pub read: f32,
    pub row: f32,
    pub col: f32,
}

/// The optically black photosites beside and above the active area.
///
/// Keeps well clear of the active area: on the 6D the dozen columns nearest
/// it see light. Photosites over 8σ are the strip's own hot photosites — the
/// same ones in every frame — and are left out, as the app's hot-pixel pass
/// removes their kin before the network sees them.
pub fn dark_border(raw: &RawImage) -> Option<Dark> {
    let (x0, y0, w, h) = (
        raw.crop.x as usize,
        raw.crop.y as usize,
        raw.crop.width as usize,
        raw.crop.height as usize,
    );
    let stride = raw.width as usize;
    let span = span(raw);
    if x0 < 40 || raw.samples_per_pixel != 1 {
        return None;
    }
    let cols = 4..x0 - 16;
    let nc = cols.len() as f32;
    // Residual after removing each row's mean and each column's mean.
    let mut row_means = Vec::with_capacity(h);
    let mut col_sum = vec![0.0f64; cols.len()];
    for y in y0..y0 + h {
        let line = &raw.data[y * stride..y * stride + x0];
        let m = cols.clone().map(|x| line[x] as f32).sum::<f32>() / nc;
        row_means.push(m);
        for (k, x) in cols.clone().enumerate() {
            col_sum[k] += (line[x] as f32 - m) as f64;
        }
    }
    let col_mean: Vec<f32> = col_sum.iter().map(|s| (*s / h as f64) as f32).collect();
    let resid = |y: usize, k: usize, x: usize| {
        raw.data[y * stride + x] as f32 - row_means[y - y0] - col_mean[k]
    };
    let (mut s1, mut n) = (0.0f64, 0usize);
    for y in y0..y0 + h {
        for (k, x) in cols.clone().enumerate() {
            s1 += (resid(y, k, x) as f64).powi(2);
            n += 1;
        }
    }
    let rough = (s1 / n as f64).sqrt() as f32;
    let (mut s2, mut n2) = (0.0f64, 0usize);
    for y in y0..y0 + h {
        for (k, x) in cols.clone().enumerate() {
            let r = resid(y, k, x);
            if r.abs() < 8.0 * rough {
                s2 += (r as f64).powi(2);
                n2 += 1;
            }
        }
    }
    let read = (s2 / n2.max(1) as f64).sqrt() as f32;
    let rm = row_means.iter().sum::<f32>() / h as f32;
    let row_var = row_means.iter().map(|m| (m - rm).powi(2)).sum::<f32>() / h as f32;
    let row = (row_var - read * read / nc).max(0.0).sqrt();

    // Columns: the masked rows above the image span every column.
    let col = if y0 >= 24 {
        let rows = 4..y0 - 12;
        let nr = rows.len() as f32;
        let means: Vec<f32> = (x0..x0 + w)
            .map(|x| {
                rows.clone()
                    .map(|y| raw.data[y * stride + x] as f32)
                    .sum::<f32>()
                    / nr
            })
            .collect();
        let mm = means.iter().sum::<f32>() / means.len() as f32;
        let var = means.iter().map(|m| (m - mm).powi(2)).sum::<f32>() / means.len() as f32;
        (var - read * read / nr).max(0.0).sqrt()
    } else {
        0.0
    };
    Some(Dark {
        read: read / span,
        row: row / span,
        col: col / span,
    })
}

/// The quietest-third bias of the patch variance, and the residual bias the
/// estimate showed against the 6D's table (0.91 at the median), in one: the
/// estimate is divided by this.
const QUIET_FACTOR: f32 = 0.85 * 0.91;

/// The frame's own noise: read noise from the border (or, lacking one, the
/// floor of the quietest patches), shot gain from flat patches of one green
/// plane, the same for every colour, as a sensor's gain is.
pub fn measured(raw: &RawImage, dark: Option<&Dark>) -> Option<NoiseModel> {
    if raw.cfa_pattern.is_xtrans() || raw.samples_per_pixel != 1 {
        return None;
    }
    let m = active(raw);
    let (h, w) = (m.h, m.w);
    // One green plane at a two-photosite pitch.
    let (gy, gx) = green_offset(raw.cfa_pattern)?;
    let ph = (h - gy) / 2;
    let pw = (w - gx) / 2;
    let g = |y: usize, x: usize| m.at(gy + 2 * y, gx + 2 * x);
    const B: usize = 8;
    let mut patches: Vec<(f32, f32)> = Vec::new(); // (level, variance)
    for by in 0..ph / B {
        for bx in 0..(pw - 2) / B {
            let (mut s, mut s2, mut lv) = (0.0f32, 0.0f32, 0.0f32);
            for y in by * B..by * B + B {
                for x in bx * B..bx * B + B {
                    // Second difference: cancels any gradient; var = 6σ².
                    let d = g(y, x + 2) - 2.0 * g(y, x + 1) + g(y, x);
                    s += d;
                    s2 += d * d;
                    lv += g(y, x + 1);
                }
            }
            let n = (B * B) as f32;
            let var = (s2 / n - (s / n).powi(2)) / 6.0;
            patches.push((lv / n, var));
        }
    }
    let floor = dark.map(|d| d.read);
    let lo = 4.0 * floor.unwrap_or(0.002);
    patches.retain(|(l, _)| *l > lo && *l < 0.7);
    if patches.len() < 500 {
        return None;
    }
    patches.sort_by(|a, b| a.0.total_cmp(&b.0));
    let bins = 12;
    let per = patches.len() / bins;
    let mut ests = Vec::new();
    let mut floors = Vec::new();
    for b in 0..bins {
        let mut bin: Vec<(f32, f32)> = patches[b * per..(b + 1) * per].to_vec();
        if bin.len() < 60 {
            continue;
        }
        bin.sort_by(|a, b| a.1.total_cmp(&b.1));
        let quiet = &bin[..bin.len() / 3];
        let read2 = floor.map(|r| r * r);
        let mut e: Vec<f32> = quiet
            .iter()
            .map(|(l, v)| (v / QUIET_FACTOR - read2.unwrap_or(0.0)) / l)
            .collect();
        e.sort_by(f32::total_cmp);
        ests.push(e[e.len() / 2]);
        floors.push(quiet[quiet.len() / 2]);
    }
    ests.sort_by(f32::total_cmp);
    let s = *ests.get(ests.len() / 2)?;
    if !(s.is_finite() && s > 0.0) {
        return None;
    }
    // No border: the read variance is what the darkest bin leaves unexplained.
    let read2 = match floor {
        Some(r) => r * r,
        None => {
            let (l, v) = floors.first().copied()?;
            (v / QUIET_FACTOR - s * l).max(1e-9)
        }
    };
    Some(NoiseModel {
        s: [s; 3],
        o: [read2; 3],
        row: dark.map_or(0.0, |d| d.row),
        col: dark.map_or(0.0, |d| d.col),
        source: Source::Measured,
    })
}

/// Black-to-white range of the frame, as the demosaic normalises it.
pub(crate) fn span(raw: &RawImage) -> f32 {
    let black = raw.black_level.iter().map(|&b| b as f32).sum::<f32>() / 4.0;
    (raw.white_level as f32 - black).max(1.0)
}

/// Where a green photosite sits in the pattern's 2×2 cell, (dy, dx).
fn green_offset(p: CfaPattern) -> Option<(usize, usize)> {
    match p {
        CfaPattern::Rggb | CfaPattern::Bggr => Some((0, 1)),
        CfaPattern::Grbg | CfaPattern::Gbrg => Some((0, 0)),
        _ => None,
    }
}

/// The active area, normalised, read lazily.
pub(crate) struct Active<'a> {
    raw: &'a RawImage,
    black: [f32; 4],
    inv: [f32; 4],
    pub h: usize,
    pub w: usize,
}

impl Active<'_> {
    /// Photosite (y, x) of the active area, black 0, white 1.
    #[inline]
    pub fn at(&self, y: usize, x: usize) -> f32 {
        let c = (y & 1) * 2 + (x & 1);
        let v = self.raw.data[(self.raw.crop.y as usize + y) * self.raw.width as usize
            + self.raw.crop.x as usize
            + x];
        (v as f32 - self.black[c]) * self.inv[c]
    }
}

/// Black levels per position of the crop's 2×2 cell, as the demosaic reads
/// them: one reported level is broadcast.
pub(crate) fn active(raw: &RawImage) -> Active<'_> {
    let b = raw.black_level;
    let black = if b[1] == 0 && b[2] == 0 && b[3] == 0 {
        [b[0] as f32; 4]
    } else {
        b.map(|v| v as f32)
    };
    let inv = black.map(|bl| 1.0 / (raw.white_level as f32 - bl).max(1.0));
    Active {
        raw,
        black,
        inv,
        h: raw.crop.height as usize,
        w: raw.crop.width as usize,
    }
}