use multiversion::multiversion;
use rayon::prelude::*;
use std::{ops::Add, time::Instant};

use crate::{
  cfa::{CFA, CFA_COLOR_B, CFA_COLOR_G, CFA_COLOR_R, PlaneColor},
  imgop::{Rect, sensor::Demosaic},
  pixarray::{Color2D, PixF32, RgbF32},
};

#[derive(Default)]
pub struct XtransFastDemosaic {}

impl XtransFastDemosaic {
  pub fn new() -> Self {
    Self {}
  }
}

/// PPG demosaic a raw image (f32 values)
///
/// The PPG - Pattern Pixel Grouping - alogrithm was developed by Chuan-kai Lin
/// Source on Internet Archive:
/// https://web.archive.org/web/20160923211135/https://sites.google.com/site/chklin/demosaic/
///
/// This is a simple algorithm but provides acceptable results.
///
/// # Panics
///
/// This function panics for CFA pattern that are not RGGB or variants. You need
/// to check the pattern before calling.
impl Demosaic<f32, 3> for XtransFastDemosaic {
  /// Debayer image by using superpixel method.
  /// Each output pixel RGB tuple is generated by 4 pixels from input.
  /// The result image is 1/4 of size.
  fn demosaic(&self, pixels: &PixF32, cfa: &CFA, _colors: &PlaneColor, roi: Rect) -> Color2D<f32, 3> {
    // PPG can only applied to pure RGGB or variants.
    if !cfa.is_rgb() {
      panic!("CFA pattern '{}' is not a RGB pattern, can not demosaic with xtrans", cfa);
    }
    // Measure time
    let now = Instant::now();

    // The ROI changes the pattern if not perfectly aligned on the origin pattern
    let cfa_roi = cfa.shift(roi.p.x, roi.p.y);

    // Expand the bayer data to full RGB channel image, but only for ROI
    let mut rgb = super::expand_xtrans_rgb(pixels.pixels(), pixels.dim(), cfa, roi);

    // Now interpolate the missing channels
    interpolate_missing(&mut rgb, &cfa_roi);

    log::debug!("xtrans total debayer time: {:.5}s", now.elapsed().as_secs_f32());
    rgb
  }
}

#[multiversion(targets("x86_64+avx+avx2", "x86+sse", "aarch64+neon"))]
fn interpolate_missing(input: &mut RgbF32, shifted: &CFA) {
  let w = input.width;
  let h = input.height;
  // Iterate over all rows
  for row in 0..h {
    for col in 0..w {
      let mut tmp = *input.at(row, col);
      for (i, p) in tmp.iter_mut().enumerate() {
        if p.is_nan() {
          if shifted.color_at(row.saturating_sub(1), col) == i {
            *p = input.at(row.saturating_sub(1), col)[i];
          }
          if shifted.color_at(row, col.saturating_sub(1)) == i {
            *p = input.at(row, col.saturating_sub(1))[i];
          }
          if shifted.color_at((h - 1).min(row + 1), col) == i {
            *p = input.at((h - 1).min(row + 1), col)[i];
          }
          if shifted.color_at(row, (w - 1).min(col + 1)) == i {
            *p = input.at(row, (w - 1).min(col + 1))[i];
          }
        }
      }
      *input.at_mut(row, col) = tmp;
    }
  }
}
