heic 0.1.4

//! HEVC conformance test suite
//!
//! Downloads and caches ITU-T JCTVC conformance bitstreams, decodes them
//! with our decoder and dec265 reference, then compares frame-by-frame.
//!
//! Run: cargo test --test conformance -- --nocapture
//! Run one: cargo test --test conformance AMVP_A -- --nocapture
//!
//! Vectors are cached in conformance/vectors/. First run downloads from ITU.
//! Reference YUV is generated by dec265 on first run.

use std::path::{Path, PathBuf};
use std::process::Command;

/// Base URL for ITU HEVC v1 conformance bitstreams
const ITU_BASE: &str =
    "https://www.itu.int/wftp3/av-arch/jctvc-site/bitstream_exchange/draft_conformance/HEVC_v1";

/// Path to dec265 reference decoder
const DEC265: &str = "/home/lilith/work/heic/libde265-src/build/dec265/dec265";

/// Conformance vectors directory
fn vectors_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("conformance/vectors")
}

/// Download and extract a conformance vector if not already cached
fn ensure_vector(name: &str) -> Option<PathBuf> {
    let dir = vectors_dir().join(name);
    let bitstream = find_bitstream(&dir);
    if bitstream.is_some() {
        return bitstream;
    }

    // Download
    let url = format!("{ITU_BASE}/{name}.zip");
    let zip_path = format!("/tmp/{name}.zip");
    eprintln!("Downloading {name}...");
    let status = Command::new("curl")
        .args(["-sSfL", &url, "-o", &zip_path])
        .status()
        .ok()?;
    if !status.success() {
        eprintln!("  FAILED to download {name}");
        return None;
    }

    // Extract
    std::fs::create_dir_all(&dir).ok()?;
    let status = Command::new("unzip")
        .args(["-o", "-q", &zip_path, "-d"])
        .arg(&dir)
        .status()
        .ok()?;
    std::fs::remove_file(&zip_path).ok();
    if !status.success() {
        eprintln!("  FAILED to extract {name}");
        return None;
    }

    find_bitstream(&dir)
}

/// Find the .bit/.bin/.265 bitstream file in a vector directory
fn find_bitstream(dir: &Path) -> Option<PathBuf> {
    if !dir.exists() {
        return None;
    }
    for entry in walkdir(dir) {
        if let Some(ext) = entry.extension()
            && (ext == "bit" || ext == "bin" || ext == "265")
        {
            return Some(entry);
        }
    }
    None
}

/// Simple recursive directory walk
fn walkdir(dir: &Path) -> Vec<PathBuf> {
    let mut result = Vec::new();
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                result.extend(walkdir(&path));
            } else {
                result.push(path);
            }
        }
    }
    result
}

/// Generate reference YUV using dec265 if not already cached
fn ensure_reference_yuv(name: &str, bitstream: &Path) -> Option<PathBuf> {
    let _dir = bitstream.parent()?;
    // Look for reference.yuv in the vector's top-level dir
    let top_dir = vectors_dir().join(name);
    let ref_path = top_dir.join("reference.yuv");
    if ref_path.exists() && std::fs::metadata(&ref_path).ok()?.len() > 0 {
        return Some(ref_path);
    }

    // Also check if the vector shipped with a .yuv
    for entry in walkdir(&top_dir) {
        if entry.extension().is_some_and(|e| e == "yuv") && !entry.to_string_lossy().contains("md5")
        {
            return Some(entry);
        }
    }

    // Use ffmpeg to generate display-ordered YUV reference
    // (dec265 -o outputs in decode order, not display order!)
    eprintln!("  Generating reference YUV with ffmpeg (display order)...");
    let status = Command::new("ffmpeg")
        .args(["-i"])
        .arg(bitstream)
        .args(["-pix_fmt", "yuv420p", "-f", "rawvideo", "-y"])
        .arg(&ref_path)
        .args(["-loglevel", "error"])
        .status()
        .ok()?;

    if status.success() && ref_path.exists() {
        Some(ref_path)
    } else {
        // Fallback to dec265
        if Path::new(DEC265).exists() {
            eprintln!("  ffmpeg failed, trying dec265 (WARNING: decode order)...");
            let status = Command::new(DEC265)
                .arg(bitstream)
                .args(["-o", ref_path.to_str()?])
                .arg("--quiet")
                .status()
                .ok()?;
            if status.success() && ref_path.exists() {
                return Some(ref_path);
            }
        }
        eprintln!("  Failed to generate reference for {name}");
        None
    }
}

