const gs : usize = 7;
const sz : usize = 3;
const sb : usize = 3;

fn gaussian_smoothing<n, m : usize>(
  input: f32[n, m],
  filter: f32[7, 7],
) -> f32[n, m] {
  @res let result : f32[n, m];

  // Define the gaussian radius as half the gaussian size
  const gr = gs / 2;

  @image_loop for row = 0 to n {
    for col = 0 to m {
      let smoothed = 0.0;

      @filter_loop for i = 0 to gs {
        for j = 0 to gs {
          let br = min!(max!(row + i, gr) - gr, n - 1);
          let bc = min!(max!(col + j, gr) - gr, m - 1);
          let val = input[br, bc];
          smoothed += val * filter[i, j];
        }
      }

      result[row, col] = smoothed;
    }
  }

  return result;
}

const MIN_BR : f32 = 0;
const MAX_BR : f32 = 1;

fn laplacian_estimate<n, m : usize>(
  input: f32[n, m],
  structure: f32[3, 3],
) -> f32[n, m] {
  const r = sz / 2;

  @res let result : f32[n, m];

  @image_loop for row = 0 to n {
    for col = 0 to m {
      // Compute pixel of dilated image
      let dilated = MIN_BR;
      @filter_loop for i = 0 to sz {
        for j = 0 to sz {
	  let filter = if row + i < r
                         || row + i - r > n - 1
                         || col + j < r
                         || col + j - r > m - 1 then MIN_BR
                      else input[row + i - r, col + j - r];
          dilated = max!(dilated, filter * structure[i, j]);
        }
      }

      // Compute pixel of eroded image
      let eroded = MAX_BR;
      @filter_loop for i = 0 to sz {
        for j = 0 to sz {
	  let filter = if row + i < r
                         || row + i - r > n - 1
                         || col + j < r
                         || col + j - r > m - 1 then MAX_BR
                      else input[row + i - r, col + j - r];
          eroded = min!(eroded, filter * structure[i, j]);
        }
      }

      let laplacian = dilated + eroded - 2 * input[row, col];
      result[row, col] = laplacian;
    }
  }

  return result;
}

fn zero_crossings<n, m : usize>(
  input: f32[n, m],
  structure: f32[3, 3],
) -> f32[n, m] {
  const r = sz / 2;

  @res let result : f32[n, m];

  @image_loop for row = 0 to n {
    for col = 0 to m {
      // Compute the pixel of dilated image
      let dilated = MIN_BR;
      @filter_loop for i = 0 to sz {
        for j = 0 to sz {
	  let filter = if row + i < r
                          || row + i - r > n - 1
                          || col + j < r
                          || col + j - r > m - 1 then MIN_BR
                       else if input[row + i - r, col + j - r] > MIN_BR then MAX_BR
                       else MIN_BR;
          dilated = max!(dilated, filter * structure[i, j]);
        }
      }

      // Compute the pixel of eroded image
      let eroded = MAX_BR;
      @filter_loop for i = 0 to sz {
        for j = 0 to sz {
	  let filter = if row + i < r
                          || row + i - r > n - 1
                          || col + j < r
                          || col + j - r > m - 1 then MAX_BR
                       else if input[row + i - r, col + j - r] > MIN_BR then MAX_BR
                       else MIN_BR;
          eroded = min!(eroded, filter * structure[i, j]);
        }
      }

      let sign = dilated - eroded;
      result[row, col] = sign;
    }
  }

  return result;
}

fn gradient<n, m : usize>(
  input: f32[n, m],
  sx: f32[3, 3],
  sy: f32[3, 3],
) -> f32[n, m] {
  const sbr = sb / 2;

  @res let result : f32[n, m];

  for row = 0 to n {
    for col = 0 to m {

      let gx = 0;
      let gy = 0;

      @filter_loop for i = 0 to sb {
        for j = 0 to sb {
          let val = input[if row + i < sbr              then 0
                          else if row + i - sbr > n - 1 then n - 1
                                                        else row + i - sbr,
                          if col + j < sbr              then 0
                          else if col + j - sbr > m - 1 then m - 1
                                                        else col + j - sbr];
          gx += val * sx[i, j];
          gy += val * sy[i, j];
        }
      }

      result[row, col] = sqrt!(gx * gx + gy * gy);
    }
  }

  return result;
}

fn max_gradient<n, m: usize>(gradient: f32[n, m]) -> f32 {
  let max = -1.0;

  for i = 0 to n {
    for j = 0 to m {
      max = max!(max, gradient[i, j]);
    }
  }

  return max;
}

fn reject_zero_crossings<n, m: usize>(
  crossings: f32[n, m],
  gradient: f32[n, m],
  max_gradient: f32,
  theta: f32,
) -> f32[n, m] {
  @res let result : f32[n, m];

  for row = 0 to n {
    for col = 0 to m {
      result[row, col] =
        if crossings[row, col] > 0 && gradient[row, col] > theta * max_gradient
        then 1.0
        else 0.0;
    }
  }

  return result;
}

#[entry]
fn edge_detection<n, m : usize>(
  input: f32[n, m],
  gaussian_filter: f32[gs, gs],
  structure: f32[3, 3],
  sx: f32[3, 3],
  sy: f32[3, 3],
  theta: f32,
) -> f32[n, m] {
  let smoothed = gaussian_smoothing::<n, m>(input, gaussian_filter);
  @le let laplacian = laplacian_estimate::<n, m>(smoothed, structure);
  @zc let zcs = zero_crossings::<n, m>(laplacian, structure);
  let gradient = gradient::<n, m>(smoothed, sx, sy);
  let maxgrad = max_gradient::<n, m>(gradient);
  return reject_zero_crossings::<n, m>(zcs, gradient, maxgrad, theta);
}