h264: cache data before use in fill_deblock()

[nihav.git] / nihav-duck / src / codecs / vp7enc / rdo.rs

use super::blocks::*;
use super::coder::*;

static mut COUNTER: usize = 0;

pub const SMALL_DIST: u32 = 256;
pub const MAX_DIST: u32 = std::u32::MAX;

const INTER_TO_INTRA_RATIO: f32 = 0.85;

pub struct RateDistMetric {
    pub lambda: f32,
}

impl RateDistMetric {
    pub fn new() -> Self {
        Self {
            lambda:     1.0,
        }
    }
    pub fn calc_metric(&self, dist: u32, nits: u32) -> u32 {
        ((dist as f32) + self.lambda * (nits as f32) + 0.5) as u32
    }
    pub fn adjust_br(&mut self, cur_size: usize, tgt_size: usize) {
        let low_limit = tgt_size - tgt_size / 8;
        let up_limit  = tgt_size + tgt_size / 8;
        if cur_size < low_limit {
            self.lambda = (self.lambda - 0.1).max(0.0);
        } else if cur_size > up_limit {
            self.lambda = (self.lambda + 0.1).min(16.0);
        }
    }

    pub fn block_dist(&self, src: &[u8; 16], new: &[u8; 16], q: usize, ctype: usize, pctx: u8, probs: &[[[u8; 11]; 3]; 8]) -> (u32, bool) {
        let mut diff = [0i16; 16];
        get_block_difference(&mut diff, src, new);
        diff.fdct();
        diff.quant(q, ctype);
        let has_nz = diff.has_nz();
        let nits = estimate_subblock_nits(&diff, ctype, pctx, probs);
        diff.dequant(q, ctype);
        diff.idct();
        let dist = get_difference_dist(src, new, &diff);
unsafe {COUNTER += 1;}
        (self.calc_metric(dist, nits), has_nz)
    }
}

#[derive(Default)]
pub struct BitRateControl {
    tb_num:     u32,
    tb_den:     u32,
    key_int:    u32,
    bitrate:    u32,
    force_q:    Option<usize>,
    bitpool:    u32,
    fpos:       u32,
    kpos:       u32,
    num_mb:     u32,
}

