path: root/mtree-test/src/encode.rs

use crate::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
use framework::BitstreamFilter;
use glam::{I16Vec2, i16vec2};
use rayon::{
    iter::{IntoParallelRefIterator, ParallelIterator},
    join,
};
use std::{collections::VecDeque, time::Instant};

pub struct Enc {
    res: I16Vec2,
    last: LastFrames,
    tree: MTree,
    frame_num: u64,

    use_bulk_insert: bool,
}
impl BitstreamFilter for Enc {
    const INPUT_CODEC_ID: &str = "V_UNCOMPRESSED";
    const OUTPUT_CODEC_ID: &str = "V_VCEMTREE";

    fn new(width: u32, height: u32) -> Self {
        Self {
            res: i16vec2(width as i16, height as i16),
            last: LastFrames {
                frame_offset: 0,
                frames: VecDeque::new(),
            },
            tree: MTree::Leaf(AbsRef {
                off: i16vec2(0, 0),
                frame: u64::MAX,
            }),
            frame_num: 0,
            use_bulk_insert: true,
        }
    }
    fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
        let mut frame = Frame(frame);
        let mut out = Vec::new();
        let mut num_refs = 0;

        let t = Instant::now();
        for bx in 0..self.res.x / BLOCK_SIZE {
            for by in 0..self.res.y / BLOCK_SIZE {
                let boff = i16vec2(bx * BLOCK_SIZE, by * BLOCK_SIZE);
                let r = self.tree.search(self.res, &self.last, &frame, boff);

                let rdist = if r.frame == u64::MAX {
                    u32::MAX
                } else {
                    distance(
                        self.res,
                        &self.last.frames[(r.frame - self.last.frame_offset) as usize],
                        &frame,
                        r.off,
                        boff,
                    )
                };
                if rdist < BLOCK_SIZE as u32 * BLOCK_SIZE as u32 * 50 {
                    out.push(1);
                    out.extend(r.frame.to_le_bytes());
                    out.extend(r.off.x.to_le_bytes());
                    out.extend(r.off.y.to_le_bytes());
                    frame_to_frame_rect_copy(
                        self.res,
                        &mut frame,
                        &self.last.frames[(r.frame - self.last.frame_offset) as usize],
                        I16Vec2::splat(BLOCK_SIZE),
                        boff,
                        r.off,
                    );
                    num_refs += 1;
                } else {
                    out.push(0);
                    frame.export_rect(self.res, boff..boff + I16Vec2::splat(BLOCK_SIZE), &mut out);
                }
            }
        }
        eprintln!("search\t{:?}", t.elapsed());

        self.last.frames.push_back(frame);
        let frame_index = self.last.frame_offset + self.last.frames.len() as u64 - 1;

        if self.frame_num % 30 == 0 {
            let t = Instant::now();
            if self.use_bulk_insert {
                let mut rs = Vec::new();
                for xo in 0..self.res.x - BLOCK_SIZE {
                    for yo in 0..self.res.y - BLOCK_SIZE {
                        rs.push(AbsRef {
                            frame: frame_index,
                            off: i16vec2(xo, yo),
                        });
                    }
                }
                self.tree.bulk_insert(self.res, &self.last, rs);
            } else {
                for xo in 0..self.res.x - BLOCK_SIZE {
                    for yo in 0..self.res.y - BLOCK_SIZE {
                        self.tree.insert(
                            self.res,
                            &self.last,
                            AbsRef {
                                frame: frame_index,
                                off: i16vec2(xo, yo),
                            },
                        )
                    }
                }
            }
            eprintln!("insert\t{:?}", t.elapsed());
        }

        eprintln!("depth={} refs={num_refs}", self.tree.depth());

        self.frame_num += 1;
        out
    }
}