/// Get stream dimensions via ffprobe
fn get_dimensions(bitstream: &Path) -> Option<(u32, u32)> {
    let output = Command::new("ffprobe")
        .args(["-v", "error", "-select_streams", "v:0"])
        .args(["-show_entries", "stream=width,height"])
        .args(["-of", "csv=p=0"])
        .arg(bitstream)
        .output()
        .ok()?;
    let s = String::from_utf8_lossy(&output.stdout);
    let parts: Vec<&str> = s.trim().split(',').collect();
    if parts.len() >= 2 {
        Some((parts[0].parse().ok()?, parts[1].parse().ok()?))
    } else {
        None
    }
}

/// Count frames by type via ffprobe
fn count_frame_types(bitstream: &Path) -> String {
    let output = Command::new("ffprobe")
        .args(["-v", "error", "-show_frames", "-select_streams", "v:0"])
        .args(["-show_entries", "frame=pict_type"])
        .args(["-of", "csv=p=0"])
        .arg(bitstream)
        .output()
        .ok();
    match output {
        Some(o) => {
            let s = String::from_utf8_lossy(&o.stdout);
            let types: Vec<&str> = s.trim().split('\n').collect();
            let i_count = types.iter().filter(|&&t| t == "I").count();
            let p_count = types.iter().filter(|&&t| t == "P").count();
            let b_count = types.iter().filter(|&&t| t == "B").count();
            format!(
                "{} frames ({}I {}P {}B)",
                types.len(),
                i_count,
                p_count,
                b_count
            )
        }
        None => "unknown".to_string(),
    }
}

/// Decode a bitstream with our decoder and return per-frame cropped Y planes.
fn decode_with_ours(bitstream: &Path) -> Result<Vec<Vec<u16>>, String> {
    let data = std::fs::read(bitstream).map_err(|e| format!("read: {e}"))?;
    let mut decoder = heic::VideoDecoder::new(16);
    let frames = decoder
        .decode_annex_b(&data)
        .map_err(|e| format!("decode: {e}"))?;
    // Extract cropped Y plane (conformance window applied)
    Ok(frames
        .into_iter()
        .map(|f| {
            let cw = f.cropped_width();
            let ch = f.cropped_height();
            let stride = f.width as usize;
            let mut cropped = Vec::with_capacity((cw * ch) as usize);
            for y in f.crop_top..(f.height - f.crop_bottom) {
                let row_start = y as usize * stride + f.crop_left as usize;
                let row_end = row_start + cw as usize;
                cropped.extend_from_slice(&f.y_plane[row_start..row_end]);
            }
            cropped
        })
        .collect())
}

/// Load reference YUV frames (420p 8-bit)
fn load_reference_yuv(path: &Path, width: u32, height: u32) -> Vec<Vec<u16>> {
    let data = std::fs::read(path).unwrap_or_default();
    let luma_size = (width * height) as usize;
    let chroma_size = (width / 2 * height / 2) as usize;
    let frame_size = luma_size + 2 * chroma_size; // YUV420

    let num_frames = data.len().checked_div(frame_size).unwrap_or(0);

    let mut frames = Vec::with_capacity(num_frames);
    for i in 0..num_frames {
        let offset = i * frame_size;
        let y_bytes = &data[offset..offset + luma_size];
        let y_plane: Vec<u16> = y_bytes.iter().map(|&b| b as u16).collect();
        frames.push(y_plane);
    }
    frames
}

/// Compare Y planes and return (psnr_db, num_diff_pixels, max_diff)
fn compare_y_planes(ours: &[u16], reference: &[u16], width: u32, height: u32) -> (f64, usize, u16) {
    let len = (width * height) as usize;
    let ours = &ours[..len.min(ours.len())];
    let reference = &reference[..len.min(reference.len())];

    let mut sse = 0u64;
    let mut num_diff = 0usize;
    let mut max_diff = 0u16;

    for (&a, &b) in ours.iter().zip(reference.iter()) {
        let diff = (a as i32 - b as i32).unsigned_abs() as u16;
        if diff > 0 {
            num_diff += 1;
            max_diff = max_diff.max(diff);
            sse += (diff as u64) * (diff as u64);
        }
    }

    let mse = if !ours.is_empty() {
        sse as f64 / ours.len() as f64
    } else {
        0.0
    };
    let psnr = if mse > 0.0 {
        10.0 * (255.0f64 * 255.0 / mse).log10()
    } else {
        f64::INFINITY
    };

    (psnr, num_diff, max_diff)
}

