//! **eICCI** (enhanced in-loop cross-component improvement) — the DWT-domain post-filter
//! of the JPEG-AI decode path. This module holds the wavelet transform its CNN operates
//! in; the network itself (per-tile model ensemble) builds on top.
//!
//! The transform is a **Haar wavelet on 2×2 blocks** with ½ scaling, applied **two
//! levels**: a plane `[1, H, W]` becomes `[4, H/2, W/2]` and then `[16, H/4, W/4]`.
//! Stacking the three planes' 16 subbands gives the 48 channels the network's first
//! convolution consumes. Perfect reconstruction: [`haar_idwt`] inverts [`haar_dwt`]
//! exactly (`LL−HL−LH+HH = 4a` ⇒ `rec1 = a`), so `idwt2 ∘ dwt2 = id` to float precision.
//!
//! Channel layout matches the reference (`cat` along the channel axis): one level emits
//! `[LL(c), HL(c), LH(c), HH(c)]`; two levels emit the four level-1 bands each expanded
//! into its own four level-2 bands, i.e. `x16[i*4 .. i*4+4]` is the level-2 decomposition
//! of level-1 band `i`.

// On the JPEG AI exact decode path: no platform math-library call (see clippy.toml).
#![deny(clippy::disallowed_methods)]

pub use crate::jpegai_synthesis::BadMember;

/// One-level Haar DWT: `[c, h, w]` → `[4c, h/2, w/2]`, channels grouped `LL, HL, LH, HH`.
/// `h` and `w` must be even. Returns the coefficients and the output dims.
pub fn haar_dwt(x: &[f32], c: usize, h: usize, w: usize) -> (Vec<f32>, usize, usize) {
    assert_eq!(x.len(), c * h * w, "input must be [c, h, w]");
    assert!(h % 2 == 0 && w % 2 == 0, "DWT needs even dims, got {h}×{w}");
    let (ho, wo) = (h / 2, w / 2);
    let plane_o = ho * wo;
    let mut out = vec![0f32; 4 * c * plane_o];
    for ch in 0..c {
        let base = ch * h * w;
        for i in 0..ho {
            for j in 0..wo {
                // the 2×2 block, pre-scaled by ½ (matching the reference)
                let a = x[base + (2 * i) * w + 2 * j] / 2.0; // even row, even col
                let b = x[base + (2 * i + 1) * w + 2 * j] / 2.0; // odd row, even col
                let cc = x[base + (2 * i) * w + 2 * j + 1] / 2.0; // even row, odd col
                let d = x[base + (2 * i + 1) * w + 2 * j + 1] / 2.0; // odd row, odd col
                let p = i * wo + j;
                out[(ch) * plane_o + p] = a + b + cc + d; // LL
                out[(c + ch) * plane_o + p] = -a - b + cc + d; // HL
                out[(2 * c + ch) * plane_o + p] = -a + b - cc + d; // LH
                out[(3 * c + ch) * plane_o + p] = a - b - cc + d; // HH
            }
        }
    }
    (out, ho, wo)
}

/// One-level Haar IDWT: `[4c, h, w]` → `[c, 2h, 2w]` — the exact inverse of [`haar_dwt`].
pub fn haar_idwt(x: &[f32], c: usize, h: usize, w: usize) -> (Vec<f32>, usize, usize) {
    assert_eq!(x.len(), 4 * c * h * w, "input must be [4c, h, w]");
    let (ho, wo) = (h * 2, w * 2);
    let plane_i = h * w;
    let mut out = vec![0f32; c * ho * wo];
    for ch in 0..c {
        let base = ch * ho * wo;
        for i in 0..h {
            for j in 0..w {
                let p = i * w + j;
                let ll = x[(ch) * plane_i + p] / 2.0;
                let hl = x[(c + ch) * plane_i + p] / 2.0;
                let lh = x[(2 * c + ch) * plane_i + p] / 2.0;
                let hh = x[(3 * c + ch) * plane_i + p] / 2.0;
                out[base + (2 * i) * wo + 2 * j] = ll - hl - lh + hh; // even row, even col
                out[base + (2 * i + 1) * wo + 2 * j] = ll - hl + lh - hh; // odd row, even col
                out[base + (2 * i) * wo + 2 * j + 1] = ll + hl - lh - hh; // even row, odd col
                out[base + (2 * i + 1) * wo + 2 * j + 1] = ll + hl + lh + hh; // odd row, odd col
            }
        }
    }
    (out, ho, wo)
}