enum MTree {
    Branch(Box<[MTree; 2]>),
    Leaf(AbsRef),
}
impl MTree {
    pub fn depth(&self) -> usize {
        (match self {
            MTree::Branch(b) => b[0].depth().max(b[1].depth()),
            MTree::Leaf(_) => 0,
        }) + 1
    }
    pub fn center(&self) -> AbsRef {
        match self {
            MTree::Branch(b) => b[0].center(),
            MTree::Leaf(r) => *r,
        }
    }
    pub fn bulk_insert(&mut self, res: I16Vec2, last: &LastFrames, mut rs: Vec<AbsRef>) {
        match self {
            MTree::Branch(b) => {
                let dirs = rs
                    .par_iter()
                    .map(|r| {
                        distance_absref(res, last, b[0].center(), *r)
                            < distance_absref(res, last, b[1].center(), *r)
                    })
                    .collect::<Vec<bool>>();

                let mut rs0 = Vec::new();
                let mut rs1 = Vec::new();
                for (r, d) in rs.into_iter().zip(dirs) {
                    if d {
                        rs0.push(r);
                    } else {
                        rs1.push(r);
                    }
                }
                let [b0, b1] = b.as_mut();
                join(
                    || b0.bulk_insert(res, last, rs0),
                    || b1.bulk_insert(res, last, rs1),
                );
            }
            MTree::Leaf(l) => {
                if rs.is_empty() {
                    return;
                }
                if l.frame == u64::MAX {
                    *l = rs.pop().unwrap();
                    if rs.is_empty() {
                        return;
                    }
                }
                let mut b =
                    MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(rs.pop().unwrap())]));
                b.bulk_insert(res, last, rs);
                *self = b;
            }
        }
    }
    pub fn insert(&mut self, res: I16Vec2, last: &LastFrames, r: AbsRef) {
        match self {
            MTree::Branch(b) => {
                if distance_absref(res, last, b[0].center(), r)
                    < distance_absref(res, last, b[1].center(), r)
                {
                    b[0].insert(res, last, r)
                } else {
                    b[1].insert(res, last, r)
                }
            }
            MTree::Leaf(l) => {
                if l.frame == u64::MAX {
                    *l = r
                } else {
                    *self = MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(r)]))
                }
            }
        }
    }
    pub fn search(
        &mut self,
        res: I16Vec2,
        last: &LastFrames,
        frame: &Frame,
        off: I16Vec2,
    ) -> AbsRef {
        match self {
            MTree::Branch(b) => {
                let acenter = b[0].center();
                let bcenter = b[1].center();
                let adist = distance(
                    res,
                    &last.frames[(acenter.frame - last.frame_offset) as usize],
                    frame,
                    acenter.off,
                    off,
                );
                let bdist = distance(
                    res,
                    &last.frames[(bcenter.frame - last.frame_offset) as usize],
                    frame,
                    bcenter.off,
                    off,
                );
                if adist < bdist {
                    b[0].search(res, last, frame, off)
                } else {
                    b[1].search(res, last, frame, off)
                }
            }
            MTree::Leaf(l) => *l,
        }
    }
}

fn distance_absref(res: I16Vec2, last: &LastFrames, a: AbsRef, b: AbsRef) -> u32 {
    let aframe = &last.frames[(a.frame - last.frame_offset) as usize];
    let bframe = &last.frames[(b.frame - last.frame_offset) as usize];
    distance(res, aframe, bframe, a.off, b.off)
}
fn distance(res: I16Vec2, aframe: &Frame, bframe: &Frame, aoff: I16Vec2, boff: I16Vec2) -> u32 {
    let mut diff = 0;
    for yo in 0..BLOCK_SIZE {
        // for xo in 0..BLOCK_SIZE {
        //     let off = i16vec2(xo, yo);
        //     let asample = aframe.get(res, aoff + off);
        //     let bsample = bframe.get(res, boff + off);

        //     diff += asample[0].abs_diff(bsample[0]) as u32;
        //     diff += asample[1].abs_diff(bsample[1]) as u32;
        //     diff += asample[2].abs_diff(bsample[2]) as u32;
        // }

        let aslice = aframe.slice(res, yo + aoff.y, aoff.x..aoff.x + BLOCK_SIZE);
        let bslice = bframe.slice(res, yo + boff.y, boff.x..boff.x + BLOCK_SIZE);

        for i in 0..BLOCK_SIZE as usize {
            diff += aslice[0][i].abs_diff(bslice[0][i]) as u32;
            diff += aslice[0][i].abs_diff(bslice[0][i]) as u32;
        }
        for i in 0..BLOCK_SIZE as usize / 2 {
            diff += aslice[1][i].abs_diff(bslice[1][i]) as u32;
            diff += aslice[1][i].abs_diff(bslice[1][i]) as u32;
        }
        for i in 0..BLOCK_SIZE as usize / 2 {
            diff += aslice[2][i].abs_diff(bslice[2][i]) as u32;
            diff += aslice[2][i].abs_diff(bslice[2][i]) as u32;
        }
    }
    diff
}