/// Run a single conformance test vector
fn run_conformance_test(name: &str) {
    eprintln!("\n=== {name} ===");

    let bitstream = match ensure_vector(name) {
        Some(b) => b,
        None => {
            eprintln!("  SKIP: could not obtain bitstream");
            return;
        }
    };

    let (width, height) = match get_dimensions(&bitstream) {
        Some(d) => d,
        None => {
            eprintln!("  SKIP: could not determine dimensions");
            return;
        }
    };

    let frame_types = count_frame_types(&bitstream);
    eprintln!("  {width}x{height}, {frame_types}");

    let ref_yuv = match ensure_reference_yuv(name, &bitstream) {
        Some(r) => r,
        None => {
            eprintln!("  SKIP: no reference YUV");
            return;
        }
    };

    // Decode with our decoder
    let our_frames = match decode_with_ours(&bitstream) {
        Ok(f) => f,
        Err(e) => {
            eprintln!("  FAIL: {e}");
            // This is useful info — how many frames did we get?
            return;
        }
    };

    // Load reference
    let ref_frames = load_reference_yuv(&ref_yuv, width, height);

    eprintln!(
        "  Decoded {} frames (ref: {} frames)",
        our_frames.len(),
        ref_frames.len()
    );

    // Compare frame by frame
    let compare_count = our_frames.len().min(ref_frames.len());
    let mut all_exact = true;
    let mut worst_psnr = f64::INFINITY;
    let mut total_diff_pixels = 0usize;

    for i in 0..compare_count {
        let (psnr, num_diff, max_diff) =
            compare_y_planes(&our_frames[i], &ref_frames[i], width, height);
        worst_psnr = worst_psnr.min(psnr);
        total_diff_pixels += num_diff;
        if num_diff > 0 {
            all_exact = false;
            if i < 5 || psnr < 40.0 {
                eprintln!(
                    "  Frame {i}: PSNR={psnr:.1}dB, {num_diff} diff pixels, max_diff={max_diff}"
                );
            }
        }
    }

    if compare_count > 0 {
        if all_exact {
            eprintln!("  PASS: all {compare_count} frames pixel-exact");
        } else {
            eprintln!(
                "  RESULT: worst PSNR={worst_psnr:.1}dB, total diff pixels={total_diff_pixels}"
            );
        }
    }

    if our_frames.len() < ref_frames.len() {
        eprintln!(
            "  WARNING: decoded only {} of {} reference frames",
            our_frames.len(),
            ref_frames.len()
        );
    }
}

// === Individual conformance tests ===

#[test]
fn girlshy() {
    // Local test: libde265 test bitstream (I, P, B frames)
    let bitstream = Path::new("/home/lilith/work/heic/libde265-src/testdata/girlshy.h265");
    if !bitstream.exists() {
        eprintln!("SKIP: girlshy.h265 not found");
        return;
    }
    let (width, height) = get_dimensions(bitstream).unwrap_or((316, 240));
    let frame_types = count_frame_types(bitstream);
    eprintln!("\n=== girlshy ===");
    eprintln!("  {width}x{height}, {frame_types}");

    // Generate reference YUV in display order using ffmpeg
    let ref_path = PathBuf::from("/tmp/girlshy_ref_display.yuv");
    if !ref_path.exists() {
        let _ = Command::new("ffmpeg")
            .args(["-i"])
            .arg(bitstream)
            .args(["-pix_fmt", "yuv420p", "-f", "rawvideo", "-y"])
            .arg(&ref_path)
            .args(["-loglevel", "error"])
            .status();
    }

    match decode_with_ours(bitstream) {
        Ok(our_frames) => {
            let ref_frames = load_reference_yuv(&ref_path, width, height);
            eprintln!(
                "  Decoded {} frames (ref: {} frames)",
                our_frames.len(),
                ref_frames.len()
            );
            let n = our_frames.len().min(ref_frames.len());
            for i in 0..n.min(15) {
                let (psnr, num_diff, max_diff) =
                    compare_y_planes(&our_frames[i], &ref_frames[i], width, height);
                let uninit = our_frames[i].iter().filter(|&&v| v == u16::MAX).count();
                let total_px = width * height;
                eprintln!(
                    "  Frame {i}: PSNR={psnr:.1}dB, diff={num_diff}/{total_px} max={max_diff}, uninit={uninit}"
                );
                // Show first differing pixel for frame 1
                if i == 1 && num_diff > 0 {
                    for j in 0..our_frames[i].len().min(ref_frames[i].len()) {
                        let a = our_frames[i][j];
                        let b = ref_frames[i][j];
                        if a != b {
                            let px = j as u32 % width;
                            let py = j as u32 / width;
                            eprintln!(
                                "    First diff at ({px},{py}): ours={a}, ref={b}, diff={}",
                                (a as i32 - b as i32).abs()
                            );
                            break;
                        }
                    }
                    // Also show pixel (0,0) comparison
                    if !our_frames[i].is_empty() && !ref_frames[i].is_empty() {
                        eprintln!(
                            "    (0,0): ours={}, ref={}",
                            our_frames[i][0], ref_frames[i][0]
                        );
                        eprintln!(
                            "    (100,100): ours={}, ref={}",
                            our_frames[i][(100 * width + 100) as usize],
                            ref_frames[i][(100 * width + 100) as usize]
                        );
                    }
                }
            }
        }
        Err(e) => eprintln!("  FAIL: {e}"),
    }
}

