use crate::diff::{diff, pixel_diff};
use crate::huff::write_huff;
use crate::split::split;
use crate::{decode::decode_block, Block, Frame, Pixel, Ref, View, P2};
use std::io::{BufReader, BufWriter, Read, Write};
use std::time::Instant;

#[derive(Debug, Clone)]
pub struct EncodeConfig {
    pub threshold: f32,
    pub max_block_size: usize,
    pub attention_split: u32,
    pub keyframe_interval: usize,
}

pub fn encode(
    config: EncodeConfig,
    size: P2,
    input: impl Read,
    output: impl Write,
) -> std::io::Result<()> {
    let mut input = BufReader::new(input);
    let mut output = BufWriter::new(output);

    let mut last_frame = Frame::new(size);
    for frame_number in 0.. {
        let mut frame = Frame::read(&mut input, size)?;

        let mut config = config.clone();
        if frame_number % config.keyframe_interval != 0 {
            config.threshold = std::f32::INFINITY;
        }

        let t = Instant::now();
        let b: Block = encode_block(&last_frame, &frame, View::all(size), &config);
        let time_encode = t.elapsed();

        let t = Instant::now();
        decode_block(&last_frame, &mut frame, View::all(size), &b);
        last_frame = frame;
        let time_decode = t.elapsed();

        if true {
            let mut buf = vec![];
            let mut bufw = std::io::Cursor::new(&mut buf);
            b.write(&mut bufw)?;
            drop(bufw);
            let t = Instant::now();
            let bits_raw = buf.len() * 8;
            let bits_huff = write_huff(&buf, &mut output)?;
            let time_huff = t.elapsed();
            drop(buf);

            eprintln!(
                "frame {frame_number}: {:?}",
                time_decode + time_huff + time_encode
            );
            eprintln!(
                "\tencode {time_encode:?} ({:.2}%)",
                (bits_raw as f32 / (size.area() * 24) as f32) * 100.0
            );
            eprintln!(
                "\thuff {time_huff:?} ({:.2}%)",
                (bits_huff as f32 / bits_raw as f32) * 100.0
            );
            eprintln!("\tdecode {time_decode:?}");
        } else {
            b.write(&mut output)?;
        }
    }
    Ok(())
}

pub fn encode_block(last_frame: &Frame, frame: &Frame, view: View, config: &EncodeConfig) -> Block {
    let view_area = view.size().area();
    if view_area > config.max_block_size
        || (view_area > 64 && attention(frame, view) > config.attention_split)
    {
        let [av, bv] = split(view);
        let (ab, bb) = rayon::join(
            || Box::new(encode_block(last_frame, frame, av, config)),
            || Box::new(encode_block(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;
    let thres = (config.threshold * view_area as f32) as u32;

    let target_average = average_color(frame, view);

    for granularity in [2, 1, 2, 1, 2, 1, 2, 1] {
        let (nd, nrp) = optimize_ref(last_frame, frame, view, r, granularity, target_average);
        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,
    r: Ref,
    g: i32,
    target_average: Pixel,
) -> (u32, Ref) {
    let g2 = g * 2;
    [
        Some(r.apply(|r| r.pos_off += P2 { x: g, y: 0 })),
        Some(r.apply(|r| r.pos_off += P2 { x: g, y: g })),
        Some(r.apply(|r| r.pos_off += P2 { x: 0, y: g })),
        Some(r.apply(|r| r.pos_off += P2 { x: -g, y: g })),
        Some(r.apply(|r| r.pos_off += P2 { x: -g, y: 0 })),
        Some(r.apply(|r| r.pos_off += P2 { x: -g, y: -g })),
        Some(r.apply(|r| r.pos_off += P2 { x: 0, y: -g })),
        Some(r.apply(|r| r.pos_off += P2 { x: g, y: -g })),
        Some(r.apply(|r| r.pos_off += P2 { x: g2, y: 0 })),
        Some(r.apply(|r| r.pos_off += P2 { x: g2, y: g2 })),
        Some(r.apply(|r| r.pos_off += P2 { x: 0, y: g2 })),
        Some(r.apply(|r| r.pos_off += P2 { x: -g2, y: g2 })),
        Some(r.apply(|r| r.pos_off += P2 { x: -g2, y: 0 })),
        Some(r.apply(|r| r.pos_off += P2 { x: -g2, y: -g2 })),
        Some(r.apply(|r| r.pos_off += P2 { x: 0, y: -g2 })),
        Some(r.apply(|r| r.pos_off += P2 { x: g2, y: -g2 })),
        {
            let mut r = r;
            let last_avr = average_color(last_frame, view);
            let diff = target_average - last_avr;
            r.color_off = diff;
            if diff != Pixel::BLACK {
                Some(r)
            } else {
                None
            }
        },
    ]
    .into_iter()
    .flatten()
    .map(|r| (diff([last_frame, frame], view, r), r))
    .min_by_key(|e| e.0)
    .unwrap()
}

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
}

pub fn average_color(frame: &Frame, view: View) -> Pixel {
    let mut r = 0u32;
    let mut g = 0u32;
    let mut b = 0u32;

    for y in view.a.y..view.b.y {
        for x in view.a.x..view.b.x {
            let p = frame[P2 { x, y }];
            r += p.r as u32;
            g += p.g as u32;
            b += p.b as u32;
        }
    }
    let area = view.size().area() as u32;
    Pixel {
        r: (r / area) as i16,
        g: (g / area) as i16,
        b: (b / area) as i16,
    }
}