use framework::{BitstreamFilter, bitstream_filter_main};
use glam::{I16Vec2, i16vec2};
use mtree_test::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
use std::{collections::VecDeque, io::Result, time::Instant};

fn main() -> Result<()> {
    bitstream_filter_main::<Enc>()
}

struct Enc {
    res: I16Vec2,
    last: LastFrames,
    tree: MTree,
}
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,
            }),
        }
    }
    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 * BLOCK_SIZE * 20) as u32 {
                    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;

        let t = Instant::now();
        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());

        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 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
}