use jxl_grid::{AlignedGrid, MutableSubgrid};
use jxl_threadpool::JxlThreadPool;

use crate::{ImageWithRegion, Region};

use super::impls::generic::gabor::gabor_row_edge;

pub fn apply_gabor_like(
    fb: &mut ImageWithRegion,
    color_padded_region: Region,
    fb_scratch: &mut [AlignedGrid<f32>; 3],
    weights: [[f32; 2]; 3],
    pool: &jxl_threadpool::JxlThreadPool,
) {
    tracing::debug!("Running gaborish");
    let region = fb.regions_and_shifts()[0].0;
    assert!(region.contains(color_padded_region));
    let left = region.left.abs_diff(color_padded_region.left) as usize;
    let top = region.top.abs_diff(color_padded_region.top) as usize;
    let right = left + color_padded_region.width as usize;
    let bottom = top + color_padded_region.height as usize;

    let buffers = fb.as_color_floats_mut();
    let buffers = buffers.map(|g| g.as_subgrid_mut().subgrid(left..right, top..bottom));

    super::impls::apply_gabor_like(buffers, fb_scratch, weights, pool);

    let left = color_padded_region.left;
    let top = color_padded_region.top;
    for (idx, grid) in fb_scratch.iter_mut().enumerate() {
        let width = grid.width() as u32;
        let height = grid.height() as u32;
        let region = Region {
            width,
            height,
            left,
            top,
        };
        fb.swap_channel_f32(idx, grid, region);
    }
}

pub(super) struct GaborRow<'buf> {
    pub input_rows: [&'buf [f32]; 3],
    pub output_row: &'buf mut [f32],
    pub weights: [f32; 2],
}

pub(super) fn run_gabor_rows<'buf>(
    input: MutableSubgrid<'buf, f32>,
    output: &'buf mut AlignedGrid<f32>,
    weights: [f32; 2],
    pool: &JxlThreadPool,
    handle_row: for<'a> fn(GaborRow<'a>),
) {
    unsafe { run_gabor_rows_unsafe(input, output, weights, pool, handle_row) }
}

pub(super) unsafe fn run_gabor_rows_unsafe<'buf>(
    input: MutableSubgrid<'buf, f32>,
    output: &'buf mut AlignedGrid<f32>,
    weights: [f32; 2],
    pool: &JxlThreadPool,
    handle_row: for<'a> unsafe fn(GaborRow<'a>),
) {
    let width = input.width();
    let height = input.height();
    let output_buf = output.buf_mut();
    assert_eq!(output_buf.len(), width * height);

    if height == 1 {
        let input_buf = input.get_row(0);
        gabor_row_edge(input_buf, None, output_buf, weights);
        return;
    }

    {
        let input_buf_c = input.get_row(0);
        let input_buf_a = input.get_row(1);
        let output_buf = &mut output_buf[..width];
        gabor_row_edge(input_buf_c, Some(input_buf_a), output_buf, weights);
    }

    let (inner_rows, bottom_row) = output_buf[width..].split_at_mut((height - 2) * width);
    let output_rows = inner_rows
        .chunks_mut(width * 8)
        .enumerate()
        .collect::<Vec<_>>();

    pool.for_each_vec(output_rows, |(y8, output_rows)| {
        let it = output_rows.chunks_exact_mut(width);
        for (dy, output_row) in it.enumerate() {
            let y_up = y8 * 8 + dy;
            let input_rows = [
                input.get_row(y_up),
                input.get_row(y_up + 1),
                input.get_row(y_up + 2),
            ];
            let row = GaborRow {
                input_rows,
                output_row,
                weights,
            };
            unsafe {
                handle_row(row);
            }
        }
    });

    {
        let input_buf_c = input.get_row(height - 1);
        let input_buf_a = input.get_row(height - 2);
        let output_buf = bottom_row;
        gabor_row_edge(input_buf_c, Some(input_buf_a), output_buf, weights);
    }
}

#[allow(unused)]
pub(crate) fn run_gabor_row_generic(row: GaborRow) {
    super::impls::generic::gabor::run_gabor_row_generic(row)
}