/// Two-level Haar DWT of a single plane: `[1, h, w]` → `[16, h/4, w/4]`.
/// Level 1 gives four bands; each is decomposed again, and the results are concatenated
/// in band order — `out[i*4 .. i*4+4]` is the level-2 decomposition of level-1 band `i`.
pub fn haar_dwt2(x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
    assert_eq!(x.len(), h * w, "input must be a single [h, w] plane");
    let (x4, h1, w1) = haar_dwt(x, 1, h, w); // [4, h/2, w/2]
    let plane1 = h1 * w1;
    let (h2, w2) = (h1 / 2, w1 / 2);
    let plane2 = h2 * w2;
    let mut out = vec![0f32; 16 * plane2];
    for band in 0..4 {
        let (sub, _, _) = haar_dwt(&x4[band * plane1..(band + 1) * plane1], 1, h1, w1); // [4, h/4, w/4]
        out[band * 4 * plane2..(band + 1) * 4 * plane2].copy_from_slice(&sub);
    }
    (out, h2, w2)
}

/// Two-level Haar IDWT: `[16, h, w]` → `[1, 4h, 4w]` — the exact inverse of [`haar_dwt2`].
pub fn haar_idwt2(x16: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
    assert_eq!(x16.len(), 16 * h * w, "input must be [16, h, w]");
    let plane = h * w;
    let (h1, w1) = (h * 2, w * 2);
    let plane1 = h1 * w1;
    let mut x4 = vec![0f32; 4 * plane1];
    for band in 0..4 {
        let (rec, _, _) = haar_idwt(&x16[band * 4 * plane..(band + 1) * 4 * plane], 1, h, w); // [1, 2h, 2w]
        x4[band * plane1..(band + 1) * plane1].copy_from_slice(&rec);
    }
    haar_idwt(&x4, 1, h1, w1) // [1, 4h, 4w]
}

/// Stack the three planes' two-level subbands into the 48-channel tensor the eICCI
/// network consumes: `cat(Y16, U16, V16)` at `[48, h/4, w/4]`. All planes must be the
/// same `h × w` (4:4:4).
pub fn icci_dwt_planes(y: &[f32], u: &[f32], v: &[f32], h: usize, w: usize)
    -> (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>, usize, usize) {
    let (y16, h2, w2) = haar_dwt2(y, h, w);
    let (u16, _, _) = haar_dwt2(u, h, w);
    let (v16, _, _) = haar_dwt2(v, h, w);
    let mut yuv48 = Vec::with_capacity(48 * h2 * w2);
    yuv48.extend_from_slice(&y16);
    yuv48.extend_from_slice(&u16);
    yuv48.extend_from_slice(&v16);
    (y16, u16, v16, yuv48, h2, w2)
}

// ---- network body ----
//
// The eICCI CNN is float in the reference, so we run it in f32 rather than quantising.
// IEEE-754 add and multiply are correctly rounded, so with the fixed accumulation order
// below — single-threaded, no fused multiply-add (Rust emits none unless asked), in the
// default floating-point environment `jpegai_decode` checks for — the port gives the
// same bytes on every conforming target, and it avoids any PTQ error. BatchNorm is
// folded into the preceding convolution, which is exact at inference:
//
//   BN(conv(x)) = γ·(w·x + b − μ)/√(σ²+ε) + β
//               = (w · γ/√(σ²+ε))·x + ((b − μ)·γ/√(σ²+ε) + β)

/// A convolution with its BatchNorm already folded in. Weights are `[cout, cin, kh, kw]`
/// row-major; `bias` is per output channel. Padding is `(ph, pw)`, stride 1.
#[derive(Clone)]
pub struct Conv {
    pub w: Vec<f32>,
    pub b: Vec<f32>,
    pub cin: usize,
    pub cout: usize,
    pub kh: usize,
    pub kw: usize,
    pub ph: usize,
    pub pw: usize,
}

