use crate::{array2d::Array2d, Crop, Image, ResizableImage};

const RESZ: u32 = 128;
const RESZU: usize = RESZ as usize;
const RESZF64: f64 = RESZ as f64;

const WHITE_THRESHOLD: u16 = 16;
const ROW_THRESHOLD: f64 = 50.0;
const MIN_CROPPED_SIZE: f64 = 60.0;
const MAX_ASYMMETRY: isize = 10;

/// Check if the image has black borders and return a crop to remove them
pub fn find_border_crop<I: Image + ResizableImage<RI>, RI: Image>(img: &I) -> Option<Crop> {
    let bw = image_to_bw(img);
    let eroded = erode(&bw, 2);
    let contours = sobel(&eroded);
    get_crop(&contours, img.width(), img.height())
}

/// Convert an image to a scaled-down 128x128 black and white representation
fn image_to_bw<I: Image + ResizableImage<RI>, RI: Image>(img: &I) -> Array2d {
    let img_small = img.resize(RESZ, RESZ);
    let mut bw_pixels = Vec::with_capacity(RESZU * RESZU);
    for y in 0..RESZ {
        for x in 0..RESZ {
            let p = img_small.get(x, y);
            if u16::from(p.r) + u16::from(p.g) + u16::from(p.b) > WHITE_THRESHOLD {
                bw_pixels.push(1.0);
            } else {
                bw_pixels.push(0.0);
            }
        }
    }
    Array2d::from_vec(bw_pixels, (RESZ, RESZ))
}

fn get_crop(contours: &Array2d, width: u32, height: u32) -> Option<Crop> {
    let filter_row = |row: &[f64]| -> Vec<bool> {
        row.iter()
            .enumerate()
            .map(|(i, &val)| val > ROW_THRESHOLD && maximum_filter1d_point(row, 2, i) == val)
            .collect()
    };

    let horiz = contours.rows().map(|r| r.sum()).collect::<Vec<f64>>();
    let horiz_filtered = filter_row(&horiz);

    let verti = contours.cols().map(|r| r.sum()).collect::<Vec<f64>>();
    let verti_filtered = filter_row(&verti);

    let mut crop = Crop {
        x: 0,
        y: 0,
        width,
        height,
    };

    if let Some((i, j)) = find_cut_points(&horiz_filtered) {
        let im = (i * height as f64 / RESZF64) as u32;
        let jm = (j * height as f64 / RESZF64) as u32;
        crop.y = im;
        crop.height = jm - im;
    }
    if let Some((i, j)) = find_cut_points(&verti_filtered) {
        let im = (i * width as f64 / RESZF64) as u32;
        let jm = (j * width as f64 / RESZF64) as u32;
        crop.x = im;
        crop.width = jm - im;
    }

    Some(crop).filter(|c| c.height != height || c.width != width)
}

fn maximum_filter1d_point(input: &[f64], ext: usize, i: usize) -> f64 {
    let window = &input[i.saturating_sub(ext)..(i + ext + 1).min(input.len())];
    *window.iter().max_by(|x, y| x.total_cmp(y)).unwrap()
}

fn find_cut_points(filtered_rows: &[bool]) -> Option<(f64, f64)> {
    // List of all row indices which are above the threshold
    let indices: Vec<usize> = filtered_rows
        .iter()
        .enumerate()
        .filter(|&(_, &val)| val)
        .map(|(idx, _)| idx)
        .collect();

    // Get a list of continous ranges of indices
    // example: [16, 17, 112, 114] => [16..18, 112..114]
    let mut ranges = Vec::new();
    let mut r = 0..0;

    for i in indices {
        if r.start == 0 {
            r.start = i;
            r.end = i + 1;
        } else if i == r.end {
            r.end = i + 1;
        } else {
            ranges.push(r);
            r = i..(i + 1);
        }
    }
    if !r.is_empty() {
        ranges.push(r);
    }

    if ranges.len() != 2 {
        return None;
    }

    let points = ((ranges[0].end - 1) as f64, ranges[1].start as f64);

    // Check if there is enough space between the borders
    if points.1 - points.0 < MIN_CROPPED_SIZE {
        return None;
    }

    // Check if the borders are nearly symmetrical
    if ((points.0 as isize) - (RESZ as isize - points.1 as isize)).abs() > MAX_ASYMMETRY {
        return None;
    }

    Some(points)
}

fn erode(image: &Array2d, size: usize) -> Array2d {
    let mask_s = size as isize;
    let mask_d = mask_s / 2;

    fn check_pixel(image: &Array2d, mask_s: isize, mask_d: isize, y: isize, x: isize) -> f64 {
        let (my_start, mx_start) = (y - mask_d, x - mask_d);
        for my in 0..mask_s {
            for mx in 0..mask_s {
                if image.get_reflect(my_start + my, mx_start + mx) == 0.0 {
                    return 0.0;
                }
            }
        }
        1.0
    }

    Array2d::from_cb(image.dimensions(), |(x, y)| {
        check_pixel(image, mask_s, mask_d, y as isize, x as isize)
    })
}

fn sobel(image: &Array2d) -> Array2d {
    let edge_weights = [1.0, 0.0, -1.0];
    let smooth_weights = [0.25, 0.5, 0.25];

    let mut output = Array2d::new(image.dimensions());

    for edge_dim in 0..2 {
        let mut kernel = reshape_nd(&edge_weights, edge_dim);
        let smooth_dim: usize = if edge_dim == 0 { 1 } else { 0 };
        let ks = reshape_nd(&smooth_weights, smooth_dim);
        kernel = kernel * ks;

        convolve(image, &kernel, &mut output);
        output.apply_op(|x| x * x);
    }

    output.apply_op(|x| x.sqrt() / std::f64::consts::SQRT_2);
    output
}

fn reshape_nd(arr: &[f64], dim: usize) -> Array2d {
    let mut shape = [1; 2];
    shape[dim] = arr.len() as u32;
    Array2d::from_vec(arr.to_vec(), (shape[1], shape[0]))
}

fn convolve(image: &Array2d, kernel: &Array2d, output: &mut Array2d) {
    // Reverse order of columns and rows
    let (kw, kh) = kernel.dimensions();
    let weights = Array2d::from_cb(kernel.dimensions(), |(x, y)| {
        kernel.get(kw - x - 1, kh - y - 1)
    });

    correlate(image, &weights, output)
}

fn correlate(image: &Array2d, weights: &Array2d, output: &mut Array2d) {
    let (w, h) = image.dimensions();
    let (mask_w, mask_h) = weights.dimensions();
    let (mask_dx, mask_dy) = (mask_w as isize / 2, mask_h as isize / 2);

    for y in 0..(h as isize) {
        for x in 0..(w as isize) {
            let (my_start, mx_start) = (y - mask_dy, x - mask_dx);
            let mut p = 0.0;
            for my in 0..mask_h {
                for mx in 0..mask_w {
                    let x = weights.get(mx, my)
                        * image.get_reflect(my_start + (my as isize), mx_start + (mx as isize));
                    p += x;
                }
            }

            let pixel = output.get_mut(x as u32, y as u32);
            *pixel += p;
        }
    }
}
