medias 0.6.0

//! Dolby Vision RPUs of a file: HEVC NAL units of type 62 read through
//! libavformat (dlopen, demux only) and parsed by the dolby_vision crate.

use std::ffi::{CString, c_char, c_int, c_uint, c_void};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use dolby_vision::rpu::dovi_rpu::DoviRpu;
use dolby_vision::rpu::rpu_data_mapping::{DoviMappingMethod, RpuDataMapping};

use crate::vk::video;

/// Sonames tried, newest first: FFmpeg 9, 8, 7, 6, 5.
const SONAMES: [&std::ffi::CStr; 5] =
    [c"libavformat.so.63", c"libavformat.so.62", c"libavformat.so.61", c"libavformat.so.60", c"libavformat.so.59"];
const AVMEDIA_TYPE_VIDEO: c_int = 0;
const AV_NOPTS_VALUE: i64 = i64::MIN;
const AVSEEK_FLAG_BACKWARD: c_int = 1;
/// HEVC nal_unit_type of a Dolby Vision RPU (UNSPEC62).
const NAL_RPU: u8 = 62;

type Ctx = *mut c_void;

struct Lib {
    open_input: unsafe extern "C" fn(*mut Ctx, *const c_char, *const c_void, *mut c_void) -> c_int,
    find_stream_info: unsafe extern "C" fn(Ctx, *mut c_void) -> c_int,
    find_best_stream: unsafe extern "C" fn(Ctx, c_int, c_int, c_int, *mut c_void, c_int) -> c_int,
    read_frame: unsafe extern "C" fn(Ctx, *mut RawPacket) -> c_int,
    seek_frame: unsafe extern "C" fn(Ctx, c_int, i64, c_int) -> c_int,
    close_input: unsafe extern "C" fn(*mut Ctx),
    packet_alloc: unsafe extern "C" fn() -> *mut RawPacket,
    packet_unref: unsafe extern "C" fn(*mut RawPacket),
    packet_free: unsafe extern "C" fn(*mut *mut RawPacket),
}

/// Head of `AVPacket` (libavcodec/packet.h), unchanged since FFmpeg 3.
#[repr(C)]
struct RawPacket {
    buf: *mut c_void,
    pts: i64,
    dts: i64,
    data: *const u8,
    size: c_int,
    stream_index: c_int,
}

/// Head of `AVFormatContext` (libavformat 59 to 63).
#[repr(C)]
struct RawContext {
    av_class: *const c_void,
    iformat: *const c_void,
    oformat: *const c_void,
    priv_data: *mut c_void,
    pb: *mut c_void,
    ctx_flags: c_int,
    nb_streams: c_uint,
    streams: *const *const RawStream,
}

/// Head of `AVStream` (libavformat 59 to 63).
#[repr(C)]
struct RawStream {
    av_class: *const c_void,
    index: c_int,
    id: c_int,
    codecpar: *mut c_void,
    priv_data: *mut c_void,
    time_base: [c_int; 2],
}

impl Lib {
    fn load() -> Result<Lib, String> {
        let handle = SONAMES
            .iter()
            .map(|name| unsafe { libc::dlopen(name.as_ptr(), libc::RTLD_NOW | libc::RTLD_GLOBAL) })
            .find(|h| !h.is_null())
            .ok_or("libavformat.so.59 to .63 not found")?;
        let sym = |name: &std::ffi::CStr| {
            let p = unsafe { libc::dlsym(handle, name.as_ptr()) };
            if p.is_null() { Err(format!("libavformat lacks {name:?}")) } else { Ok(p) }
        };
        // SAFETY: signatures as declared in avformat.h and packet.h.
        unsafe {
            Ok(Lib {
                open_input: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(*mut Ctx, *const c_char, *const c_void, *mut c_void) -> c_int>(
                    sym(c"avformat_open_input")?,
                ),
                find_stream_info: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(Ctx, *mut c_void) -> c_int>(sym(c"avformat_find_stream_info")?),
                find_best_stream: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(Ctx, c_int, c_int, c_int, *mut c_void, c_int) -> c_int>(
                    sym(c"av_find_best_stream")?,
                ),
                read_frame: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(Ctx, *mut RawPacket) -> c_int>(sym(c"av_read_frame")?),
                seek_frame: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(Ctx, c_int, i64, c_int) -> c_int>(sym(c"av_seek_frame")?),
                close_input: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(*mut Ctx)>(sym(c"avformat_close_input")?),
                packet_alloc: std::mem::transmute::<*mut c_void, unsafe extern "C" fn() -> *mut RawPacket>(sym(c"av_packet_alloc")?),
                packet_unref: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(*mut RawPacket)>(sym(c"av_packet_unref")?),
                packet_free: std::mem::transmute::<*mut c_void, unsafe extern "C" fn(*mut *mut RawPacket)>(sym(c"av_packet_free")?),
            })
        }
    }
}