impl Conv {
    /// Fold a BatchNorm into a convolution's weights/bias (exact at inference).
    /// `gamma`/`beta`/`mean`/`var` are per output channel.
    pub fn fold_bn(mut self, gamma: &[f32], beta: &[f32], mean: &[f32], var: &[f32], eps: f32) -> Self {
        assert_eq!(gamma.len(), self.cout, "BN params must be per output channel");
        let tap = self.cin * self.kh * self.kw;
        for o in 0..self.cout {
            let s = gamma[o] / (var[o] + eps).sqrt();
            for t in 0..tap {
                self.w[o * tap + t] *= s;
            }
            self.b[o] = (self.b[o] - mean[o]) * s + beta[o];
        }
        self
    }

    /// Stride-1 zero-padded convolution: `[cin, h, w]` → `[cout, h, w]`.
    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
        assert_eq!(x.len(), self.cin * h * w, "conv input must be [cin, h, w]");
        let mut out = vec![0f32; self.cout * h * w];
        for o in 0..self.cout {
            for oy in 0..h {
                for ox in 0..w {
                    let mut acc = self.b[o];
                    for i in 0..self.cin {
                        for ky in 0..self.kh {
                            let iy = oy + ky;
                            if iy < self.ph || iy - self.ph >= h { continue; }
                            let sy = iy - self.ph;
                            for kx in 0..self.kw {
                                let ix = ox + kx;
                                if ix < self.pw || ix - self.pw >= w { continue; }
                                let sx = ix - self.pw;
                                acc += self.w[((o * self.cin + i) * self.kh + ky) * self.kw + kx]
                                     * x[(i * h + sy) * w + sx];
                            }
                        }
                    }
                    out[(o * h + oy) * w + ox] = acc;
                }
            }
        }
        out
    }
}

/// One `ResidualBlock_BN_RectKernel`: separable rectangular kernels with a scaled residual.
///
/// ```text
/// out = x + scale · BN2( conv2_{3×1}( relu( BN1( conv1_{1×3}(x) ) ) ) )
/// ```
///
/// Both BatchNorms are folded into their convolutions, so this is two convs + a ReLU.
#[derive(Clone)]
pub struct ResBlock {
    pub conv1: Conv, // [1,3] kernel, pad [0,1]
    pub conv2: Conv, // [3,1] kernel, pad [1,0]
    pub scale: f32,
}

impl ResBlock {
    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
        let mut t = self.conv1.forward(x, h, w);
        for v in t.iter_mut() { if *v < 0.0 { *v = 0.0; } }
        let t = self.conv2.forward(&t, h, w);
        x.iter().zip(&t).map(|(&a, &b)| a + b * self.scale).collect()
    }
}

/// One plane-path of the eICCI network: `conv_first` + BN + ReLU → N residual blocks →
/// `conv_last`, then a scaled residual add onto the plane's own subbands.
///
/// ```text
/// fea      = relu( BN( conv_first(yuv48) ) )      // BN folded into conv_first
/// fea      = resblocks(fea)
/// out16    = conv_last(fea) · scale + subbands16
/// plane    = idwt2(out16)
/// ```
#[derive(Clone)]
pub struct IcciPath {
    pub conv_first: Conv, // 48 → nf, 3×3, pad 1 (BN folded)
    pub blocks: Vec<ResBlock>,
    pub conv_last: Conv, // nf → 16, 3×3, pad 1
    pub scale: f32,
}

impl IcciPath {
    /// Run the path and return the refined 16 subbands (before the inverse transform).
    pub fn refine_subbands(&self, yuv48: &[f32], subbands16: &[f32], h: usize, w: usize) -> Vec<f32> {
        let mut fea = self.conv_first.forward(yuv48, h, w);
        for v in fea.iter_mut() { if *v < 0.0 { *v = 0.0; } }
        for blk in &self.blocks {
            fea = blk.forward(&fea, h, w);
        }
        let res = self.conv_last.forward(&fea, h, w);
        assert_eq!(res.len(), subbands16.len(), "conv_last must emit the 16 subbands");
        res.iter().zip(subbands16).map(|(&r, &s)| r * self.scale + s).collect()
    }

