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

pub struct Enc {
    res: IVec2,
    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: ivec2(width as i32, height as i32),
            last: LastFrames {
                frame_offset: 0,
                frames: VecDeque::new(),
            },
            tree: MTree::Leaf(AbsRef {
                off: ivec2(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 {
            data: frame,
            res: self.res,
        };
        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 = ivec2(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.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 as i16).to_le_bytes());
                    out.extend((r.off.y as i16).to_le_bytes());
                    frame_to_frame_rect_copy(
                        &self.last.frames[(r.frame - self.last.frame_offset) as usize],
                        &mut frame,
                        IVec2::splat(BLOCK_SIZE),
                        r.off,
                        boff,
                    );
                    num_refs += 1;
                } else {
                    out.push(0);
                    frame.export_rect(boff, IVec2::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: ivec2(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: ivec2(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: IVec2, last: &LastFrames, mut rs: Vec<AbsRef>) {
        match self {
            MTree::Branch(b) => {
                let dirs = rs
                    .par_iter()
                    .map(|r| {
                        distance_absref(last, b[0].center(), *r)
                            < distance_absref(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: IVec2, last: &LastFrames, r: AbsRef) {
        match self {
            MTree::Branch(b) => {
                if distance_absref(last, b[0].center(), r) < distance_absref(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: IVec2, last: &LastFrames, frame: &Frame, off: IVec2) -> AbsRef {
        match self {
            MTree::Branch(b) => {
                let acenter = b[0].center();
                let bcenter = b[1].center();
                let adist = distance(
                    &last.frames[(acenter.frame - last.frame_offset) as usize],
                    frame,
                    acenter.off,
                    off,
                );
                let bdist = distance(
                    &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(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(aframe, bframe, a.off, b.off)
}
fn distance(aframe: &Frame, bframe: &Frame, aoff: IVec2, boff: IVec2) -> u32 {
    assert_eq!(aframe.res, bframe.res);
    let res = aframe.res;
    let mut diff = 0;

    // Luma
    for y in 0..BLOCK_SIZE {
        let ay_off = aoff.x + (y + aoff.y) * res.x;
        let by_off = boff.x + (y + boff.y) * res.x;

        for x in 0..BLOCK_SIZE {
            diff += aframe.data[(ay_off + x) as usize].abs_diff(bframe.data[(by_off + x) as usize])
                as u32
        }
    }

    // Chroma
    let uvplane_off = res.x * res.y;
    for y in 0..BLOCK_SIZE / 2 {
        let ay_off = uvplane_off + (aoff.x & !1) + (y + aoff.y / 2) * res.x;
        let by_off = uvplane_off + (boff.x & !1) + (y + boff.y / 2) * res.x;

        for x in 0..BLOCK_SIZE {
            diff += aframe.data[(ay_off + x) as usize].abs_diff(bframe.data[(by_off + x) as usize])
                as u32
                * 2;
        }
    }

    diff
}