/// The RPU NAL unit (with its two header bytes) of one HEVC access unit:
/// 4-byte length prefixes (Matroska, MP4) or Annex B start codes.
pub fn rpu_nal(packet: &[u8]) -> Option<&[u8]> {
    let is_rpu = |nal: &[u8]| nal.len() > 2 && (nal[0] >> 1) & 0x3f == NAL_RPU;
    if packet.starts_with(&[0, 0, 1]) || packet.starts_with(&[0, 0, 0, 1]) {
        let starts: Vec<usize> = (0..packet.len().saturating_sub(2)).filter(|&i| packet[i..i + 3] == [0, 0, 1]).map(|i| i + 3).collect();
        return starts.iter().enumerate().find_map(|(k, &start)| {
            let end = starts.get(k + 1).map_or(packet.len(), |next| next - 3);
            Some(&packet[start..end]).filter(|nal| is_rpu(nal))
        });
    }
    let mut rest = packet;
    while rest.len() > 4 {
        let len = u32::from_be_bytes([rest[0], rest[1], rest[2], rest[3]]) as usize;
        let nal = rest.get(4..4 + len)?;
        if is_rpu(nal) {
            return Some(nal);
        }
        rest = &rest[4 + len..];
    }
    None
}

/// Demuxer over the video stream of one file.
pub struct Reader {
    lib: Lib,
    ctx: Ctx,
    packet: *mut RawPacket,
    stream: c_int,
    /// Seconds per pts unit.
    tick: f64,
}

// SAFETY: the context and packet are used by one thread at a time.
unsafe impl Send for Reader {}

impl Reader {
    pub fn open(url: &str) -> Result<Reader, String> {
        let lib = Lib::load()?;
        let c_url = CString::new(url).map_err(|_| "path with NUL")?;
        let mut ctx: Ctx = std::ptr::null_mut();
        unsafe {
            if (lib.open_input)(&mut ctx, c_url.as_ptr(), std::ptr::null(), std::ptr::null_mut()) < 0 {
                return Err(format!("libavformat cannot open {url}"));
            }
            let fail = |mut ctx: Ctx, message: String| {
                (lib.close_input)(&mut ctx);
                message
            };
            if (lib.find_stream_info)(ctx, std::ptr::null_mut()) < 0 {
                return Err(fail(ctx, "no stream information".into()));
            }
            let stream = (lib.find_best_stream)(ctx, AVMEDIA_TYPE_VIDEO, -1, -1, std::ptr::null_mut(), 0);
            let raw = &*ctx.cast::<RawContext>();
            if stream < 0 || stream as c_uint >= raw.nb_streams || raw.streams.is_null() {
                return Err(fail(ctx, "no video stream".into()));
            }
            let [num, den] = (**raw.streams.add(stream as usize)).time_base;
            // Guards the struct layouts above: a time base is a small positive fraction.
            if num <= 0 || den <= 0 || f64::from(num) / f64::from(den) > 1.0 {
                return Err(fail(ctx, format!("unexpected stream time base {num}/{den}")));
            }
            let packet = (lib.packet_alloc)();
            if packet.is_null() {
                return Err(fail(ctx, "av_packet_alloc failed".into()));
            }
            Ok(Reader { lib, ctx, packet, stream, tick: f64::from(num) / f64::from(den) })
        }
    }

    /// Next video access unit: presentation time in seconds and its RPU NAL
    /// unit, if it has one. `None` at the end of the file or on a read error.
    pub fn next(&mut self) -> Option<(f64, Option<Vec<u8>>)> {
        loop {
            unsafe {
                if (self.lib.read_frame)(self.ctx, self.packet) < 0 {
                    return None;
                }
                let p = &*self.packet;
                let unit = (p.stream_index == self.stream && !p.data.is_null() && p.size > 0).then(|| {
                    let pts = if p.pts != AV_NOPTS_VALUE { p.pts } else { p.dts };
                    let data = std::slice::from_raw_parts(p.data, p.size as usize);
                    (pts as f64 * self.tick, rpu_nal(data).map(<[u8]>::to_vec))
                });
                (self.lib.packet_unref)(self.packet);
                if unit.is_some() {
                    return unit;
                }
            }
        }
    }
}