    /// Full plane: refine the subbands then invert the two-level transform.
    pub fn forward_plane(&self, yuv48: &[f32], subbands16: &[f32], h: usize, w: usize) -> Vec<f32> {
        let refined = self.refine_subbands(yuv48, subbands16, h, w);
        haar_idwt2(&refined, h, w).0
    }
}

/// The whole eICCI filter for one tile: three planes in, three refined planes out.
///
/// Shape follows the exported model, **not** the reference Python class (which says
/// `nf=64, nbY=8, nbUV=8`): the shipped nets are `nf=48`, **2** luma blocks and **4**
/// chroma blocks. The chroma trunk is shared — the reference traces it once per chroma
/// plane, computing identical features twice; we evaluate it once and feed both heads.
///
/// ```text
/// (Y16,U16,V16), yuv48 = dwt2(Y), dwt2(U), dwt2(V), cat(..)
/// Y_o = idwt2( conv_last_Y(luma_trunk(yuv48))  · scale + Y16 )
/// fea = chroma_trunk(yuv48)                      // evaluated once
/// U_o = idwt2( conv_last_u(fea) · scale + U16 )
/// V_o = idwt2( conv_last_v(fea) · scale + V16 )
/// ```
#[derive(Clone)]
pub struct IcciNet {
    /// luma: `conv_first_Y` (BN pre-folded by the exporter) + its residual blocks
    pub luma: IcciPath,
    /// chroma trunk: `conv_firstUV` (BN pre-folded) + its residual blocks; `conv_last`
    /// is the U head and `chroma.scale` is the **U** residual scale (`scaleU`).
    pub chroma: IcciPath,
    /// the V head — shares `chroma`'s trunk, differs only in this final convolution
    pub conv_last_v: Conv,
    /// the V residual scale (`scaleV`). The reference signals the U and V scales
    /// **separately** (`comp_model.scaleU` / `scaleV`); on the shipped models they
    /// differ by ~1%, so reusing one scale for both heads shifts V by more than a
    /// pixel LSB.
    pub scale_v: f32,
}

impl IcciNet {
    /// Refine one tile's three planes. `y`/`u`/`v` are `[h × w]` (4:4:4); `h` and `w`
    /// must be multiples of 4 (two wavelet levels).
    pub fn forward(&self, y: &[f32], u: &[f32], v: &[f32], h: usize, w: usize)
        -> (Vec<f32>, Vec<f32>, Vec<f32>) {
        let (y16, u16, v16, yuv48, h2, w2) = icci_dwt_planes(y, u, v, h, w);

        let y_o = self.luma.forward_plane(&yuv48, &y16, h2, w2);

        // shared chroma trunk, evaluated once
        let mut fea = self.chroma.conv_first.forward(&yuv48, h2, w2);
        for t in fea.iter_mut() { if *t < 0.0 { *t = 0.0; } }
        for blk in &self.chroma.blocks {
            fea = blk.forward(&fea, h2, w2);
        }
        let head = |conv: &Conv, band: &[f32], scale: f32| -> Vec<f32> {
            let res = conv.forward(&fea, h2, w2);
            let refined: Vec<f32> = res.iter().zip(band).map(|(&r, &s)| r * scale + s).collect();
            haar_idwt2(&refined, h2, w2).0
        };
        let u_o = head(&self.chroma.conv_last, &u16, self.chroma.scale);
        let v_o = head(&self.conv_last_v, &v16, self.scale_v);
        (y_o, u_o, v_o)
    }

    /// Parse the whole eICCI net from the `icci` weight export
    /// (`{luma, chroma, conv_last_v, scale_v}`, produced by
    /// `tools/wai_jpegai_export_weights.py icci`). A member missing or of another type is
    /// refused, named by its path ([`BadMember`]); the loader used to panic on one.
    pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
        Ok(IcciNet {
            luma: IcciPath::from_json(&v["luma"]).map_err(|e| e.under("luma"))?,
            chroma: IcciPath::from_json(&v["chroma"]).map_err(|e| e.under("chroma"))?,
            conv_last_v: Conv::from_json(&v["conv_last_v"]).map_err(|e| e.under("conv_last_v"))?,
            scale_v: v["scale_v"].as_f64().ok_or_else(|| BadMember("scale_v".into()))? as f32,
        })
    }
}