// Inter prediction vectors
#[test]
fn amvp_a() {
    run_conformance_test("AMVP_A_MTK_4");
}
#[test]
fn amvp_b() {
    run_conformance_test("AMVP_B_MTK_4");
}
#[test]
fn amvp_c() {
    run_conformance_test("AMVP_C_Samsung_7");
}
#[test]
fn merge_a() {
    run_conformance_test("MERGE_A_TI_3");
}
#[test]
fn merge_b() {
    run_conformance_test("MERGE_B_TI_3");
}
#[test]
fn merge_c() {
    run_conformance_test("MERGE_C_TI_3");
}
#[test]
fn merge_d() {
    run_conformance_test("MERGE_D_TI_3");
}
#[test]
fn merge_e() {
    run_conformance_test("MERGE_E_TI_3");
}
#[test]
fn merge_f() {
    run_conformance_test("MERGE_F_MTK_4");
}
#[test]
fn tmvp_a() {
    run_conformance_test("TMVP_A_MS_3");
}
#[test]
fn amp_a() {
    run_conformance_test("AMP_A_Samsung_7");
}
#[test]
fn amp_b() {
    run_conformance_test("AMP_B_Samsung_7");
}
#[test]
fn amp_d() {
    run_conformance_test("AMP_D_Hisilicon_3");
}
#[test]
fn pmerge_a() {
    run_conformance_test("PMERGE_A_TI_3");
}
#[test]
fn pmerge_b() {
    run_conformance_test("PMERGE_B_TI_3");
}
#[test]
fn pmerge_c() {
    run_conformance_test("PMERGE_C_TI_3");
}
#[test]
fn mvclip_a() {
    run_conformance_test("MVCLIP_A_qualcomm_3");
}
#[test]
fn mvdl1zero() {
    run_conformance_test("MVDL1ZERO_A_docomo_4");
}

// Reference picture sets
#[test]
fn rps_a() {
    run_conformance_test("RPS_A_docomo_5");
}
#[test]
fn rps_b() {
    run_conformance_test("RPS_B_qualcomm_5");
}
#[test]
fn rps_c() {
    run_conformance_test("RPS_C_ericsson_5");
}
#[test]
fn rap_a() {
    run_conformance_test("RAP_A_docomo_6");
}
#[test]
fn rap_b() {
    run_conformance_test("RAP_B_Bossen_2");
}

// Deblocking & SAO
#[test]
fn dblk_a() {
    run_conformance_test("DBLK_A_SONY_3");
}
#[test]
fn dblk_b() {
    run_conformance_test("DBLK_B_SONY_3");
}
#[test]
fn dblk_d() {
    run_conformance_test("DBLK_D_VIXS_2");
}
#[test]
fn sao_a() {
    run_conformance_test("SAO_A_MediaTek_4");
}
#[test]
fn sao_b() {
    run_conformance_test("SAO_B_MediaTek_5");
}

// Weighted prediction
#[test]
fn wp_a() {
    run_conformance_test("WP_A_Toshiba_3");
}
#[test]
fn wp_b() {
    run_conformance_test("WP_B_Toshiba_3");
}

// Intra & transform (should pass — existing I-slice support)
#[test]
fn ipred_a() {
    run_conformance_test("IPRED_A_Qualcomm_3");
}
#[test]
fn ipred_b() {
    run_conformance_test("IPRED_B_Nokia_3");
}
#[test]
fn rqt_a() {
    run_conformance_test("RQT_A_HHI_4");
}
#[test]
fn rqt_b() {
    run_conformance_test("RQT_B_HHI_4");
}
#[test]
fn sdh_a() {
    run_conformance_test("SDH_A_Orange_4");
}
#[test]
fn slist_a() {
    run_conformance_test("SLIST_A_Sony_5");
}