impl Reader {
    /// Positions the reader at the keyframe at or before `seconds`.
    pub fn seek(&mut self, seconds: f64) -> Result<(), String> {
        let ts = (seconds.max(0.0) / self.tick) as i64;
        if unsafe { (self.lib.seek_frame)(self.ctx, self.stream, ts, AVSEEK_FLAG_BACKWARD) } < 0 {
            return Err(format!("seek to {seconds:.3} s failed"));
        }
        Ok(())
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            (self.lib.packet_free)(&mut self.packet);
            (self.lib.close_input)(&mut self.ctx);
        }
    }
}

/// One reshaping curve: pieces between ascending pivots (base-layer code / (2^depth - 1)),
/// each c0 + c1 x + c2 x^2 on the normalised code x.
#[derive(Clone, Debug, PartialEq)]
pub struct Curve {
    pivots: Vec<f32>,
    pieces: Vec<[f32; 3]>,
}

impl Curve {
    fn at(&self, x: f32) -> f32 {
        let x = x.clamp(self.pivots[0], self.pivots[self.pivots.len() - 1]);
        let piece = self.pivots[1..self.pivots.len() - 1].iter().take_while(|p| x >= **p).count();
        let [c0, c1, c2] = self.pieces[piece];
        c0 + x * (c1 + x * c2)
    }
}

/// Profile 5 conversion of one RPU: per-component polynomial reshaping to I, Ct, Cp,
/// `nonlinear` (ycc_to_rgb) to PQ L'M'S', PQ EOTF, `linear` (rgb_to_lms) to LMS.
#[derive(Clone, Debug, PartialEq)]
pub struct Mapping {
    curves: [Curve; 3],
    nonlinear: [[f32; 3]; 3],
    offset: [f32; 3],
    linear: [[f32; 3]; 3],
    /// Mastering peak, cd/m2 (source_max_pq).
    pub peak: f32,
}

/// BT.2020 RGB to XYZ (D65) and XYZ to Hunt-Pointer-Estevez LMS normalised to D65.
const BT2020_TO_XYZ: [[f64; 3]; 3] = [[0.636958, 0.144617, 0.168881], [0.262700, 0.677998, 0.059302], [0.0, 0.028073, 1.060985]];
const XYZ_TO_LMS: [[f64; 3]; 3] = [[0.4002, 0.7076, -0.0808], [-0.2263, 1.1653, 0.0457], [0.0, 0.0, 0.9182]];

fn mul(a: [[f64; 3]; 3], b: [[f64; 3]; 3]) -> [[f64; 3]; 3] {
    std::array::from_fn(|i| std::array::from_fn(|j| (0..3).map(|k| a[i][k] * b[k][j]).sum()))
}

fn inverse(m: [[f64; 3]; 3]) -> [[f64; 3]; 3] {
    let cof = |i: usize, j: usize| m[(i + 1) % 3][(j + 1) % 3] * m[(i + 2) % 3][(j + 2) % 3] - m[(i + 1) % 3][(j + 2) % 3] * m[(i + 2) % 3][(j + 1) % 3];
    let det: f64 = (0..3).map(|j| m[0][j] * cof(0, j)).sum();
    std::array::from_fn(|i| std::array::from_fn(|j| cof(j, i) / det))
}

/// LMS (as in `XYZ_TO_LMS`) to linear BT.2020 RGB.
fn lms_to_bt2020() -> [[f32; 3]; 3] {
    inverse(mul(XYZ_TO_LMS, BT2020_TO_XYZ)).map(|r| r.map(|v| v as f32))
}