impl BitRateControl {
    pub fn new() -> Self { Self::default() }
    fn reset(&mut self) {
        self.fpos = 0;
        self.kpos = 0;
        self.bitpool = self.bitrate;
    }
    pub fn set_params(&mut self, tb_num: u32, tb_den: u32, bitrate: u32, key_int: u8, num_mb: usize) {
        self.tb_num  = tb_num;
        self.tb_den  = tb_den;
        self.bitrate = bitrate;
        self.key_int = u32::from(key_int);
        self.num_mb  = num_mb as u32;
        self.reset();
    }
    pub fn has_bitrate(&self) -> bool { self.bitrate != 0 }
    pub fn get_quant(&self) -> Option<usize> { self.force_q }
    pub fn set_quant(&mut self, q: Option<usize>) {
        if self.force_q != q {
            self.force_q = q;
            self.reset();
        }
    }
    pub fn set_key_interval(&mut self, key_int: u8) {
        let key_int = u32::from(key_int);
        if self.key_int != key_int {
            self.key_int = key_int;
            self.reset();
        }
    }
    pub fn get_target_size(&self, is_intra: bool) -> u32 {
        if self.bitrate != 0 && self.force_q.is_none() {
            let pool_frames = self.tb_den - self.fpos;
            if self.key_int <= 1 { // all intra
                if self.bitpool == 0 || pool_frames == 0 {
                    self.bitrate * self.tb_num / self.tb_den
                } else {
                    self.bitpool / pool_frames
                }
            } else {
                let full_gop_weight = 1.0 + ((self.key_int - 1) as f32) * INTER_TO_INTRA_RATIO;
                let i_bits = if self.bitpool == 0 || pool_frames == 0 {
                        let gop_size = self.bitrate * self.tb_num * self.key_int / self.tb_den;
                        (gop_size as f32) / full_gop_weight
                    } else {
                        let full_gops = pool_frames / self.key_int;
                        let weight = (full_gops as f32) * full_gop_weight + ((pool_frames % self.key_int) as f32) * INTER_TO_INTRA_RATIO;
                        (self.bitpool as f32) / weight
                    };
                if is_intra {
                    (i_bits + 0.5) as u32
                } else {
                    (i_bits * INTER_TO_INTRA_RATIO + 0.5) as u32
                }
            }
        } else {
            0
        }
    }
    fn pred_nits_per_mb(is_intra: bool, q: usize) -> f32 {
        let fq = q as f32;
        match (is_intra, q) {
            (true, 0..=6) => 3434.0 + fq * fq * 7.5 - fq * 195.0,
            (true, _) => 2500.0 - (fq - 6.0).ln() * 500.0,
            (false, 0..=10) => 1595.0 + fq * fq * 3.4 - fq * 125.0,
            (false, _) => 800.0 - (fq - 8.0).ln() * 155.0,
        }
    }
    #[allow(dead_code)]
    // todo use for refining maybe
    pub fn predict_size(&self, is_intra: bool, q: usize) -> u32 {
        let min_size = if is_intra { 200 * 8 } else { 50 * 8 };
        let nits_per_mb = Self::pred_nits_per_mb(is_intra, q);
        ((nits_per_mb * (self.num_mb as f32) / 8.0) as u32).max(min_size)
    }
    pub fn get_frame_quant(&self, is_intra: bool) -> usize {
        if let Some(q) = self.force_q {
            q
        } else {
            let expected_size = self.get_target_size(is_intra);
            let nits_per_mb = ((expected_size * 8) as f32) / (self.num_mb as f32);
            if is_intra {
                if nits_per_mb > 2500.0 { // simple search
                    if nits_per_mb > Self::pred_nits_per_mb(is_intra, 3) {
                        if nits_per_mb > Self::pred_nits_per_mb(is_intra, 1) {
                            0
                        } else if nits_per_mb > Self::pred_nits_per_mb(is_intra, 2) {
                            1
                        } else {
                            2
                        }
                    } else {
                        if nits_per_mb > Self::pred_nits_per_mb(is_intra, 4) {
                            3
                        } else if nits_per_mb > Self::pred_nits_per_mb(is_intra, 5) {
                            4
                        } else {
                            5
                        }
                    }
                } else {
                    ((((2500.0 - nits_per_mb) / 500.0).exp() + 6.0) as usize).min(127)
                }
            } else {
                if nits_per_mb > 680.0 { // simple search
                    let (start, end) = if nits_per_mb > Self::pred_nits_per_mb(is_intra, 5) {
                            if nits_per_mb > Self::pred_nits_per_mb(is_intra, 3) {
                                (0, 3)
                            } else {
                                (3, 5)
                            }
                        } else if nits_per_mb > Self::pred_nits_per_mb(is_intra, 7) {
                            (5, 7)
                        } else {
                            (7, 10)
                        };
                    let mut q = end;
                    for qq in start..end {
                        if nits_per_mb > Self::pred_nits_per_mb(is_intra, qq) {
                            q = qq;
                            break;
                        }
                    }
                    q
                } else {
                    ((((800.0 - nits_per_mb) / 155.0).exp() + 6.0) as usize).max(10).min(127)
                }
            }
        }
    }
    pub fn update(&mut self, size: usize) {
        self.kpos += 1;
        if self.kpos == self.key_int {
            self.kpos = 0;
        }
        if self.bitrate == 0 || self.force_q.is_some() {
            return;
        }
        self.fpos += self.tb_num;
        while self.fpos >= self.tb_den {
            self.fpos -= self.tb_den;
            self.bitpool += self.bitrate;
        }
        self.bitpool = self.bitpool.saturating_sub((size * 8) as u32);
    }
}