// CABAC init, structure, misc
#[test]
fn cainit_a() {
    run_conformance_test("CAINIT_A_SHARP_4");
}
#[test]
fn cainit_b() {
    run_conformance_test("CAINIT_B_SHARP_4");
}
#[test]
fn struct_a() {
    run_conformance_test("STRUCT_A_Samsung_7");
}
#[test]
fn struct_b() {
    run_conformance_test("STRUCT_B_Samsung_7");
}
#[test]
fn confwin_a() {
    run_conformance_test("CONFWIN_A_Sony_1");
}
#[test]
fn deltaqp_a() {
    run_conformance_test("DELTAQP_A_BRCM_4");
}
#[test]
fn slices_a() {
    run_conformance_test("SLICES_A_Rovi_3");
}
#[test]
fn tiles_a() {
    run_conformance_test("TILES_A_Cisco_2");
}
#[test]
fn tiles_b() {
    run_conformance_test("TILES_B_Cisco_1");
}
#[test]
fn poc_a() {
    run_conformance_test("POC_A_Bossen_3");
}
#[test]
fn wpp_a() {
    run_conformance_test("WPP_A_ericsson_MAIN_2");
}

/// Debug test: dump detailed diff locations for MERGE_A to find deblocking issues
#[test]
fn merge_a_deblock_debug() {
    let name = "MERGE_A_TI_3";
    let bitstream = match ensure_vector(name) {
        Some(b) => b,
        None => {
            eprintln!("SKIP");
            return;
        }
    };
    let (width, height) = get_dimensions(&bitstream).unwrap();
    let ref_yuv = match ensure_reference_yuv(name, &bitstream) {
        Some(r) => r,
        None => {
            eprintln!("SKIP: no ref");
            return;
        }
    };

    let data = std::fs::read(&bitstream).unwrap();
    let mut decoder = heic::VideoDecoder::new(16);
    let frames = decoder.decode_annex_b(&data).unwrap();
    let ref_frames = load_reference_yuv(&ref_yuv, width, height);

    for frame_idx in 0..frames.len().min(ref_frames.len()) {
        let f = &frames[frame_idx];
        let cw = f.cropped_width();
        let _ch = f.cropped_height();
        let stride = f.width as usize;
        let mut our_y = Vec::new();
        for y in f.crop_top..(f.height - f.crop_bottom) {
            let start = y as usize * stride + f.crop_left as usize;
            our_y.extend_from_slice(&f.y_plane[start..start + cw as usize]);
        }
        let ref_y = &ref_frames[frame_idx];

        let mut diffs: Vec<(u32, u32, i32)> = Vec::new();
        let len = (width * height) as usize;
        for i in 0..len.min(our_y.len()).min(ref_y.len()) {
            let d = our_y[i] as i32 - ref_y[i] as i32;
            if d != 0 {
                diffs.push((i as u32 % width, i as u32 / width, d));
            }
        }

        if diffs.is_empty() {
            eprintln!("Frame {frame_idx}: EXACT");
            continue;
        }

        let max_abs = diffs.iter().map(|d| d.2.abs()).max().unwrap();
        eprintln!(
            "Frame {frame_idx}: {} diffs, max_abs={max_abs}",
            diffs.len()
        );

        // Group diffs by nearest edge
        for (x, y, d) in &diffs {
            // For each diff pixel, identify which deblocking edge could be responsible:
            // Vertical edges at multiples of 8 affect pixels in range [edge-3, edge+3]
            // Horizontal edges at multiples of 8 affect pixels in range [edge-3, edge+3]
            let vert_edge = (*x / 8) * 8; // nearest lower vert edge
            let vert_edge_hi = vert_edge + 8; // nearest upper vert edge
            let horiz_edge = (*y / 8) * 8;
            let horiz_edge_hi = horiz_edge + 8;

            let near_vert = if (*x as i32 - vert_edge as i32).abs() <= 3 {
                Some(vert_edge)
            } else if (vert_edge_hi as i32 - *x as i32).abs() <= 3 {
                Some(vert_edge_hi)
            } else {
                None
            };
            let near_horiz = if (*y as i32 - horiz_edge as i32).abs() <= 3 {
                Some(horiz_edge)
            } else if (horiz_edge_hi as i32 - *y as i32).abs() <= 3 {
                Some(horiz_edge_hi)
            } else {
                None
            };

            eprintln!("  ({x},{y}) diff={d:+} vert_edge={near_vert:?} horiz_edge={near_horiz:?}");
        }
    }
}