impl Mapping {
    /// `Err` names what this pass does not implement (MMR, float coefficients, a missing part).
    pub fn from_rpu(rpu: &DoviRpu) -> Result<Mapping, String> {
        let header = &rpu.header;
        if header.coefficient_data_type != 0 {
            return Err("floating-point coefficients".into());
        }
        let mapping = rpu.rpu_data_mapping.as_ref().ok_or("an RPU without mapping data")?;
        let dm = rpu.vdr_dm_data.as_ref().ok_or("an RPU without display management data")?;
        let denom = (1u64 << header.coefficient_log2_denom) as f32;
        let code_max = ((1u32 << (header.bl_bit_depth_minus8 + 8)) - 1) as f32;
        let mut curves = Vec::with_capacity(3);
        for curve in &mapping.curves {
            let poly = match (curve.mapping_idc, &curve.polynomial) {
                (DoviMappingMethod::Polynomial, Some(poly)) => poly,
                _ => return Err("MMR reshaping".into()),
            };
            // Pivots are coded as differences.
            let mut sum = 0u32;
            let pivots: Vec<f32> = curve
                .pivots
                .iter()
                .map(|p| {
                    sum += u32::from(*p);
                    sum as f32 / code_max
                })
                .collect();
            let count = pivots.len().saturating_sub(1);
            if count == 0 || poly.poly_coef.len() < count || poly.poly_coef_int.len() < count || poly.linear_interp_flag.iter().any(|f| *f) {
                return Err("an unsupported polynomial curve".into());
            }
            let pieces = (0..count)
                .map(|i| {
                    let mut c = [0.0f32; 3];
                    for (k, (int, frac)) in poly.poly_coef_int[i].iter().zip(poly.poly_coef[i].iter()).enumerate().take(3) {
                        c[k] = *int as f32 + *frac as f32 / denom;
                    }
                    c
                })
                .collect();
            curves.push(Curve { pivots, pieces });
        }
        let curves: [Curve; 3] = curves.try_into().map_err(|_| "three components expected")?;
        let m = |v: [i16; 9], scale: f32| -> [[f32; 3]; 3] { std::array::from_fn(|i| std::array::from_fn(|j| f32::from(v[i * 3 + j]) / scale)) };
        Ok(Mapping {
            curves,
            nonlinear: m(
                [dm.ycc_to_rgb_coef0, dm.ycc_to_rgb_coef1, dm.ycc_to_rgb_coef2, dm.ycc_to_rgb_coef3, dm.ycc_to_rgb_coef4,
                 dm.ycc_to_rgb_coef5, dm.ycc_to_rgb_coef6, dm.ycc_to_rgb_coef7, dm.ycc_to_rgb_coef8],
                8192.0,
            ),
            offset: [dm.ycc_to_rgb_offset0, dm.ycc_to_rgb_offset1, dm.ycc_to_rgb_offset2].map(|v| v as f32 / (1u64 << 28) as f32),
            linear: m(
                [dm.rgb_to_lms_coef0, dm.rgb_to_lms_coef1, dm.rgb_to_lms_coef2, dm.rgb_to_lms_coef3, dm.rgb_to_lms_coef4,
                 dm.rgb_to_lms_coef5, dm.rgb_to_lms_coef6, dm.rgb_to_lms_coef7, dm.rgb_to_lms_coef8],
                16384.0,
            ),
            peak: video::pq_eotf(f32::from(dm.source_max_pq) / 4095.0),
        })
    }

    /// Base-layer codes (Y, Cb, Cr) / (2^depth - 1) to linear BT.2020 light in cd/m2.
    pub fn light(&self, codes: [f32; 3], to_rgb: &[[f32; 3]; 3]) -> [f32; 3] {
        let dot = |m: &[[f32; 3]; 3], v: [f32; 3]| m.map(|r| r[0] * v[0] + r[1] * v[1] + r[2] * v[2]);
        let ipt: [f32; 3] = std::array::from_fn(|k| self.curves[k].at(codes[k]) - self.offset[k]);
        let lms = dot(&self.linear, dot(&self.nonlinear, ipt).map(video::pq_eotf));
        dot(to_rgb, lms)
    }

    /// Table for `vk::video::Colour::Cube`: codes to SDR BT.709, or to PQ BT.2020 R'G'B' for `hdr`.
    pub fn table(&self, hdr: bool) -> video::Table {
        let to_rgb = lms_to_bt2020();
        let peak = self.peak.clamp(video::SDR_WHITE, 10000.0);
        video::Table::of(hdr, |codes| {
            let light = self.light(codes, &to_rgb);
            if hdr { light.map(|v| video::pq_oetf(v.clamp(0.0, 10000.0))) } else { video::tone_map_light(light, peak, video::SDR_WHITE) }
        })
    }
}