impl Conv {
    /// Parse a BN-folded conv from an export node `{w, b, cin, cout, kh, kw, ph, pw}`. A
    /// field missing or of another type is refused ([`BadMember`], named by the field).
    pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
        let g = |k: &str| v[k].as_u64().and_then(|x| usize::try_from(x).ok()).ok_or_else(|| BadMember(k.into()));
        let f = |k: &str| -> Result<Vec<f32>, BadMember> {
            v[k].as_array().and_then(|a| a.iter().map(|x| x.as_f64().map(|x| x as f32)).collect::<Option<Vec<f32>>>()).ok_or_else(|| BadMember(k.into()))
        };
        Ok(Conv { w: f("w")?, b: f("b")?, cin: g("cin")?, cout: g("cout")?,
                  kh: g("kh")?, kw: g("kw")?, ph: g("ph")?, pw: g("pw")? })
    }
}

impl IcciPath {
    /// Parse one plane-path (`{conv_first, blocks:[{c1,c2,scale}], conv_last, scale}`).
    pub fn from_json(v: &serde_json::Value) -> Result<Self, BadMember> {
        let blocks = v["blocks"].as_array().ok_or_else(|| BadMember("blocks".into()))?;
        Ok(IcciPath {
            conv_first: Conv::from_json(&v["conv_first"]).map_err(|e| e.under("conv_first"))?,
            blocks: blocks.iter().enumerate().map(|(i, b)| -> Result<ResBlock, BadMember> {
                let at = |e: BadMember| e.under(&format!("blocks.{i}"));
                Ok(ResBlock {
                    conv1: Conv::from_json(&b["c1"]).map_err(|e| at(e.under("c1")))?,
                    conv2: Conv::from_json(&b["c2"]).map_err(|e| at(e.under("c2")))?,
                    scale: b["scale"].as_f64().ok_or_else(|| at(BadMember("scale".into())))? as f32,
                })
            }).collect::<Result<_, _>>()?,
            conv_last: Conv::from_json(&v["conv_last"]).map_err(|e| e.under("conv_last"))?,
            scale: v["scale"].as_f64().ok_or_else(|| BadMember("scale".into()))? as f32,
        })
    }
}

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

    fn ramp(n: usize) -> Vec<f32> {
        (0..n).map(|i| ((i * 37 % 251) as f32) - 125.0).collect()
    }

    #[test]
    fn haar_dwt_single_block_matches_hand_math() {
        // one 2×2 block: a=(0,0)=1, c=(0,1)=2, b=(1,0)=3, d=(1,1)=4  (row-major [1,2,3,4])
        let (o, ho, wo) = haar_dwt(&[1.0, 2.0, 3.0, 4.0], 1, 2, 2);
        assert_eq!((ho, wo), (1, 1));
        // halves: a=.5 b=1.5 c=1 d=2  →  LL=5, HL=1, LH=2, HH=0
        assert!((o[0] - 5.0).abs() < 1e-6, "LL {}", o[0]);
        assert!((o[1] - 1.0).abs() < 1e-6, "HL {}", o[1]);
        assert!((o[2] - 2.0).abs() < 1e-6, "LH {}", o[2]);
        assert!((o[3] - 0.0).abs() < 1e-6, "HH {}", o[3]);
    }

    #[test]
    fn haar_idwt_inverts_dwt_exactly() {
        let (h, w) = (8usize, 6usize);
        let x = ramp(h * w);
        let (d, dh, dw) = haar_dwt(&x, 1, h, w);
        let (r, rh, rw) = haar_idwt(&d, 1, dh, dw);
        assert_eq!((rh, rw), (h, w));
        for i in 0..x.len() {
            assert!((r[i] - x[i]).abs() < 1e-4, "idwt∘dwt != id at {i}: {} vs {}", r[i], x[i]);
        }
    }

    #[test]
    fn two_level_round_trips_and_has_16_bands() {
        let (h, w) = (16usize, 16usize);
        let x = ramp(h * w);
        let (x16, h2, w2) = haar_dwt2(&x, h, w);
        assert_eq!((h2, w2), (4, 4));
        assert_eq!(x16.len(), 16 * 4 * 4, "two levels must give 16 subbands at h/4");
        let (r, rh, rw) = haar_idwt2(&x16, h2, w2);
        assert_eq!((rh, rw), (h, w));
        for i in 0..x.len() {
            assert!((r[i] - x[i]).abs() < 1e-3, "idwt2∘dwt2 != id at {i}: {} vs {}", r[i], x[i]);
        }
    }

    #[test]
    fn planes_stack_to_48_channels() {
        let (h, w) = (8usize, 8usize);
        let (y, u, v) = (ramp(h * w), ramp(h * w), ramp(h * w));
        let (y16, _, _, yuv48, h2, w2) = icci_dwt_planes(&y, &u, &v, h, w);
        assert_eq!((h2, w2), (2, 2));
        assert_eq!(yuv48.len(), 48 * h2 * w2, "network input must be 48 channels");
        // the Y block is the leading 16 channels
        assert_eq!(&yuv48[..y16.len()], &y16[..]);
    }

    fn ident_conv(c: usize, kh: usize, kw: usize, ph: usize, pw: usize) -> Conv {
        // identity: centre tap 1 on the matching in/out channel, everything else 0
        let mut w = vec![0f32; c * c * kh * kw];
        for o in 0..c {
            w[((o * c + o) * kh + kh / 2) * kw + kw / 2] = 1.0;
        }
        Conv { w, b: vec![0f32; c], cin: c, cout: c, kh, kw, ph, pw }
    }

    #[test]
    fn conv_identity_and_zero_padding() {
        let x = vec![1.0f32, 2.0, 3.0, 4.0];
        // 1×3 identity with pad [0,1] reproduces the input
        let o = ident_conv(1, 1, 3, 0, 1).forward(&x, 2, 2);
        assert_eq!(o, x, "identity conv must be a no-op");

        // a pure left-shift tap reads the zero pad at the left edge
        let mut c = Conv { w: vec![0.0, 0.0, 1.0], b: vec![0.0], cin: 1, cout: 1, kh: 1, kw: 3, ph: 0, pw: 1 };
        c.b[0] = 0.0;
        let o = c.forward(&[5.0, 7.0], 1, 2);
        // tap at kx=2 reads x[ox+2-1] = x[ox+1]: out[0]=x[1]=7, out[1]=pad=0
        assert_eq!(o, vec![7.0, 0.0], "zero padding at the right edge");
    }

    #[test]
    fn fold_bn_matches_explicit_batchnorm() {
        // conv = identity, BN with known params → folded conv must equal BN(x) directly
        let (gamma, beta, mean, var, eps) = ([2.0f32], [0.5f32], [1.0f32], [4.0f32], 1e-5f32);
        let folded = ident_conv(1, 1, 1, 0, 0).fold_bn(&gamma, &beta, &mean, &var, eps);
        let x = [3.0f32, -1.0];
        let got = folded.forward(&x, 1, 2);
        for i in 0..x.len() {
            let want = gamma[0] * (x[i] - mean[0]) / (var[0] + eps).sqrt() + beta[0];
            assert!((got[i] - want).abs() < 1e-5, "fold != BN: {} vs {want}", got[i]);
        }
    }

    #[test]
    fn resblock_is_identity_when_scale_is_zero() {
        let c = 2usize;
        let (h, w) = (3usize, 3usize);
        let x = ramp(c * h * w);
        let blk = ResBlock {
            conv1: ident_conv(c, 1, 3, 0, 1),
            conv2: ident_conv(c, 3, 1, 1, 0),
            scale: 0.0,
        };
        assert_eq!(blk.forward(&x, h, w), x, "scale=0 must leave the identity path alone");

        // scale=1 with identity convs adds relu(x) back onto x
        let blk1 = ResBlock { scale: 1.0, ..blk };
        let got = blk1.forward(&x, h, w);
        for i in 0..x.len() {
            let want = x[i] + if x[i] > 0.0 { x[i] } else { 0.0 };
            assert!((got[i] - want).abs() < 1e-5, "at {i}: {} vs {want}", got[i]);
        }
    }
}
