Skip to main content

yscv_video/
hevc_filter.rs

1//! HEVC in-loop filters (deblocking + SAO) and chroma decoding.
2//!
3//! Implements the deblocking filter per ITU-T H.265 section 8.7.2 and the
4//! Sample Adaptive Offset (SAO) per section 8.7.3, along with chroma-plane
5//! reconstruction utilities for 4:2:0 subsampling.
6
7use super::hevc_cabac::CabacDecoder;
8use super::hevc_decoder::HevcPredMode;
9use super::hevc_syntax::CodingUnitData;
10
11// ---------------------------------------------------------------------------
12// HEVC tc / beta threshold tables (ITU-T H.265, Tables 8-11, 8-12)
13// ---------------------------------------------------------------------------
14
15/// `tc` (clipping threshold) table indexed by `Q = Clip3(0, 53, qP_L + 2*(bs-1) + tc_offset)`.
16/// ITU-T H.265 Table 8-11.
17#[rustfmt::skip]
18pub const TC_TABLE: [i32; 54] = [
19     0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
20     0,  0,  0,  0,  0,  0,  0,  0,  1,  1,
21     1,  1,  1,  1,  1,  1,  1,  2,  2,  2,
22     2,  3,  3,  3,  3,  4,  4,  4,  5,  5,
23     6,  6,  7,  8,  9, 10, 11, 13, 14, 16,
24    18, 20, 22, 24,
25];
26
27/// `beta` (decision threshold) table indexed by `Q' = Clip3(0, 51, qP_L + beta_offset)`.
28/// ITU-T H.265 Table 8-12.
29#[rustfmt::skip]
30pub const BETA_TABLE: [i32; 52] = [
31     0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
32     0,  0,  0,  0,  0,  0,  6,  7,  8,  9,
33    10, 11, 12, 13, 14, 15, 16, 17, 18, 20,
34    22, 24, 26, 28, 30, 32, 34, 36, 38, 40,
35    42, 44, 46, 48, 50, 52, 54, 56, 58, 60,
36    62, 64,
37];
38
39// ---------------------------------------------------------------------------
40// Chroma QP mapping (ITU-T H.265, Table 8-10)
41// ---------------------------------------------------------------------------
42
43/// Map luma QP to chroma QP for 4:2:0 (Table 8-10).
44#[rustfmt::skip]
45const CHROMA_QP_TABLE: [u8; 58] = [
46     0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
47    10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
48    20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
49    29, 30, 31, 32, 33, 33, 34, 34, 35, 35,
50    36, 36, 37, 37, 38, 39, 40, 41, 42, 43,
51    44, 45, 46, 47, 48, 49, 50, 51,
52];
53
54/// Derive chroma QP from luma QP (clamped).
55pub fn derive_chroma_qp(luma_qp: u8) -> u8 {
56    let idx = (luma_qp as usize).min(CHROMA_QP_TABLE.len() - 1);
57    CHROMA_QP_TABLE[idx]
58}
59
60// ---------------------------------------------------------------------------
61// Threshold lookup helpers
62// ---------------------------------------------------------------------------
63
64/// Look up the `tc` threshold from QP and boundary strength.
65pub fn derive_tc(qp: u8, bs: u8) -> i32 {
66    if bs == 0 {
67        return 0;
68    }
69    let q = (qp as i32 + 2 * (bs as i32 - 1)).clamp(0, 53);
70    TC_TABLE[q as usize]
71}
72
73/// Look up the `beta` threshold from QP.
74pub fn derive_beta(qp: u8) -> i32 {
75    let idx = (qp as usize).min(51);
76    BETA_TABLE[idx]
77}
78
79// ---------------------------------------------------------------------------
80// Boundary strength (ITU-T H.265, 8.7.2.4)
81// ---------------------------------------------------------------------------
82
83/// Compute boundary strength (bs) for an edge between blocks P and Q.
84///
85/// - `bs = 2`: at least one side is intra-coded.
86/// - `bs = 1`: different reference indices or motion-vector difference >= 1 integer pel.
87/// - `bs = 0`: no filtering.
88pub fn hevc_boundary_strength(
89    is_intra_p: bool,
90    is_intra_q: bool,
91    ref_idx_p: i8,
92    ref_idx_q: i8,
93    mv_p: (i16, i16),
94    mv_q: (i16, i16),
95) -> u8 {
96    if is_intra_p || is_intra_q {
97        return 2;
98    }
99    if ref_idx_p != ref_idx_q {
100        return 1;
101    }
102    // MV difference >= 1 integer pel (4 quarter-pel units)
103    let dx = (mv_p.0 as i32 - mv_q.0 as i32).unsigned_abs();
104    let dy = (mv_p.1 as i32 - mv_q.1 as i32).unsigned_abs();
105    if dx >= 4 || dy >= 4 {
106        return 1;
107    }
108    0
109}
110
111// ---------------------------------------------------------------------------
112// Luma edge filtering (ITU-T H.265, 8.7.2.5.3 / 8.7.2.5.4)
113// ---------------------------------------------------------------------------
114
115/// Filter one 4-sample luma edge.
116///
117/// `samples` is a contiguous buffer containing `p3, p2, p1, p0, q0, q1, q2, q3`
118/// at positions `[offset - 3*stride .. offset + 4*stride]` where `offset` points
119/// to p0, but here we pass a small slice of 8 entries at `stride = 1` for
120/// clarity.
121///
122/// For frame-level usage the caller extracts the 8 relevant samples, calls this
123/// function, then writes them back.
124pub fn hevc_filter_edge_luma(samples: &mut [u8], stride: usize, bs: u8, qp: u8, bit_depth: u8) {
125    if bs == 0 || samples.len() < 8 * stride {
126        return;
127    }
128    // samples layout: p3 p2 p1 p0 q0 q1 q2 q3
129    // indices:         0  1  2  3  4  5  6  7  (when stride=1)
130    let p3_idx = 0;
131    let p2_idx = stride;
132    let p1_idx = 2 * stride;
133    let p0_idx = 3 * stride;
134    let q0_idx = 4 * stride;
135    let q1_idx = 5 * stride;
136    let q2_idx = 6 * stride;
137    let q3_idx = 7 * stride;
138
139    // Bounds check
140    if q3_idx >= samples.len() {
141        return;
142    }
143
144    let p0 = samples[p0_idx] as i32;
145    let p1 = samples[p1_idx] as i32;
146    let p2 = samples[p2_idx] as i32;
147    let p3 = samples[p3_idx] as i32;
148    let q0 = samples[q0_idx] as i32;
149    let q1 = samples[q1_idx] as i32;
150    let q2 = samples[q2_idx] as i32;
151    let q3 = samples[q3_idx] as i32;
152
153    let tc = derive_tc(qp, bs);
154    let beta = derive_beta(qp);
155    let max_val = (1i32 << bit_depth) - 1;
156
157    // Decision thresholds (ITU-T H.265, 8.7.2.5.2)
158    // d = dp0 + dq0; if d >= beta, don't filter.
159    let dp0 = (p2 - 2 * p1 + p0).abs();
160    let dq0 = (q2 - 2 * q1 + q0).abs();
161    let dp3 = (p3 - 2 * p2 + p1).abs();
162    let dq3 = (q3 - 2 * q2 + q1).abs();
163    let d = dp0 + dq0;
164
165    if d >= beta {
166        return;
167    }
168
169    let d_strong = d + dp3 + dq3;
170
171    let strong = d_strong < (beta >> 3) && (p0 - q0).abs() < ((5 * tc + 1) >> 1);
172
173    if strong {
174        // Strong filter (ITU-T H.265, 8.7.2.5.4)
175        samples[p0_idx] = ((p2 + 2 * p1 + 2 * p0 + 2 * q0 + q1 + 4) >> 3).clamp(0, max_val) as u8;
176        samples[p1_idx] = ((p2 + p1 + p0 + q0 + 2) >> 2).clamp(0, max_val) as u8;
177        samples[p2_idx] = ((2 * p3 + 3 * p2 + p1 + p0 + q0 + 4) >> 3).clamp(0, max_val) as u8;
178        samples[q0_idx] = ((q2 + 2 * q1 + 2 * q0 + 2 * p0 + p1 + 4) >> 3).clamp(0, max_val) as u8;
179        samples[q1_idx] = ((q2 + q1 + q0 + p0 + 2) >> 2).clamp(0, max_val) as u8;
180        samples[q2_idx] = ((2 * q3 + 3 * q2 + q1 + q0 + p0 + 4) >> 3).clamp(0, max_val) as u8;
181    } else {
182        // Weak filter (ITU-T H.265, 8.7.2.5.3)
183        let delta = (9 * (q0 - p0) - 3 * (q1 - p1) + 8) >> 4;
184        if delta.abs() < tc * 10 {
185            let delta_clamped = delta.clamp(-tc, tc);
186            samples[p0_idx] = (p0 + delta_clamped).clamp(0, max_val) as u8;
187            samples[q0_idx] = (q0 - delta_clamped).clamp(0, max_val) as u8;
188
189            // Conditional modification of p1 and q1
190            let tc2 = tc >> 1;
191            if dp0 + dp3 < (beta + (beta >> 1)) >> 3 {
192                let delta_p = ((((p2 + p0 + 1) >> 1) - p1) + delta_clamped) >> 1;
193                samples[p1_idx] = (p1 + delta_p.clamp(-tc2, tc2)).clamp(0, max_val) as u8;
194            }
195            if dq0 + dq3 < (beta + (beta >> 1)) >> 3 {
196                let delta_q = ((((q2 + q0 + 1) >> 1) - q1) - delta_clamped) >> 1;
197                samples[q1_idx] = (q1 + delta_q.clamp(-tc2, tc2)).clamp(0, max_val) as u8;
198            }
199        }
200    }
201}
202
203/// Filter one 2-sample chroma edge (ITU-T H.265, 8.7.2.5.5).
204///
205/// Chroma edges are only filtered when `bs == 2`. Layout uses 4 samples:
206/// `p1, p0 | q0, q1` at indices `[0, stride, 2*stride, 3*stride]`.
207pub fn hevc_filter_edge_chroma(samples: &mut [u8], stride: usize, bs: u8, qp: u8, bit_depth: u8) {
208    // HEVC spec: chroma edges are only filtered for bs == 2
209    if bs < 2 {
210        return;
211    }
212    let p1_idx = 0;
213    let p0_idx = stride;
214    let q0_idx = 2 * stride;
215    let q1_idx = 3 * stride;
216
217    if q1_idx >= samples.len() {
218        return;
219    }
220
221    let p1 = samples[p1_idx] as i32;
222    let p0 = samples[p0_idx] as i32;
223    let q0 = samples[q0_idx] as i32;
224    let q1 = samples[q1_idx] as i32;
225
226    let tc = derive_tc(qp, bs);
227    if tc == 0 {
228        return;
229    }
230
231    let max_val = (1i32 << bit_depth) - 1;
232    let delta = ((((q0 - p0) * 4) + p1 - q1 + 4) >> 3).clamp(-tc, tc);
233    samples[p0_idx] = (p0 + delta).clamp(0, max_val) as u8;
234    samples[q0_idx] = (q0 - delta).clamp(0, max_val) as u8;
235}
236
237// ---------------------------------------------------------------------------
238// Frame-level deblocking (ITU-T H.265, 8.7.2)
239// ---------------------------------------------------------------------------
240
241/// Apply deblocking filter to an entire frame (all three planes).
242///
243/// Processes vertical edges first, then horizontal edges (as specified in the
244/// HEVC standard). `qp_map` provides the QP for each CTU (row-major, indexed
245/// by `(ctu_y / ctu_size) * ctu_cols + (ctu_x / ctu_size)`). `cu_edges` marks
246/// CU/TU boundaries on a `min_cu_size` grid, stored row-major as
247/// `(y / min_cu_size) * grid_cols + (x / min_cu_size)`.
248pub fn hevc_deblock_frame(
249    luma: &mut [u8],
250    cb: &mut [u8],
251    cr: &mut [u8],
252    width: usize,
253    height: usize,
254    qp_map: &[u8],
255    cu_edges: &[bool],
256    min_cu_size: usize,
257) {
258    if width == 0 || height == 0 || min_cu_size == 0 {
259        return;
260    }
261
262    let grid_cols = width.div_ceil(min_cu_size);
263    let ctu_size = 64usize.min(width).min(height);
264    let ctu_cols = width.div_ceil(ctu_size);
265
266    // Helper to look up QP for a pixel position.
267    let qp_at = |px: usize, py: usize| -> u8 {
268        let cx = (px / ctu_size).min(ctu_cols.saturating_sub(1));
269        let cy = py / ctu_size;
270        let idx = cy * ctu_cols + cx;
271        if idx < qp_map.len() {
272            qp_map[idx]
273        } else {
274            26 // fallback
275        }
276    };
277
278    // Helper: is there a CU/TU edge at grid position (gx, gy)?
279    let has_edge = |gx: usize, gy: usize| -> bool {
280        let idx = gy * grid_cols + gx;
281        if idx < cu_edges.len() {
282            cu_edges[idx]
283        } else {
284            true // picture boundaries are always edges
285        }
286    };
287
288    // --- Vertical edges (process column by column on the min_cu grid) ---
289    for gy in 0..(height / min_cu_size) {
290        for gx in 1..(width / min_cu_size) {
291            if !has_edge(gx, gy) {
292                continue;
293            }
294            let edge_x = gx * min_cu_size;
295            let edge_y = gy * min_cu_size;
296            let qp = qp_at(edge_x, edge_y);
297
298            // Luma: filter each row within this min_cu_size block
299            for row in 0..min_cu_size {
300                let y = edge_y + row;
301                if y >= height || edge_x + 3 >= width || edge_x < 4 {
302                    continue;
303                }
304                // Extract 8 samples: p3 p2 p1 p0 q0 q1 q2 q3
305                let mut buf = [0u8; 8];
306                for i in 0..4 {
307                    buf[i] = luma[y * width + edge_x - 4 + i];
308                }
309                for i in 0..4 {
310                    buf[4 + i] = luma[y * width + edge_x + i];
311                }
312                hevc_filter_edge_luma(&mut buf, 1, 2, qp, 8);
313                for i in 0..4 {
314                    luma[y * width + edge_x - 4 + i] = buf[i];
315                }
316                for i in 0..4 {
317                    luma[y * width + edge_x + i] = buf[4 + i];
318                }
319            }
320
321            // Chroma (4:2:0): half resolution edges
322            let chroma_w = width / 2;
323            let chroma_h = height / 2;
324            let cx = edge_x / 2;
325            let cy = edge_y / 2;
326            let c_rows = min_cu_size / 2;
327            let chroma_qp = derive_chroma_qp(qp);
328            if cx >= 2 && cx + 1 < chroma_w {
329                for row in 0..c_rows {
330                    let cy_r = cy + row;
331                    if cy_r >= chroma_h {
332                        continue;
333                    }
334                    // Cb
335                    let mut buf_c = [0u8; 4];
336                    buf_c[0] = cb[cy_r * chroma_w + cx - 2];
337                    buf_c[1] = cb[cy_r * chroma_w + cx - 1];
338                    buf_c[2] = cb[cy_r * chroma_w + cx];
339                    buf_c[3] = cb[cy_r * chroma_w + cx + 1];
340                    hevc_filter_edge_chroma(&mut buf_c, 1, 2, chroma_qp, 8);
341                    cb[cy_r * chroma_w + cx - 2] = buf_c[0];
342                    cb[cy_r * chroma_w + cx - 1] = buf_c[1];
343                    cb[cy_r * chroma_w + cx] = buf_c[2];
344                    cb[cy_r * chroma_w + cx + 1] = buf_c[3];
345                    // Cr
346                    buf_c[0] = cr[cy_r * chroma_w + cx - 2];
347                    buf_c[1] = cr[cy_r * chroma_w + cx - 1];
348                    buf_c[2] = cr[cy_r * chroma_w + cx];
349                    buf_c[3] = cr[cy_r * chroma_w + cx + 1];
350                    hevc_filter_edge_chroma(&mut buf_c, 1, 2, chroma_qp, 8);
351                    cr[cy_r * chroma_w + cx - 2] = buf_c[0];
352                    cr[cy_r * chroma_w + cx - 1] = buf_c[1];
353                    cr[cy_r * chroma_w + cx] = buf_c[2];
354                    cr[cy_r * chroma_w + cx + 1] = buf_c[3];
355                }
356            }
357        }
358    }
359
360    // --- Horizontal edges (process row by row on the min_cu grid) ---
361    for gy in 1..(height / min_cu_size) {
362        for gx in 0..(width / min_cu_size) {
363            if !has_edge(gx, gy) {
364                continue;
365            }
366            let edge_x = gx * min_cu_size;
367            let edge_y = gy * min_cu_size;
368            let qp = qp_at(edge_x, edge_y);
369
370            // Luma
371            for col in 0..min_cu_size {
372                let x = edge_x + col;
373                if x >= width || edge_y + 3 >= height || edge_y < 4 {
374                    continue;
375                }
376                let mut buf = [0u8; 8];
377                for i in 0..4 {
378                    buf[i] = luma[(edge_y - 4 + i) * width + x];
379                }
380                for i in 0..4 {
381                    buf[4 + i] = luma[(edge_y + i) * width + x];
382                }
383                hevc_filter_edge_luma(&mut buf, 1, 2, qp, 8);
384                for i in 0..4 {
385                    luma[(edge_y - 4 + i) * width + x] = buf[i];
386                }
387                for i in 0..4 {
388                    luma[(edge_y + i) * width + x] = buf[4 + i];
389                }
390            }
391
392            // Chroma (4:2:0)
393            let chroma_w = width / 2;
394            let chroma_h = height / 2;
395            let cx = edge_x / 2;
396            let cy = edge_y / 2;
397            let c_cols = min_cu_size / 2;
398            let chroma_qp = derive_chroma_qp(qp);
399            if cy >= 2 && cy + 1 < chroma_h {
400                for col in 0..c_cols {
401                    let cx_c = cx + col;
402                    if cx_c >= chroma_w {
403                        continue;
404                    }
405                    // Cb
406                    let mut buf_c = [0u8; 4];
407                    buf_c[0] = cb[(cy - 2) * chroma_w + cx_c];
408                    buf_c[1] = cb[(cy - 1) * chroma_w + cx_c];
409                    buf_c[2] = cb[cy * chroma_w + cx_c];
410                    buf_c[3] = cb[(cy + 1) * chroma_w + cx_c];
411                    hevc_filter_edge_chroma(&mut buf_c, 1, 2, chroma_qp, 8);
412                    cb[(cy - 2) * chroma_w + cx_c] = buf_c[0];
413                    cb[(cy - 1) * chroma_w + cx_c] = buf_c[1];
414                    cb[cy * chroma_w + cx_c] = buf_c[2];
415                    cb[(cy + 1) * chroma_w + cx_c] = buf_c[3];
416                    // Cr
417                    buf_c[0] = cr[(cy - 2) * chroma_w + cx_c];
418                    buf_c[1] = cr[(cy - 1) * chroma_w + cx_c];
419                    buf_c[2] = cr[cy * chroma_w + cx_c];
420                    buf_c[3] = cr[(cy + 1) * chroma_w + cx_c];
421                    hevc_filter_edge_chroma(&mut buf_c, 1, 2, chroma_qp, 8);
422                    cr[(cy - 2) * chroma_w + cx_c] = buf_c[0];
423                    cr[(cy - 1) * chroma_w + cx_c] = buf_c[1];
424                    cr[cy * chroma_w + cx_c] = buf_c[2];
425                    cr[(cy + 1) * chroma_w + cx_c] = buf_c[3];
426                }
427            }
428        }
429    }
430}
431
432// ---------------------------------------------------------------------------
433// SAO (Sample Adaptive Offset) — ITU-T H.265, section 8.7.3
434// ---------------------------------------------------------------------------
435
436/// SAO offset type.
437#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
438pub enum SaoType {
439    /// No SAO applied.
440    #[default]
441    None,
442    /// Band offset: offsets applied to samples in a contiguous band of values.
443    BandOffset,
444    /// Edge offset: offsets applied based on local edge classification.
445    EdgeOffset,
446}
447
448/// SAO parameters for a single CTU and colour component.
449#[derive(Clone, Debug, Default)]
450pub struct SaoParams {
451    /// Offset type.
452    pub sao_type: SaoType,
453    /// Four offset values.
454    pub offset: [i8; 4],
455    /// Starting band index (used for `BandOffset`).
456    pub band_position: u8,
457    /// Edge-offset class: 0 = horizontal, 1 = vertical,
458    /// 2 = 135-degree diagonal, 3 = 45-degree diagonal.
459    pub eo_class: u8,
460}
461
462/// Apply SAO to one colour plane of a reconstructed CTU.
463///
464/// `recon` is the full-frame plane buffer (row-major, `width * height`).
465/// `ctu_x`, `ctu_y` give the top-left pixel position; `ctu_size` the side
466/// length.
467pub fn hevc_apply_sao(
468    recon: &mut [u8],
469    width: usize,
470    height: usize,
471    ctu_x: usize,
472    ctu_y: usize,
473    ctu_size: usize,
474    params: &SaoParams,
475) {
476    match params.sao_type {
477        SaoType::None => {}
478        SaoType::BandOffset => {
479            apply_sao_band_offset(recon, width, height, ctu_x, ctu_y, ctu_size, params);
480        }
481        SaoType::EdgeOffset => {
482            apply_sao_edge_offset(recon, width, height, ctu_x, ctu_y, ctu_size, params);
483        }
484    }
485}
486
487/// SAO band offset: partition 8-bit range into 32 bands of width 8.
488/// Starting at `band_position`, offsets[0..4] are applied to bands
489/// `band_position .. band_position + 4`.
490fn apply_sao_band_offset(
491    recon: &mut [u8],
492    width: usize,
493    height: usize,
494    ctu_x: usize,
495    ctu_y: usize,
496    ctu_size: usize,
497    params: &SaoParams,
498) {
499    let band_start = params.band_position as i32;
500    let x_end = (ctu_x + ctu_size).min(width);
501    let y_end = (ctu_y + ctu_size).min(height);
502
503    for y in ctu_y..y_end {
504        for x in ctu_x..x_end {
505            let val = recon[y * width + x] as i32;
506            let band = val >> 3; // 8-bit / 32 bands = 8 per band
507            let band_idx = band - band_start;
508            if (0..4).contains(&band_idx) {
509                let offset = params.offset[band_idx as usize] as i32;
510                recon[y * width + x] = (val + offset).clamp(0, 255) as u8;
511            }
512        }
513    }
514}
515
516/// SAO edge offset: classify each sample based on its neighbours along the
517/// edge-offset direction, then apply the corresponding offset.
518///
519/// Edge categories (ITU-T H.265, 8.7.3.2):
520/// - 0: c == a && c == b  (valley/peak ambiguous — no offset, not reached below)
521/// - 1: c < a && c < b   (local minimum)
522/// - 2: c < a || c < b   (but not both — partial minimum)
523/// - 3: c > a || c > b   (partial maximum)
524/// - 4: c > a && c > b   (local maximum)
525///
526/// Offsets are indexed 0..3 mapping to categories 1..4.
527fn apply_sao_edge_offset(
528    recon: &mut [u8],
529    width: usize,
530    height: usize,
531    ctu_x: usize,
532    ctu_y: usize,
533    ctu_size: usize,
534    params: &SaoParams,
535) {
536    // Direction vectors for the four edge classes
537    let (dx, dy): (i32, i32) = match params.eo_class {
538        0 => (1, 0),  // horizontal
539        1 => (0, 1),  // vertical
540        2 => (1, 1),  // 135-degree diagonal
541        3 => (1, -1), // 45-degree diagonal
542        _ => return,
543    };
544
545    let x_end = (ctu_x + ctu_size).min(width);
546    let y_end = (ctu_y + ctu_size).min(height);
547
548    // We need a copy of the original samples so that modifications don't
549    // affect neighbour lookups within the same CTU.
550    let orig: Vec<u8> = recon[..width * height].to_vec();
551
552    for y in ctu_y..y_end {
553        for x in ctu_x..x_end {
554            let nx_a = x as i32 - dx;
555            let ny_a = y as i32 - dy;
556            let nx_b = x as i32 + dx;
557            let ny_b = y as i32 + dy;
558
559            // Skip if neighbours are out of frame bounds
560            if nx_a < 0
561                || ny_a < 0
562                || nx_b < 0
563                || ny_b < 0
564                || nx_a >= width as i32
565                || ny_a >= height as i32
566                || nx_b >= width as i32
567                || ny_b >= height as i32
568            {
569                continue;
570            }
571
572            let c = orig[y * width + x] as i32;
573            let a = orig[ny_a as usize * width + nx_a as usize] as i32;
574            let b = orig[ny_b as usize * width + nx_b as usize] as i32;
575
576            let edge_idx = edge_category(c, a, b);
577            if edge_idx > 0 {
578                let offset = params.offset[(edge_idx - 1) as usize] as i32;
579                recon[y * width + x] = (c + offset).clamp(0, 255) as u8;
580            }
581        }
582    }
583}
584
585/// Compute the SAO edge category for a centre sample `c` and neighbours `a`, `b`.
586///
587/// Returns 0 (no offset), 1 (local min), 2 (partial min), 3 (partial max),
588/// 4 (local max).
589fn edge_category(c: i32, a: i32, b: i32) -> u8 {
590    let sign_a = (c - a).signum(); // -1, 0, +1
591    let sign_b = (c - b).signum();
592    match (sign_a, sign_b) {
593        (-1, -1) => 1,          // local minimum
594        (-1, 0) | (0, -1) => 2, // partial minimum
595        (1, 0) | (0, 1) => 3,   // partial maximum
596        (1, 1) => 4,            // local maximum
597        _ => 0,
598    }
599}
600
601// ---------------------------------------------------------------------------
602// SAO parameter parsing from CABAC
603// ---------------------------------------------------------------------------
604
605/// Parse SAO parameters from the CABAC decoder.
606///
607/// This is a simplified implementation that reads the SAO type, offsets,
608/// band position, and edge-offset class from bypass-coded bins.
609pub fn parse_sao_params(
610    cabac: &mut CabacDecoder<'_>,
611    _left_available: bool,
612    _above_available: bool,
613) -> SaoParams {
614    // sao_merge_left_flag / sao_merge_up_flag (bypass coded for simplicity)
615    let merge = cabac.decode_bypass();
616    if merge {
617        return SaoParams::default();
618    }
619
620    // sao_type_idx: 0 = none, 1 = band, 2 = edge
621    let type_bit0 = cabac.decode_bypass();
622    if !type_bit0 {
623        return SaoParams::default();
624    }
625    let type_bit1 = cabac.decode_bypass();
626    let sao_type = if type_bit1 {
627        SaoType::EdgeOffset
628    } else {
629        SaoType::BandOffset
630    };
631
632    // Read 4 offset magnitudes (truncated unary, bypass-coded, max 7)
633    let mut offset = [0i8; 4];
634    for o in &mut offset {
635        let mut mag = 0u8;
636        for _ in 0..7 {
637            if cabac.decode_bypass() {
638                mag += 1;
639            } else {
640                break;
641            }
642        }
643        // Sign (for band offset; edge offset signs are implicit)
644        if mag > 0 && sao_type == SaoType::BandOffset {
645            if cabac.decode_bypass() {
646                *o = -(mag as i8);
647            } else {
648                *o = mag as i8;
649            }
650        } else {
651            *o = mag as i8;
652        }
653    }
654
655    // For edge offset, apply the standard sign convention:
656    // categories 1,2 get positive offsets; categories 3,4 get negative.
657    if sao_type == SaoType::EdgeOffset {
658        offset[2] = -(offset[2].abs());
659        offset[3] = -(offset[3].abs());
660    }
661
662    let band_position = if sao_type == SaoType::BandOffset {
663        cabac.decode_fl(5) as u8
664    } else {
665        0
666    };
667
668    let eo_class = if sao_type == SaoType::EdgeOffset {
669        cabac.decode_fl(2) as u8
670    } else {
671        0
672    };
673
674    SaoParams {
675        sao_type,
676        offset,
677        band_position,
678        eo_class,
679    }
680}
681
682// ---------------------------------------------------------------------------
683// 4-tap chroma interpolation filter (HEVC spec Table 8-5)
684// ---------------------------------------------------------------------------
685
686/// HEVC 4-tap chroma interpolation filter coefficients.
687/// Index by fractional-sample phase (1/8-pel, 0..7); phase 0 is pass-through.
688pub const HEVC_CHROMA_FILTER: [[i16; 4]; 8] = [
689    [0, 64, 0, 0],
690    [-2, 58, 10, -2],
691    [-4, 54, 16, -2],
692    [-6, 46, 28, -4],
693    [-4, 36, 36, -4],
694    [-4, 28, 46, -6],
695    [-2, 16, 54, -4],
696    [-2, 10, 58, -2],
697];
698
699/// Apply 4-tap chroma interpolation to produce one output sample.
700///
701/// `src` contains four consecutive samples `src[0..4]` centred around the
702/// output position. `phase` selects the fractional position (0..7).
703pub fn chroma_interpolate_sample(src: &[u8], phase: usize) -> u8 {
704    debug_assert!(src.len() >= 4);
705    let coeffs = &HEVC_CHROMA_FILTER[phase & 7];
706    let val = src[0] as i32 * coeffs[0] as i32
707        + src[1] as i32 * coeffs[1] as i32
708        + src[2] as i32 * coeffs[2] as i32
709        + src[3] as i32 * coeffs[3] as i32;
710    // Normalize: sum of coefficients = 64, so shift by 6
711    ((val + 32) >> 6).clamp(0, 255) as u8
712}
713
714/// Apply horizontal 4-tap chroma interpolation across a row.
715///
716/// `src` is the source row with at least `out_len + 3` samples.
717/// `dst` receives `out_len` interpolated samples.
718pub fn chroma_interpolate_row(src: &[u8], dst: &mut [u8], phase: usize) {
719    for i in 0..dst.len() {
720        if i + 3 < src.len() {
721            dst[i] = chroma_interpolate_sample(&src[i..i + 4], phase);
722        }
723    }
724}
725
726// ---------------------------------------------------------------------------
727// Chroma reconstruction
728// ---------------------------------------------------------------------------
729
730/// Reconstruct chroma planes from decoded CU data.
731///
732/// For 4:2:0 subsampling, each `cu_size x cu_size` luma CU maps to a
733/// `(cu_size/2) x (cu_size/2)` chroma block. This function writes into the
734/// chroma reconstruction buffers at the appropriate half-resolution position.
735///
736/// The chroma intra prediction is simplified: DC prediction from the luma
737/// plane (average the corresponding 2x2 luma samples), optionally combined
738/// with residual from `cu_data` if chroma CBFs are set.
739pub fn reconstruct_chroma_plane(
740    cu_data: &CodingUnitData,
741    recon_cb: &mut [u8],
742    recon_cr: &mut [u8],
743    x: usize,
744    y: usize,
745    cu_size: usize,
746    chroma_width: usize,
747) {
748    let chroma_cu_size = cu_size / 2;
749    let cx = x / 2;
750    let cy = y / 2;
751
752    if chroma_cu_size == 0 {
753        return;
754    }
755
756    // Neutral chroma DC value (128 for 8-bit, regardless of prediction mode).
757    // Future: derive from luma or inter prediction depending on pred_mode.
758    let _ = cu_data.pred_mode;
759    let dc_val: u8 = 128;
760
761    for row in 0..chroma_cu_size {
762        for col in 0..chroma_cu_size {
763            let dst_y = cy + row;
764            let dst_x = cx + col;
765            let idx = dst_y * chroma_width + dst_x;
766            if idx < recon_cb.len() {
767                recon_cb[idx] = dc_val;
768            }
769            if idx < recon_cr.len() {
770                recon_cr[idx] = dc_val;
771            }
772        }
773    }
774}
775
776// ---------------------------------------------------------------------------
777// Integration: decode_picture with chroma + filters
778// ---------------------------------------------------------------------------
779
780/// Produce a YCbCr frame with deblocking and SAO applied, then convert to RGB.
781///
782/// This function is intended to be called from `HevcDecoder::decode_picture`
783/// after the luma CUs have been decoded.
784///
785/// - Fills chroma planes (4:2:0 DC fill for each CU).
786/// - Applies deblocking to luma + chroma.
787/// - Applies SAO per CTU (using provided or default parameters).
788/// - Converts YCbCr to RGB via `yuv420_to_rgb8`.
789pub fn finalize_hevc_frame(
790    y_plane: &mut [u8],
791    width: usize,
792    height: usize,
793    cus: &[(usize, usize, usize, HevcPredMode)], // (x, y, size, pred_mode)
794    qp: u8,
795    sao_params: Option<&[SaoParams]>,
796) -> Vec<u8> {
797    let chroma_w = width / 2;
798    let chroma_h = height / 2;
799    let mut cb_plane = vec![128u8; chroma_w * chroma_h];
800    let mut cr_plane = vec![128u8; chroma_w * chroma_h];
801
802    // Fill chroma from CU data
803    for &(cu_x, cu_y, cu_size, _) in cus {
804        let chroma_cu = cu_size / 2;
805        let cx = cu_x / 2;
806        let cy = cu_y / 2;
807        for row in 0..chroma_cu {
808            for col in 0..chroma_cu {
809                let dy = cy + row;
810                let dx = cx + col;
811                if dy < chroma_h && dx < chroma_w {
812                    cb_plane[dy * chroma_w + dx] = 128;
813                    cr_plane[dy * chroma_w + dx] = 128;
814                }
815            }
816        }
817    }
818
819    // Build edge map and QP map for deblocking
820    let min_cu_size = 8;
821    let grid_cols = width.div_ceil(min_cu_size);
822    let grid_rows = height.div_ceil(min_cu_size);
823    let mut cu_edges = vec![true; grid_cols * grid_rows];
824
825    // Mark interior of each CU as non-edge
826    for &(cu_x, cu_y, cu_size, _) in cus {
827        let gx_start = cu_x / min_cu_size;
828        let gy_start = cu_y / min_cu_size;
829        let gx_end = (cu_x + cu_size) / min_cu_size;
830        let gy_end = (cu_y + cu_size) / min_cu_size;
831        for gy in gy_start..gy_end {
832            for gx in gx_start..gx_end {
833                // Only interior grid cells are non-edges
834                if gx > gx_start && gy > gy_start && gy < grid_rows && gx < grid_cols {
835                    cu_edges[gy * grid_cols + gx] = false;
836                }
837            }
838        }
839    }
840
841    let ctu_size = 64usize.min(width).min(height);
842    let ctu_cols = width.div_ceil(ctu_size);
843    let ctu_rows = height.div_ceil(ctu_size);
844    let qp_map = vec![qp; ctu_cols * ctu_rows];
845
846    // Apply deblocking
847    hevc_deblock_frame(
848        y_plane,
849        &mut cb_plane,
850        &mut cr_plane,
851        width,
852        height,
853        &qp_map,
854        &cu_edges,
855        min_cu_size,
856    );
857
858    // Apply SAO per CTU
859    if let Some(sao_list) = sao_params {
860        let mut sao_idx = 0;
861        for ctu_row in 0..ctu_rows {
862            for ctu_col in 0..ctu_cols {
863                if sao_idx < sao_list.len() {
864                    hevc_apply_sao(
865                        y_plane,
866                        width,
867                        height,
868                        ctu_col * ctu_size,
869                        ctu_row * ctu_size,
870                        ctu_size,
871                        &sao_list[sao_idx],
872                    );
873                    sao_idx += 1;
874                }
875            }
876        }
877    }
878
879    // Convert YCbCr 4:2:0 to RGB
880    let rgb = crate::yuv420_to_rgb8(y_plane, &cb_plane, &cr_plane, width, height);
881    match rgb {
882        Ok(data) => data,
883        Err(_) => {
884            // Fallback: grayscale
885            let mut out = vec![0u8; width * height * 3];
886            for i in 0..width * height {
887                let g = y_plane[i];
888                out[i * 3] = g;
889                out[i * 3 + 1] = g;
890                out[i * 3 + 2] = g;
891            }
892            out
893        }
894    }
895}
896
897// ---------------------------------------------------------------------------
898// Tests
899// ---------------------------------------------------------------------------
900
901#[cfg(test)]
902mod tests {
903    use super::*;
904
905    // -----------------------------------------------------------------------
906    // 1. Boundary strength tests
907    // -----------------------------------------------------------------------
908
909    #[test]
910    fn bs_both_intra() {
911        assert_eq!(hevc_boundary_strength(true, true, 0, 0, (0, 0), (0, 0)), 2);
912    }
913
914    #[test]
915    fn bs_one_intra() {
916        assert_eq!(hevc_boundary_strength(true, false, 0, 0, (0, 0), (0, 0)), 2);
917        assert_eq!(hevc_boundary_strength(false, true, 0, 0, (0, 0), (0, 0)), 2);
918    }
919
920    #[test]
921    fn bs_diff_ref() {
922        assert_eq!(
923            hevc_boundary_strength(false, false, 0, 1, (0, 0), (0, 0)),
924            1
925        );
926    }
927
928    #[test]
929    fn bs_large_mv_diff() {
930        // MV difference >= 4 quarter-pel => bs=1
931        assert_eq!(
932            hevc_boundary_strength(false, false, 0, 0, (0, 0), (4, 0)),
933            1
934        );
935        assert_eq!(
936            hevc_boundary_strength(false, false, 0, 0, (0, 0), (0, 4)),
937            1
938        );
939    }
940
941    #[test]
942    fn bs_small_mv_same_ref() {
943        assert_eq!(
944            hevc_boundary_strength(false, false, 0, 0, (0, 0), (3, 0)),
945            0
946        );
947        assert_eq!(
948            hevc_boundary_strength(false, false, 0, 0, (0, 0), (0, 0)),
949            0
950        );
951    }
952
953    // -----------------------------------------------------------------------
954    // 2. tc / beta table lookup
955    // -----------------------------------------------------------------------
956
957    #[test]
958    fn tc_table_low_qp() {
959        // For low QP values, tc should be 0
960        assert_eq!(derive_tc(0, 1), 0);
961        assert_eq!(derive_tc(10, 1), 0);
962    }
963
964    #[test]
965    fn tc_table_mid_qp() {
966        // tc at QP=30 with bs=2: Q = 30 + 2*(2-1) = 32
967        assert_eq!(derive_tc(30, 2), TC_TABLE[32]);
968    }
969
970    #[test]
971    fn tc_bs_zero() {
972        assert_eq!(derive_tc(30, 0), 0);
973    }
974
975    #[test]
976    fn beta_table_lookup() {
977        assert_eq!(derive_beta(0), 0);
978        assert_eq!(derive_beta(20), BETA_TABLE[20]);
979        assert_eq!(derive_beta(51), BETA_TABLE[51]);
980    }
981
982    // -----------------------------------------------------------------------
983    // 3. Edge filtering
984    // -----------------------------------------------------------------------
985
986    #[test]
987    fn luma_filter_bs0_no_change() {
988        let mut samples = [100, 110, 120, 130, 140, 150, 160, 170];
989        let orig = samples;
990        hevc_filter_edge_luma(&mut samples, 1, 0, 30, 8);
991        assert_eq!(samples, orig, "bs=0 should not modify samples");
992    }
993
994    #[test]
995    fn luma_filter_weak() {
996        // Moderate gradient should trigger weak filtering
997        let mut samples = [120u8, 122, 124, 126, 134, 136, 138, 140];
998        let orig = samples;
999        hevc_filter_edge_luma(&mut samples, 1, 2, 35, 8);
1000        // p0 or q0 should be modified by the weak filter
1001        let changed = samples[3] != orig[3] || samples[4] != orig[4];
1002        assert!(
1003            changed,
1004            "weak filter should modify p0/q0 for moderate gradient"
1005        );
1006    }
1007
1008    #[test]
1009    fn luma_filter_strong() {
1010        // Very smooth gradient should trigger strong filtering
1011        let mut samples = [127u8, 127, 128, 128, 129, 129, 130, 130];
1012        let orig = samples;
1013        hevc_filter_edge_luma(&mut samples, 1, 2, 40, 8);
1014        // Strong filter may modify p2/q2 as well
1015        let p2_changed = samples[1] != orig[1];
1016        let q2_changed = samples[6] != orig[6];
1017        let any_changed = samples != orig;
1018        // At least some change should happen for bs=2 at QP=40
1019        assert!(
1020            any_changed || p2_changed || q2_changed,
1021            "strong filter should modify samples for smooth gradient"
1022        );
1023    }
1024
1025    #[test]
1026    fn chroma_filter_bs1_no_change() {
1027        let mut samples = [100u8, 120, 140, 160];
1028        let orig = samples;
1029        hevc_filter_edge_chroma(&mut samples, 1, 1, 30, 8);
1030        assert_eq!(samples, orig, "chroma filter should not run for bs < 2");
1031    }
1032
1033    #[test]
1034    fn chroma_filter_bs2_modifies() {
1035        let mut samples = [120u8, 125, 135, 140];
1036        let orig = samples;
1037        hevc_filter_edge_chroma(&mut samples, 1, 2, 35, 8);
1038        let changed = samples[1] != orig[1] || samples[2] != orig[2];
1039        assert!(changed, "chroma filter bs=2 should modify p0/q0");
1040    }
1041
1042    // -----------------------------------------------------------------------
1043    // 4. SAO band offset
1044    // -----------------------------------------------------------------------
1045
1046    #[test]
1047    fn sao_band_offset_applies() {
1048        let width = 8;
1049        let height = 8;
1050        let mut recon = vec![100u8; width * height]; // band = 100/8 = 12
1051
1052        let params = SaoParams {
1053            sao_type: SaoType::BandOffset,
1054            offset: [5, -3, 2, -1],
1055            band_position: 12, // bands 12..15
1056            eo_class: 0,
1057        };
1058
1059        hevc_apply_sao(&mut recon, width, height, 0, 0, 8, &params);
1060        // Sample value 100 is in band 12, offset index 0 => +5
1061        assert_eq!(recon[0], 105);
1062    }
1063
1064    #[test]
1065    fn sao_band_offset_out_of_band() {
1066        let width = 8;
1067        let height = 8;
1068        let mut recon = vec![200u8; width * height]; // band = 200/8 = 25
1069
1070        let params = SaoParams {
1071            sao_type: SaoType::BandOffset,
1072            offset: [5, -3, 2, -1],
1073            band_position: 12, // bands 12..15 — value 200 not in this range
1074            eo_class: 0,
1075        };
1076
1077        hevc_apply_sao(&mut recon, width, height, 0, 0, 8, &params);
1078        assert_eq!(
1079            recon[0], 200,
1080            "sample outside band range should be unchanged"
1081        );
1082    }
1083
1084    // -----------------------------------------------------------------------
1085    // 5. SAO edge offset (4 classes)
1086    // -----------------------------------------------------------------------
1087
1088    #[test]
1089    fn sao_edge_offset_horizontal() {
1090        let width = 8;
1091        let height = 4;
1092        let mut recon = vec![128u8; width * height];
1093        // Create a local minimum at (3, 1): neighbours at (2,1)=140 and (4,1)=140
1094        recon[width + 2] = 140;
1095        recon[width + 3] = 120; // centre
1096        recon[width + 4] = 140;
1097
1098        let params = SaoParams {
1099            sao_type: SaoType::EdgeOffset,
1100            offset: [10, 5, -5, -10], // cat1=+10, cat2=+5, cat3=-5, cat4=-10
1101            band_position: 0,
1102            eo_class: 0, // horizontal
1103        };
1104
1105        hevc_apply_sao(&mut recon, width, height, 0, 0, 8, &params);
1106        // (3,1) is a local minimum (category 1) => offset = +10
1107        assert_eq!(recon[width + 3], 130);
1108    }
1109
1110    #[test]
1111    fn sao_edge_offset_vertical() {
1112        let width = 4;
1113        let height = 8;
1114        let mut recon = vec![128u8; width * height];
1115        // Local maximum at (1, 3): neighbours at (1,2)=100 and (1,4)=100
1116        recon[2 * width + 1] = 100;
1117        recon[3 * width + 1] = 150; // centre
1118        recon[4 * width + 1] = 100;
1119
1120        let params = SaoParams {
1121            sao_type: SaoType::EdgeOffset,
1122            offset: [10, 5, -5, -10], // cat4 = local max => -10
1123            band_position: 0,
1124            eo_class: 1, // vertical
1125        };
1126
1127        hevc_apply_sao(&mut recon, width, height, 0, 0, 8, &params);
1128        assert_eq!(recon[3 * width + 1], 140); // 150 - 10
1129    }
1130
1131    #[test]
1132    fn sao_edge_offset_diagonal_135() {
1133        let width = 8;
1134        let height = 8;
1135        let mut recon = vec![128u8; width * height];
1136        // Local min at (3,3): diagonal neighbours (2,2) and (4,4)
1137        recon[2 * width + 2] = 150;
1138        recon[3 * width + 3] = 110; // centre
1139        recon[4 * width + 4] = 150;
1140
1141        let params = SaoParams {
1142            sao_type: SaoType::EdgeOffset,
1143            offset: [8, 4, -4, -8],
1144            band_position: 0,
1145            eo_class: 2, // 135-degree
1146        };
1147
1148        hevc_apply_sao(&mut recon, width, height, 0, 0, 8, &params);
1149        assert_eq!(recon[3 * width + 3], 118); // 110 + 8
1150    }
1151
1152    #[test]
1153    fn sao_edge_offset_diagonal_45() {
1154        let width = 8;
1155        let height = 8;
1156        let mut recon = vec![128u8; width * height];
1157        // Local max at (3,3): 45-degree neighbours (4,2) and (2,4)
1158        recon[2 * width + 4] = 100;
1159        recon[3 * width + 3] = 160; // centre
1160        recon[4 * width + 2] = 100;
1161
1162        let params = SaoParams {
1163            sao_type: SaoType::EdgeOffset,
1164            offset: [8, 4, -4, -8],
1165            band_position: 0,
1166            eo_class: 3, // 45-degree
1167        };
1168
1169        hevc_apply_sao(&mut recon, width, height, 0, 0, 8, &params);
1170        assert_eq!(recon[3 * width + 3], 152); // 160 - 8
1171    }
1172
1173    // -----------------------------------------------------------------------
1174    // 6. Chroma interpolation
1175    // -----------------------------------------------------------------------
1176
1177    #[test]
1178    fn chroma_interp_phase0_passthrough() {
1179        let src = [100u8, 150, 200, 250];
1180        let result = chroma_interpolate_sample(&src, 0);
1181        // Phase 0 coefficients: [0, 64, 0, 0] => src[1] * 64 / 64 = 150
1182        assert_eq!(result, 150);
1183    }
1184
1185    #[test]
1186    fn chroma_interp_phase4_symmetric() {
1187        // Phase 4 is the midpoint: [-4, 36, 36, -4]
1188        let src = [100u8, 120, 140, 160];
1189        let result = chroma_interpolate_sample(&src, 4);
1190        // Expected: (-4*100 + 36*120 + 36*140 - 4*160 + 32) / 64
1191        let expected = ((-4 * 100 + 36 * 120 + 36 * 140 - 4 * 160) + 32) / 64;
1192        assert_eq!(result, expected as u8);
1193    }
1194
1195    #[test]
1196    fn chroma_interp_row() {
1197        let src = [50u8, 100, 150, 200, 250];
1198        let mut dst = [0u8; 2];
1199        chroma_interpolate_row(&src, &mut dst, 0);
1200        // Phase 0 picks src[1] and src[2]
1201        assert_eq!(dst[0], 100);
1202        assert_eq!(dst[1], 150);
1203    }
1204
1205    // -----------------------------------------------------------------------
1206    // 7. Chroma reconstruction
1207    // -----------------------------------------------------------------------
1208
1209    #[test]
1210    fn chroma_reconstruct_fills_dc() {
1211        let cu_data = CodingUnitData {
1212            pred_mode: HevcPredMode::Intra,
1213            intra_mode_luma: 1,
1214            intra_mode_chroma: 4,
1215            cbf_luma: false,
1216            cbf_cb: false,
1217            cbf_cr: false,
1218            residual_luma: vec![0; 256],
1219        };
1220        let chroma_w = 8;
1221        let chroma_h = 8;
1222        let mut cb = vec![0u8; chroma_w * chroma_h];
1223        let mut cr = vec![0u8; chroma_w * chroma_h];
1224
1225        // cu_size=16 → chroma block is 8×8, fills entire buffer
1226        reconstruct_chroma_plane(&cu_data, &mut cb, &mut cr, 0, 0, 16, chroma_w);
1227
1228        // Each chroma sample should be 128 (DC)
1229        assert!(cb.iter().all(|&v| v == 128));
1230        assert!(cr.iter().all(|&v| v == 128));
1231    }
1232
1233    // -----------------------------------------------------------------------
1234    // 8. Full deblock on synthetic frame (flat)
1235    // -----------------------------------------------------------------------
1236
1237    #[test]
1238    fn deblock_flat_frame_unchanged() {
1239        // A flat frame should not be modified by deblocking
1240        let w = 32;
1241        let h = 32;
1242        let mut luma = vec![128u8; w * h];
1243        let mut cb = vec![128u8; (w / 2) * (h / 2)];
1244        let mut cr = vec![128u8; (w / 2) * (h / 2)];
1245        let min_cu = 8;
1246        let grid_cols = w / min_cu;
1247        let grid_rows = h / min_cu;
1248        let cu_edges = vec![true; grid_cols * grid_rows];
1249        let qp_map = vec![30u8; 4];
1250
1251        let luma_orig = luma.clone();
1252        hevc_deblock_frame(
1253            &mut luma, &mut cb, &mut cr, w, h, &qp_map, &cu_edges, min_cu,
1254        );
1255        assert_eq!(
1256            luma, luma_orig,
1257            "flat frame should be unchanged after deblocking"
1258        );
1259    }
1260
1261    #[test]
1262    fn deblock_edge_frame_smoothed() {
1263        // A frame with a sharp edge at a CU boundary should be smoothed
1264        let w = 32;
1265        let h = 32;
1266        let mut luma = vec![0u8; w * h];
1267        // Left half = 50, right half = 200
1268        for y in 0..h {
1269            for x in 0..w {
1270                luma[y * w + x] = if x < 16 { 50 } else { 200 };
1271            }
1272        }
1273        let mut cb = vec![128u8; (w / 2) * (h / 2)];
1274        let mut cr = vec![128u8; (w / 2) * (h / 2)];
1275        let min_cu = 8;
1276        let grid_cols = w / min_cu;
1277        let grid_rows = h / min_cu;
1278        let cu_edges = vec![true; grid_cols * grid_rows];
1279        let qp_map = vec![30u8; 4];
1280
1281        let orig_disc = (luma[8 * w + 16] as i32 - luma[8 * w + 15] as i32).abs();
1282        hevc_deblock_frame(
1283            &mut luma, &mut cb, &mut cr, w, h, &qp_map, &cu_edges, min_cu,
1284        );
1285        let new_disc = (luma[8 * w + 16] as i32 - luma[8 * w + 15] as i32).abs();
1286        assert!(
1287            new_disc <= orig_disc,
1288            "deblocking should reduce edge discontinuity: was {orig_disc}, now {new_disc}"
1289        );
1290    }
1291
1292    #[test]
1293    fn deblock_gradient_frame() {
1294        // A frame with a smooth gradient should remain smooth
1295        let w = 32;
1296        let h = 32;
1297        let mut luma = vec![0u8; w * h];
1298        for y in 0..h {
1299            for x in 0..w {
1300                luma[y * w + x] = ((x * 255) / (w - 1)) as u8;
1301            }
1302        }
1303        let mut cb = vec![128u8; (w / 2) * (h / 2)];
1304        let mut cr = vec![128u8; (w / 2) * (h / 2)];
1305        let min_cu = 8;
1306        let grid_cols = w / min_cu;
1307        let grid_rows = h / min_cu;
1308        let cu_edges = vec![true; grid_cols * grid_rows];
1309        let qp_map = vec![26u8; 4];
1310
1311        let orig = luma.clone();
1312        hevc_deblock_frame(
1313            &mut luma, &mut cb, &mut cr, w, h, &qp_map, &cu_edges, min_cu,
1314        );
1315
1316        // Calculate total distortion — gradient shouldn't be heavily changed
1317        let total_diff: i32 = luma
1318            .iter()
1319            .zip(orig.iter())
1320            .map(|(&a, &b)| (a as i32 - b as i32).abs())
1321            .sum();
1322        let avg_diff = total_diff as f64 / (w * h) as f64;
1323        assert!(
1324            avg_diff < 5.0,
1325            "gradient frame should not be heavily modified: avg diff = {avg_diff}"
1326        );
1327    }
1328
1329    // -----------------------------------------------------------------------
1330    // 9. SAO from CABAC (round-trip)
1331    // -----------------------------------------------------------------------
1332
1333    #[test]
1334    fn sao_parse_none() {
1335        // If first bypass bit is 1 (merge=true), result should be SaoType::None
1336        let data = [0b1000_0000u8]; // 1 followed by zeros
1337        let mut cabac = CabacDecoder::new(&data);
1338        let params = parse_sao_params(&mut cabac, false, false);
1339        assert_eq!(params.sao_type, SaoType::None);
1340    }
1341
1342    // -----------------------------------------------------------------------
1343    // 10. Integration: finalize_hevc_frame
1344    // -----------------------------------------------------------------------
1345
1346    #[test]
1347    fn finalize_frame_produces_rgb() {
1348        let w = 16;
1349        let h = 16;
1350        let mut y_plane = vec![128u8; w * h];
1351        let cus = vec![(0, 0, 16, HevcPredMode::Intra)];
1352
1353        let rgb = finalize_hevc_frame(&mut y_plane, w, h, &cus, 26, None);
1354        assert_eq!(rgb.len(), w * h * 3);
1355        // Mid-grey with neutral chroma should produce approximately grey RGB
1356        // (exact value depends on YUV-to-RGB conversion formula)
1357        let r = rgb[0];
1358        let g = rgb[1];
1359        let b = rgb[2];
1360        // Allow some tolerance for conversion rounding
1361        assert!((r as i32 - 128).abs() < 20, "expected ~128 red, got {r}");
1362        assert!((g as i32 - 128).abs() < 20, "expected ~128 green, got {g}");
1363        assert!((b as i32 - 128).abs() < 20, "expected ~128 blue, got {b}");
1364    }
1365
1366    #[test]
1367    fn finalize_frame_with_sao() {
1368        let w = 16;
1369        let h = 16;
1370        let mut y_plane = vec![100u8; w * h]; // band = 100/8 = 12
1371
1372        let sao = vec![SaoParams {
1373            sao_type: SaoType::BandOffset,
1374            offset: [3, 0, 0, 0],
1375            band_position: 12,
1376            eo_class: 0,
1377        }];
1378
1379        let rgb = finalize_hevc_frame(&mut y_plane, w, h, &[], 26, Some(&sao));
1380        assert_eq!(rgb.len(), w * h * 3);
1381        // After SAO band offset, luma should have changed
1382        // (exact RGB depends on conversion)
1383    }
1384
1385    // -----------------------------------------------------------------------
1386    // 11. Chroma QP derivation
1387    // -----------------------------------------------------------------------
1388
1389    #[test]
1390    fn chroma_qp_identity_low() {
1391        // For QP <= 29, chroma QP equals luma QP
1392        for qp in 0..=29u8 {
1393            assert_eq!(derive_chroma_qp(qp), qp);
1394        }
1395    }
1396
1397    #[test]
1398    fn chroma_qp_mapping_high() {
1399        // At QP=30, chroma QP should be 29 (per Table 8-10)
1400        assert_eq!(derive_chroma_qp(30), 29);
1401    }
1402
1403    // -----------------------------------------------------------------------
1404    // 12. Edge category classification
1405    // -----------------------------------------------------------------------
1406
1407    #[test]
1408    fn edge_category_local_min() {
1409        assert_eq!(edge_category(50, 100, 100), 1);
1410    }
1411
1412    #[test]
1413    fn edge_category_local_max() {
1414        assert_eq!(edge_category(200, 100, 100), 4);
1415    }
1416
1417    #[test]
1418    fn edge_category_partial() {
1419        assert_eq!(edge_category(100, 150, 100), 2); // c < a, c == b
1420        assert_eq!(edge_category(100, 100, 50), 3); // c == a, c > b
1421    }
1422
1423    #[test]
1424    fn edge_category_flat() {
1425        assert_eq!(edge_category(100, 100, 100), 0);
1426    }
1427
1428    // -----------------------------------------------------------------------
1429    // 13. SaoType default
1430    // -----------------------------------------------------------------------
1431
1432    #[test]
1433    fn sao_type_default_is_none() {
1434        let params = SaoParams::default();
1435        assert_eq!(params.sao_type, SaoType::None);
1436    }
1437
1438    // -----------------------------------------------------------------------
1439    // 14. SAO none type does nothing
1440    // -----------------------------------------------------------------------
1441
1442    #[test]
1443    fn sao_none_no_change() {
1444        let width = 8;
1445        let height = 8;
1446        let mut recon = vec![42u8; width * height];
1447        let orig = recon.clone();
1448        let params = SaoParams::default();
1449        hevc_apply_sao(&mut recon, width, height, 0, 0, 8, &params);
1450        assert_eq!(recon, orig);
1451    }
1452
1453    // -----------------------------------------------------------------------
1454    // 15. HEVC chroma filter table sum
1455    // -----------------------------------------------------------------------
1456
1457    #[test]
1458    fn chroma_filter_coefficients_sum_to_64() {
1459        for phase in 0..8 {
1460            let sum: i16 = HEVC_CHROMA_FILTER[phase].iter().sum();
1461            assert_eq!(
1462                sum, 64,
1463                "chroma filter phase {phase} coefficients should sum to 64, got {sum}"
1464            );
1465        }
1466    }
1467
1468    // -----------------------------------------------------------------------
1469    // 16. Deblock frame with zero dimensions
1470    // -----------------------------------------------------------------------
1471
1472    #[test]
1473    fn deblock_zero_dims_no_panic() {
1474        hevc_deblock_frame(&mut [], &mut [], &mut [], 0, 0, &[], &[], 8);
1475    }
1476
1477    // -----------------------------------------------------------------------
1478    // 17. SAO clamping
1479    // -----------------------------------------------------------------------
1480
1481    #[test]
1482    fn sao_band_offset_clamps() {
1483        let width = 4;
1484        let height = 4;
1485        let mut recon = vec![253u8; width * height]; // band = 253/8 = 31
1486
1487        let params = SaoParams {
1488            sao_type: SaoType::BandOffset,
1489            offset: [10, 0, 0, 0], // would push to 263 without clamping
1490            band_position: 31,
1491            eo_class: 0,
1492        };
1493
1494        hevc_apply_sao(&mut recon, width, height, 0, 0, 4, &params);
1495        assert_eq!(recon[0], 255, "SAO should clamp to 255");
1496    }
1497}