/// Seconds the reader stays ahead of the playback position.
const READ_AHEAD: f64 = 4.0;
/// Tables older than this many seconds before the position are dropped, except the newest of them.
const KEEP_BEHIND: f64 = 2.0;

struct Entry {
    pts: f64,
    id: u64,
    table: Arc<video::Table>,
}

#[derive(Default)]
struct Shared {
    position: f64,
    /// Ascending pts; an entry holds from its pts to the next entry's.
    entries: Vec<Entry>,
    error: Option<String>,
    stop: bool,
}

/// Background reader of a file's RPUs: keeps the colour tables around the playback position.
pub struct Feed {
    shared: Arc<Mutex<Shared>>,
}

impl Feed {
    /// `hdr`: tables for PQ BT.2020 output, else SDR BT.709.
    pub fn start(url: &str, hdr: bool) -> Feed {
        let shared = Arc::new(Mutex::new(Shared::default()));
        let (state, url) = (shared.clone(), url.to_owned());
        let spawned = std::thread::Builder::new().name("dovi".into()).spawn(move || {
            if let Err(error) = feed(&url, hdr, &state) {
                state.lock().unwrap().error = Some(error);
            }
        });
        if let Err(error) = spawned {
            shared.lock().unwrap().error = Some(format!("thread: {error}"));
        }
        Feed { shared }
    }

    /// Sets the playback position; returns the table for it (id, contents), if read yet.
    pub fn table(&self, position: f64) -> Result<Option<(u64, Arc<video::Table>)>, String> {
        let mut s = self.shared.lock().unwrap();
        if let Some(error) = &s.error {
            return Err(error.clone());
        }
        s.position = position;
        // Half a frame of slack at 60 fps for rounding of the position.
        Ok(s.entries.iter().rev().find(|e| e.pts <= position + 0.008).map(|e| (e.id, e.table.clone())))
    }
}

impl Drop for Feed {
    fn drop(&mut self) {
        // Not joined: the thread may be inside a blocking read.
        self.shared.lock().unwrap().stop = true;
    }
}

fn feed(url: &str, hdr: bool, shared: &Mutex<Shared>) -> Result<(), String> {
    let mut reader = Reader::open(url)?;
    // Read position, first pts after the last seek, last mapping in decode order.
    let (mut read, mut start, mut last): (f64, Option<f64>, Option<Mapping>) = (0.0, None, None);
    // Recent mappings with their tables: leading pictures after a scene cut repeat the previous one.
    let mut recent: Vec<(Mapping, u64, Arc<video::Table>)> = Vec::new();
    let mut next_id = 1u64;
    let mut ended = false;
    loop {
        let position = {
            let s = shared.lock().unwrap();
            if s.stop {
                return Ok(());
            }
            s.position
        };
        let outside = start.is_some_and(|first| position < first - 0.5) || position > read + READ_AHEAD;
        if outside {
            reader.seek(position - 1.0)?;
            shared.lock().unwrap().entries.clear();
            (read, start, last, ended) = (position, None, None, false);
        } else if ended || (start.is_some() && read > position + READ_AHEAD) {
            std::thread::sleep(Duration::from_millis(20));
            continue;
        }
        let Some((pts, nal)) = reader.next() else {
            ended = true;
            continue;
        };
        start.get_or_insert(pts);
        read = read.max(pts);
        let Some(nal) = nal else { continue };
        let rpu = DoviRpu::parse_unspec62_nalu(&nal).map_err(|e| format!("RPU at {pts:.3} s: {e}"))?;
        // An RPU without mapping data reuses the previous one.
        if rpu.rpu_data_mapping.is_none() && last.is_some() {
            continue;
        }
        let mapping = Mapping::from_rpu(&rpu).map_err(|e| format!("RPU at {pts:.3} s has {e}"))?;
        if last.as_ref() == Some(&mapping) {
            continue;
        }
        let (id, table) = match recent.iter().find(|(m, ..)| *m == mapping) {
            Some((_, id, table)) => (*id, table.clone()),
            None => {
                let entry = (next_id, Arc::new(mapping.table(hdr)));
                next_id += 1;
                recent.push((mapping.clone(), entry.0, entry.1.clone()));
                if recent.len() > 4 {
                    recent.remove(0);
                }
                entry
            }
        };
        last = Some(mapping);
        let mut s = shared.lock().unwrap();
        let at = s.entries.partition_point(|e| e.pts <= pts);
        s.entries.insert(at, Entry { pts, id, table });
        let old = s.entries.partition_point(|e| e.pts < s.position - KEEP_BEHIND);
        if old > 1 {
            s.entries.drain(..old - 1);
        }
    }
}

