use crate::diff::{diff, pixel_diff};
use crate::split::split;
use crate::{Block, Frame, Ref, View, P2};

pub struct EncodeConfig {
    pub threshold: u32,
    pub max_block_size: usize,
    pub iters: usize,
}

pub fn encode(last_frame: &Frame, frame: &Frame, view: View, config: &EncodeConfig) -> Block {
    let view_area = view.size().area();
    if view_area > config.max_block_size {
        let [av, bv] = split(view);
        let (ab, bb) = rayon::join(
            || Box::new(encode(last_frame, frame, av, config)),
            || Box::new(encode(last_frame, frame, bv, config)),
        );
        return Block::Split(ab, bb);
    }

    let mut r = Ref::default();
    let mut d = diff([last_frame, frame], view, r);

    let att = 1. - attention(frame, view) as f32 * 0.000001;
    let thres = (config.threshold as f32 * att.clamp(0.2, 1.0)) as u32;

    for granularity in [4, 2, 1] {
        for _ in 0..config.iters {
            let (nd, nrp) = optimize_ref(last_frame, frame, view, d, r, granularity);
            if nd < d {
                r = nrp;
                d = nd;
            } else {
                break;
            }
        }
    }

    if d < thres {
        return Block::Ref(r);
    } else {
        Block::Lit(frame.export(view))
    }
}

pub fn optimize_ref(
    last_frame: &Frame,
    frame: &Frame,
    view: View,
    mut d: u32,
    mut r: Ref,
    granularity: i32,
) -> (u32, Ref) {
    let mut n = |f: fn(&mut Ref, i32)| {
        let mut r2 = r;
        f(&mut r2, granularity);

        let d2 = diff([last_frame, frame], view, r2);
        if d2 < d {
            d = d2;
            r = r2;
        }
    };

    n(|r, g| r.pos_off.x += g);
    n(|r, g| r.pos_off.x -= g);
    n(|r, g| r.pos_off.y += g);
    n(|r, g| r.pos_off.y -= g);
    n(|r, g| r.color_off.r += (g as i16) << 2);
    n(|r, g| r.color_off.r -= (g as i16) << 2);
    n(|r, g| r.color_off.g += (g as i16) << 2);
    n(|r, g| r.color_off.g -= (g as i16) << 2);
    n(|r, g| r.color_off.b += (g as i16) << 2);
    n(|r, g| r.color_off.b -= (g as i16) << 2);

    (d, r)
}

pub fn attention(frame: &Frame, view: View) -> u32 {
    let mut k = 0;
    for y in view.a.y..view.b.y - 1 {
        for x in view.a.x..view.b.x - 1 {
            let p = P2 { x, y };
            k += pixel_diff(frame[p], frame[p + P2::X]).pow(2);
            k += pixel_diff(frame[p], frame[p + P2::Y]).pow(2);
        }
    }
    k
}