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

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

const BLOCK_SIZE: i16 = 16;

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: 0,
            }),
        }
    }
    fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
        let frame = Frame(frame);

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

        for xo in 0..self.res.x - BLOCK_SIZE {
            eprintln!("aaaaaaa {}  aaaa", self.tree.depth());
            for yo in 0..self.res.y - BLOCK_SIZE {
                self.tree.insert(
                    self.res,
                    &self.last,
                    AbsRef {
                        frame: frame_index,
                        off: i16vec2(xo, yo),
                    },
                )
            }
        }

        Vec::new()
    }
}

struct LastFrames {
    frame_offset: u64,
    frames: VecDeque<Frame>,
}

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) => *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;
        }
    }
    diff
}