/// One line per component: mapping method, pieces and orders.
fn describe_mapping(mapping: &RpuDataMapping) -> String {
    let parts: Vec<String> = mapping
        .curves
        .iter()
        .map(|curve| {
            let pieces = curve.pivots.len().saturating_sub(1);
            match curve.mapping_idc {
                DoviMappingMethod::Polynomial => {
                    let orders: Vec<u64> = curve.polynomial.as_ref().map(|p| p.poly_order_minus1.iter().map(|o| o + 1).collect()).unwrap_or_default();
                    format!("polynomial, {pieces} pieces, orders {orders:?}")
                }
                DoviMappingMethod::MMR => {
                    let orders: Vec<u8> = curve.mmr.as_ref().map(|m| m.mmr_order_minus1.iter().map(|o| o + 1).collect()).unwrap_or_default();
                    format!("MMR, {pieces} pieces, orders {orders:?}")
                }
                DoviMappingMethod::Invalid => "invalid".into(),
            }
        })
        .collect();
    parts.join("; ")
}

/// `medias dv-info <file> [frames]`: prints the RPU structure of the first
/// `frames` access units and how often the mapping and the matrices change.
pub fn print_info(url: &str, frames: usize) -> Result<(), String> {
    let mut reader = Reader::open(url)?;
    let (mut units, mut with_rpu, mut errors) = (0usize, 0usize, 0usize);
    let (mut mapping_changes, mut matrix_changes, mut reused) = (0usize, 0usize, 0usize);
    let mut last_mapping: Option<String> = None;
    let mut last_matrix: Option<[i64; 21]> = None;
    while units < frames {
        let Some((pts, nal)) = reader.next() else { break };
        units += 1;
        let Some(nal) = nal else { continue };
        with_rpu += 1;
        let rpu = match DoviRpu::parse_unspec62_nalu(&nal) {
            Ok(rpu) => rpu,
            Err(error) => {
                errors += 1;
                if errors == 1 {
                    println!("frame {units} ({pts:.3} s): RPU parse error: {error}");
                }
                continue;
            }
        };
        if with_rpu == 1 {
            let h = &rpu.header;
            println!(
                "profile {}; base layer {} bit, reconstructed {} bit; coefficient type {} (log2 denominator {})",
                rpu.dovi_profile,
                h.bl_bit_depth_minus8 + 8,
                h.vdr_bit_depth_minus8 + 8,
                h.coefficient_data_type,
                h.coefficient_log2_denom,
            );
        }
        match &rpu.rpu_data_mapping {
            Some(mapping) => {
                let text = format!("{mapping:?}");
                if last_mapping.as_ref() != Some(&text) {
                    if mapping_changes < 3 {
                        println!("frame {units} ({pts:.3} s): {}", describe_mapping(mapping));
                    }
                    mapping_changes += 1;
                    last_mapping = Some(text);
                }
            }
            None => reused += 1,
        }
        if let Some(dm) = &rpu.vdr_dm_data {
            let matrix = [
                dm.ycc_to_rgb_coef0, dm.ycc_to_rgb_coef1, dm.ycc_to_rgb_coef2, dm.ycc_to_rgb_coef3, dm.ycc_to_rgb_coef4,
                dm.ycc_to_rgb_coef5, dm.ycc_to_rgb_coef6, dm.ycc_to_rgb_coef7, dm.ycc_to_rgb_coef8,
                dm.rgb_to_lms_coef0, dm.rgb_to_lms_coef1, dm.rgb_to_lms_coef2, dm.rgb_to_lms_coef3, dm.rgb_to_lms_coef4,
                dm.rgb_to_lms_coef5, dm.rgb_to_lms_coef6, dm.rgb_to_lms_coef7, dm.rgb_to_lms_coef8,
            ]
            .map(i64::from);
            let mut all = [0i64; 21];
            all[..18].copy_from_slice(&matrix);
            all[18..].copy_from_slice(&[dm.ycc_to_rgb_offset0, dm.ycc_to_rgb_offset1, dm.ycc_to_rgb_offset2].map(i64::from));
            if last_matrix != Some(all) {
                if matrix_changes < 3 {
                    println!(
                        "frame {units} ({pts:.3} s): ycc_to_rgb {:?} offsets {:?}; rgb_to_lms {:?}; source PQ {}..{}, full range {}",
                        &all[..9], &all[18..], &all[9..18], dm.source_min_pq, dm.source_max_pq, dm.signal_full_range_flag,
                    );
                }
                matrix_changes += 1;
                last_matrix = Some(all);
            }
        }
    }
    println!(
        "{units} access units, {with_rpu} with an RPU, {errors} parse errors; mapping: {mapping_changes} distinct in sequence, {reused} reuse the previous; matrices: {matrix_changes} distinct in sequence"
    );
    Ok(())
}

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

    #[test]
    fn rpu_nal_is_found_behind_length_prefixes_and_start_codes() {
        let slice = [0x26u8, 0x01, 0xaa, 0xbb]; // nal_unit_type 19
        let rpu = [0x7cu8, 0x01, 0x19, 0x08, 0x09];
        let mut prefixed = Vec::new();
        for nal in [&slice[..], &rpu[..]] {
            prefixed.extend((nal.len() as u32).to_be_bytes());
            prefixed.extend(nal);
        }
        assert_eq!(rpu_nal(&prefixed), Some(&rpu[..]));
        assert_eq!(rpu_nal(&prefixed[..8]), None);
        // A length running past the packet ends the scan.
        assert_eq!(rpu_nal(&[0, 0, 1, 0, 0x7c, 0x01]), None);

        let mut annexb = vec![0, 0, 0, 1];
        annexb.extend(slice);
        annexb.extend([0, 0, 1]);
        annexb.extend(rpu);
        assert_eq!(rpu_nal(&annexb), Some(&rpu[..]));
        annexb.extend([0, 0, 1, 0x26, 0x01]);
        assert_eq!(rpu_nal(&annexb), Some(&rpu[..]));
    }

    #[test]
    fn generated_profile_5_rpu_parses_from_a_packet() {
        use dolby_vision::rpu::generate::{GenerateConfig, GenerateProfile};
        let config = GenerateConfig { profile: GenerateProfile::Profile5, ..Default::default() };
        let nal = DoviRpu::profile5_config(&config).unwrap().write_hevc_unspec62_nalu().unwrap();
        let mut packet = (nal.len() as u32).to_be_bytes().to_vec();
        packet.extend(&nal);
        let rpu = DoviRpu::parse_unspec62_nalu(rpu_nal(&packet).unwrap()).unwrap();
        assert_eq!(rpu.dovi_profile, 5);
        assert!(rpu.rpu_data_mapping.is_some() && rpu.vdr_dm_data.is_some());
    }

    fn profile_5_rpu() -> DoviRpu {
        use dolby_vision::rpu::generate::{GenerateConfig, GenerateProfile};
        DoviRpu::profile5_config(&GenerateConfig { profile: GenerateProfile::Profile5, ..Default::default() }).unwrap()
    }

    #[test]
    fn lms_matrix_maps_white_to_white() {
        let m = lms_to_bt2020();
        for row in m {
            assert!((row.iter().sum::<f32>() - 1.0).abs() < 2e-3, "{m:?}");
        }
        // libplacebo's dovi_lms2rgb, first row.
        assert!((m[0][0] - 3.0644).abs() < 0.01 && (m[0][1] + 2.1660).abs() < 0.01, "{m:?}");
    }

    #[test]
    fn identity_mapping_decodes_grey_as_pq() {
        let mapping = Mapping::from_rpu(&profile_5_rpu()).unwrap();
        let to_rgb = lms_to_bt2020();
        // Neutral chroma: I is the PQ signal of the luminance.
        for nits in [1.0f32, 100.0, 1000.0] {
            let light = mapping.light([video::pq_oetf(nits), 0.5, 0.5], &to_rgb);
            for v in light {
                assert!((v / nits - 1.0).abs() < 0.03, "{nits}: {light:?}");
            }
        }
        // Positive Cp (third component) raises red against green.
        let light = mapping.light([0.5, 0.5, 0.6], &to_rgb);
        assert!(light[0] > light[1], "{light:?}");
    }

    #[test]
    fn curve_selects_the_piece_by_pivot() {
        let curve = Curve { pivots: vec![0.1, 0.5, 0.9], pieces: vec![[0.0, 1.0, 0.0], [1.0, 0.0, 0.0]] };
        assert_eq!(curve.at(0.3), 0.3);
        assert_eq!(curve.at(0.0), 0.1);
        assert_eq!(curve.at(0.5), 1.0);
        assert_eq!(curve.at(2.0), 1.0);
    }

    /// Writes one profile 5 RPU per frame as an Annex B file for x265's --dolby-vision-rpu:
    /// `MEDIAS_RPU_OUT=<path> cargo test --release -- --ignored write_rpu_file`.
    #[test]
    #[ignore]
    fn write_rpu_file() {
        let Ok(path) = std::env::var("MEDIAS_RPU_OUT") else { return };
        let nal = profile_5_rpu().write_hevc_unspec62_nalu().unwrap();
        let mut out = Vec::new();
        for _ in 0..48 {
            out.extend([0, 0, 0, 1]);
            out.extend(&nal);
        }
        std::fs::write(path, out).unwrap();
    }

    /// `MEDIAS_DV_FILE=<profile 5 file> cargo test --release -- --ignored --nocapture feed_delivers`.
    #[test]
    #[ignore]
    fn feed_delivers_a_table_and_follows_a_seek() {
        let Ok(path) = std::env::var("MEDIAS_DV_FILE") else { return };
        let feed = Feed::start(&path, false);
        let wait = |position: f64| {
            let start = std::time::Instant::now();
            loop {
                if let Some(table) = feed.table(position).unwrap() {
                    println!("table {} for {position} s after {:.0} ms", table.0, start.elapsed().as_secs_f64() * 1e3);
                    return table;
                }
                assert!(start.elapsed() < Duration::from_secs(10), "no table for {position} s");
                std::thread::sleep(Duration::from_millis(5));
            }
        };
        let (_, table) = wait(0.0);
        // Mid grey with neutral chroma stays neutral.
        let n = 65;
        let [r, g, b] = std::array::from_fn(|k| table.bytes[4 * (n / 2 + n * (n / 2 + n * (n / 2))) + k]);
        assert!(r.abs_diff(g) <= 2 && g.abs_diff(b) <= 2, "{r} {g} {b}");
        wait(1.5);
        wait(0.2);
    }

    /// Against libplacebo: MEDIAS_DV_CODES = one decoded frame, yuv444p10le; MEDIAS_DV_REF = the
    /// same frame through ffmpeg's libplacebo filter to PQ BT.2020 full range, gbrp16le
    /// (tonemapping=clip, gamut_mode=clip, peak_detect=false). RPU: `profile_5_rpu`.
    #[test]
    #[ignore]
    fn light_matches_libplacebo() {
        let (Ok(codes), Ok(reference)) = (std::env::var("MEDIAS_DV_CODES"), std::env::var("MEDIAS_DV_REF")) else { return };
        let words = |path: String| -> Vec<u16> { std::fs::read(path).unwrap().chunks_exact(2).map(|b| u16::from_le_bytes([b[0], b[1]])).collect() };
        let (codes, reference) = (words(codes), words(reference));
        let n = codes.len() / 3;
        assert_eq!(reference.len(), 3 * n);
        let mapping = Mapping::from_rpu(&profile_5_rpu()).unwrap();
        let to_rgb = lms_to_bt2020();
        let (mut worst, mut sum) = (0.0f32, 0.0f64);
        for i in 0..n {
            let ycc = [codes[i], codes[n + i], codes[2 * n + i]].map(|v| f32::from(v) / 1023.0);
            let got = mapping.light(ycc, &to_rgb).map(|v| video::pq_oetf(v.clamp(0.0, 10000.0)));
            let want = [reference[2 * n + i], reference[i], reference[n + i]].map(|v| f32::from(v) / 65535.0);
            for k in 0..3 {
                let error = (got[k] - want[k]).abs();
                if error > worst {
                    println!("pixel {i} ch {k}: codes {ycc:?} got {got:?} want {want:?}");
                }
                worst = worst.max(error);
                sum += f64::from(error);
            }
        }
        println!("largest PQ signal difference {worst:.4}, mean {:.5}", sum / (3 * n) as f64);
        assert!(worst < 0.02, "{worst}");
    }
}