use crate::{
    block::{AdvancedReference, Block},
    frame::Frame,
    helpers::{pixel::Pixel, threading::both_par, vector::Vec2},
    refsampler::Sampler,
    view::View,
};
use clap::ValueEnum;

#[derive(Debug, Clone)]
pub struct EncodeConfig {
    pub mode: EncodeMode,
    pub ref_thres: f64,
    pub max_diff_area: isize,
    pub min_block_size: isize,
    pub max_threads: usize,
    pub weight_factor: f64,
    pub do_value_scale: bool,
    pub do_matrix_transform: bool,
}

#[derive(Debug, Clone, ValueEnum)]
pub enum EncodeMode {
    Trivial,
    Fast,
    Default,
    Advanced,
}

#[inline]
pub fn pk<F: FnMut(&mut AdvancedReference) -> ()>(
    view: &View<&Frame>,
    prev: &View<&Frame>,
    diff: &mut f64,
    params: &mut AdvancedReference,
    initial_params: &AdvancedReference,
    mut f: F,
) {
    let mut p = initial_params.clone();
    f(&mut p);
    let sampler = Sampler::from_refblock(prev.clone(), &p);
    let d = View::diff_sampler(view, &sampler);
    if d < *diff {
        *diff = d;
        *params = p;
    }
}

pub fn encode_block(view: View<&Frame>, prev: View<&Frame>, config: &EncodeConfig) -> Block {
    let (diff, refblock) = if view.area() > config.max_diff_area {
        (
            f64::INFINITY,
            Block::Reference {
                translation: Vec2::<isize>::ZERO,
            },
        )
    } else {
        let weight = importance(&view).max(0.5);
        let irrelevance = config.weight_factor / weight;
        let (diff, refblock) = match config.mode {
            EncodeMode::Trivial => (
                View::diff(&view, &prev) / view.area() as f64,
                Block::Reference {
                    translation: Vec2::<isize>::ZERO,
                },
            ),
            EncodeMode::Default => {
                let mut diff = f64::INFINITY;
                let mut translation = Vec2::<isize>::ZERO;
                const OFFSETS: &[isize] =
                    &[-64, -32, -16, -8, -4, -2, -1, 0, 1, 2, 4, 8, 16, 32, 64];
                for x in OFFSETS {
                    for y in OFFSETS {
                        let t = Vec2 { x: *x, y: *y };
                        let d = View::diff(&view, &prev.offset(t));
                        if d < diff {
                            translation = t;
                            diff = d;
                        }
                    }
                }
                (diff, Block::Reference { translation })
            }
            EncodeMode::Fast => {
                let mut pm = AdvancedReference::default();
                let sampler = Sampler::from_refblock(prev.clone(), &pm);
                let mut diff = View::diff_sampler(&view, &sampler);
                if diff - irrelevance < config.ref_thres {
                    (diff, Block::REFZERO)
                } else {
                    loop {
                        let (mut d, mut p) = (diff, pm.clone());
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x += 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x -= 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y += 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y -= 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x += 1);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x -= 1);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y += 1);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y -= 1);
                        if d >= diff {
                            break (diff, Block::AdvancedReference(pm));
                        }
                        diff = d;
                        pm = p;
                    }
                }
            }
            EncodeMode::Advanced => {
                let mut pm = AdvancedReference::default();
                let sampler = Sampler::from_refblock(prev.clone(), &pm);
                let mut diff = View::diff_sampler(&view, &sampler);
                if diff - irrelevance < config.ref_thres {
                    (diff, Block::REFZERO)
                } else {
                    loop {
                        let (mut d, mut p) = (diff, pm.clone());

                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x += 4);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x -= 4);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y += 4);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y -= 4);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x += 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x -= 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y += 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y -= 2);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x += 1);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.x -= 1);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y += 1);
                        pk(&view, &prev, &mut d, &mut p, &pm, |p| p.translation.y -= 1);
                        if config.do_value_scale {
                            // pk(&view, &prev, &mut d, &mut p, &pm, |p| p.value_scale += 1);
                            // pk(&view, &prev, &mut d, &mut p, &pm, |p| p.value_scale -= 1);
                        }
                        if config.do_matrix_transform {
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.a -= 1);
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.a += 1);
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.b -= 1);
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.b += 1);
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.c -= 1);
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.c += 1);
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.d -= 1);
                            pk(&view, &prev, &mut d, &mut p, &pm, |p| p.transform.d += 1);
                        }

                        if d >= diff {
                            break (diff, Block::AdvancedReference(pm));
                        }

                        diff = d;
                        pm = p;
                    }
                }
            }
        };
        (diff - irrelevance, refblock)
    };
    if diff < config.ref_thres {
        refblock
    } else {
        if view.size.x < config.min_block_size || view.size.y < config.min_block_size {
            Block::Literal(view.pixels())
        } else {
            let [av, bv] =
                unsafe { std::mem::transmute::<_, [View<&'static Frame>; 2]>(view.split()) };
            let [ap, bp] =
                unsafe { std::mem::transmute::<_, [View<&'static Frame>; 2]>(prev.split()) };
            let config = unsafe { std::mem::transmute::<_, &'static EncodeConfig>(config) };

            // only bother to do multithreading, when the block is big.
            let (a, b) =  // if view.size.x > 64 {
                both_par(
                    || encode_block(av, ap, config),
                    || encode_block(bv, bp, config),
                    config.max_threads,
                )
            // } else {
            //     (encode_block(av, ap, config), encode_block(bv, bp, config))
            // }
            ;

            if a.is_literal() && b.is_literal() {
                Block::Literal(view.pixels())
            } else if a.is_ref_without_translation() && b.is_ref_without_translation() {
                Block::Reference {
                    translation: Vec2::<isize>::ZERO,
                }
            } else {
                Block::Split(Box::new([a, b]))
            }
        }
    }
}

pub fn importance(view: &View<&Frame>) -> f64 {
    let mut acc = 0;
    for x in 0..view.size.x {
        for y in 0..view.size.y {
            let p = Vec2 { x, y };
            if x > 0 {
                acc += Pixel::distance(view[p], view[p + Vec2::<isize>::LEFT]);
            }
            if y > 0 {
                acc += Pixel::distance(view[p], view[p + Vec2::<isize>::UP]);
            }
        }
    }
    (acc / view.area() as usize) as f64
}