//! JPEG AI (ISO/IEC 6048-1) decoder-side post-filters.
//!
//! Four filters run between synthesis and colour conversion, in order:
//! EFElinear → eICCI → EFEnonlinear → LEF. eICCI is neural and executes
//! through ORT (faithful per-tile wrappers); the other three are
//! deterministic signalled math and are implemented natively here.
//!
//! All filters operate on full-resolution YUV planes in the [0,255] float
//! domain (the codec's internal 8-bit range).

/// One chroma plane's EFElinear parameters at 4:4:4 (`fL = 1`): a per-phase
/// affine cross-component correction. The 2×2 spatial phases of the chroma
/// plane each get `out = w_uv·(c − mean) + bias + w_y·y`, with the co-sited
/// luma sample as `y`. Coefficients arrive Q15-signalled in the picture
/// header; here they are already de-integered floats (reference
/// `deinteger`: `(q − 32767.5… ) / 2^(wP−5)` — see the VM).
pub struct EfeLinearPlane {
    pub mean: f32,
    /// Per-phase chroma weight (4 phases: TL, TR, BL, BR of the 2×2 grid).
    pub w_uv: [f32; 4],
    pub bias: [f32; 4],
    /// Per-phase luma weight.
    pub w_y: [f32; 4],
}

/// Apply EFElinear at 4:4:4 with single-tap signalled filters (the shape
/// this codec's default operating point emits). `y`/`c` are `h×w` planes;
/// returns the corrected chroma plane. Luma is never modified.
pub fn efe_linear_plane(y: &[f32], c: &[f32], h: usize, w: usize, p: &EfeLinearPlane) -> Vec<f32> {
    let mut out = vec![0f32; h * w];
    for row in 0..h {
        for col in 0..w {
            // phase index in raster order of the 2×2 grid
            let ph = (row & 1) * 2 + (col & 1);
            let i = row * w + col;
            out[i] = p.w_uv[ph] * (c[i] - p.mean) + p.bias[ph] + p.w_y[ph] * y[i];
        }
    }
    out
}

/// LEF (luma enhancement filter): σ-conditioned unsharp masking of the luma
/// plane only. `y` is the [0,255] plane; `sigma` is the final Y scale
/// surface at latent resolution (`h/16 × w/16`), nearest-upsampled here.
/// Per pixel: blur with `[[5,5,5],[5,24,5],[5,5,5]]/64` (borders keep the
/// source value), pick a magnitude by σ bucket — `[thr0,thr1)→mag0,
/// [thr1,thr2)→mag1, ≥thr2→mag2`, else 1 — then
/// `out = blur + (src − blur)·mag`, clamped to [0,1] in the normalized
/// domain. Thresholds/magnitudes are fixed per operating point
/// (`recSharpThrList` / `recSharpMagList` in the VM).
pub fn lef_apply(
    y: &[f32],
    sigma: &[i64],
    h: usize,
    w: usize,
    thr: [i64; 3],
    mag: [f32; 3],
) -> Vec<f32> {
    let sh = h / 16;
    let sw = w / 16;
    let x: Vec<f32> = y.iter().map(|&v| v / 255.0).collect();
    let mut result = x.clone();
    const K: [[f32; 3]; 3] = [[5.0, 5.0, 5.0], [5.0, 24.0, 5.0], [5.0, 5.0, 5.0]];
    for row in 1..h - 1 {
        for col in 1..w - 1 {
            let mut acc = 0f32;
            for (dy, krow) in K.iter().enumerate() {
                for (dx, &kv) in krow.iter().enumerate() {
                    acc += kv * x[(row + dy - 1) * w + col + dx - 1];
                }
            }
            result[row * w + col] = (acc / 64.0).clamp(0.0, 1.0);
        }
    }
    let mut out = vec![0f32; h * w];
    for row in 0..h {
        for col in 0..w {
            let s = sigma[(row / (h / sh)) * sw + col / (w / sw)];
            let f = if s >= thr[2] {
                mag[2]
            } else if s >= thr[1] {
                mag[1]
            } else if s >= thr[0] {
                mag[0]
            } else {
                1.0
            };
            let i = row * w + col;
            out[i] = (result[i] + (x[i] - result[i]) * f).clamp(0.0, 1.0) * 255.0;
        }
    }
    out
}

/// EFEnonlinear: luma-conditioned piecewise-linear chroma correction (one
/// tile). Signalled per-tile: 8 Q16 weights (`deinteger`:
/// `(q − 32767) / 2^11`), luma min/max ×100. Breakpoints split
/// `[min, max]` into `hist_split_num = 8` segments;
/// `out = relu(c + w0·(y − bp0)) + Σ_{i=1..7} wi·relu(y − bpi)`.
pub fn efe_nonlinear_plane(
    y: &[f32],
    c: &[f32],
    weights_q: &[i64; 8],
    luma_min_q: i64,
    luma_max_q: i64,
) -> Vec<f32> {
    let w: Vec<f32> = weights_q.iter().map(|&q| (q - 32767) as f32 / 2048.0).collect();
    let mn = luma_min_q as f32 / 100.0;
    let gap = (luma_max_q as f32 / 100.0 - mn) / 8.0;
    let bp: Vec<f32> = (0..8).map(|i| mn + i as f32 * gap).collect();
    y.iter()
        .zip(c)
        .map(|(&yy, &cc)| {
            let mut acc = (cc + w[0] * (yy - bp[0])).max(0.0);
            for i in 1..8 {
                acc += w[i] * (yy - bp[i]).max(0.0);
            }
            acc
        })
        .collect()
}
