eorst 0.1.0

An Earth Observation and Remote Sensing Toolkit
use gdal::{raster::RasterCreationOption, spatial_ref::SpatialRef, Driver};
use ndarray::{s, Array3, ArrayView2, ArrayView3};
use std::path::Path;
use std::path::PathBuf;
use uuid::Uuid;

#[cfg(feature = "use_opencv")]
use opencv::{
    core::{Mat_AUTO_STEP, Scalar},
    imgproc,
    prelude::*,
};

// A location on an array
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Index2d {
    pub col: usize,
    pub row: usize,
}

#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub(crate) struct Coordinates {
    // x and y coordinates of a location.
    pub(crate) x: f64,
    pub(crate) y: f64,
}
#[derive(Debug, Copy, Clone)]
pub struct Rectangle {
    // I need better names for this!
    pub left: usize,
    pub top: usize,
    pub right: usize,
    pub bottom: usize,
}
#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub struct BlockSize {
    pub rows: usize,
    pub cols: usize,
}

#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub struct ImageSize {
    pub rows: usize,
    pub cols: usize,
}

#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub struct GeoTransform {
    pub x_ul: f64,
    pub x_res: f64,
    pub x_rot: f64,
    pub y_ul: f64,
    pub y_rot: f64,
    pub y_res: f64,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub struct BlockAttributes {
    pub block_index: usize,
    pub read_window: ReadWindow,
    pub overlap_size: usize,
    pub geo_transform: GeoTransform,
    pub overlap: Overlap,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub struct Offset {
    pub rows: isize,
    pub cols: isize,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub struct Size {
    pub rows: isize,
    pub cols: isize,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub struct ReadWindow {
    pub offset: Offset,
    pub size: Size,
}

#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Overlap {
    // Describes how many cells have to be added to the read array
    // to produce a window. For example if overlap_size=1, then the top left
    // block will need one row to the left and one to the top to be added.
    // overlap should be {1,1,0,0}; for the lower left corner in general
    // overlap shoud be {0,0,1,1}
    pub left: usize,
    pub top: usize,
    pub right: usize,
    pub bottom: usize,
}

// extends an array to the left by n columns filling the new rows
// with the data of the first n columns
pub fn add_cols_left<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
    T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
    let n = n.to_owned();
    if n > 0 {
        let r = original.shape()[1];
        let c = original.shape()[2];
        let b = original.shape()[0];
        let mut result: Array3<T> = ndarray::Array::zeros((b, r, c + n));
        result.slice_mut(s![.., .., n..]).assign(&original);

        let first_data_cols = result.slice(s![.., .., n..n + n]).to_owned();
        result.slice_mut(s![.., .., 0..n]).assign(&first_data_cols);
        return result;
    } else {
        return original.to_owned();
    }
}

pub fn add_rows_top<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
    T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
    let n = n.to_owned();
    if n > 0 {
        let r = original.shape()[1];
        let c = original.shape()[2];
        let b = original.shape()[0];

        let mut result: Array3<T> = ndarray::Array::zeros((b, r + n, c));
        result.slice_mut(s![.., n.., ..]).assign(&original);

        let first_data_rows = result.slice(s![.., n..n + n, ..]).to_owned();
        result.slice_mut(s![.., 0..n, ..]).assign(&first_data_rows);
        return result;
    } else {
        return original.to_owned();
    }
}

pub fn add_cols_right<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
    T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
    let n = n.to_owned();
    if n > 0 {
        let r = original.shape()[1];
        let c = original.shape()[2];
        let b = original.shape()[0];

        let mut result: Array3<T> = ndarray::Array::zeros((b, r, c + n));
        let idx = -1 * n as isize;
        result.slice_mut(s![.., .., ..idx]).assign(&original);

        let f = c as isize - n as isize;
        let t = c as isize;
        let last_data_cols = result.slice(s![.., .., f..t]).to_owned();
        result.slice_mut(s![.., .., c..]).assign(&last_data_cols);
        return result;
    } else {
        return original.to_owned();
    }
}

pub fn add_rows_bottom<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
    T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
    let n = n.to_owned();
    if n > 0 {
        let r = original.shape()[1];
        let c = original.shape()[2];
        let b = original.shape()[0];

        let mut result: Array3<T> = ndarray::Array::zeros((b, r + n, c));
        let idx = -1 * n as isize;
        result.slice_mut(s![.., ..idx, ..]).assign(&original);
        let f = r as isize - n as isize;
        let t = r as isize;
        let last_data_row = result.slice(s![.., f..t, ..]).to_owned();
        result.slice_mut(s![.., r.., ..]).assign(&last_data_row);
        return result;
    } else {
        return original.to_owned();
    }
}

pub fn trimm_array3<'a, T>(array: &'a Array3<T>, overlap_size: usize) -> ArrayView3<'a, T> {
    let min_row = overlap_size;
    let max_row = array.shape()[1] - overlap_size;

    let min_col = overlap_size;
    let max_col = array.shape()[2] - overlap_size;

    array.slice(s![.., min_row..max_row, min_col..max_col])
}

impl GeoTransform {
    pub fn to_array(&self) -> [f64; 6] {
        let array: [f64; 6] = [
            self.x_ul, self.x_res, self.x_rot, self.y_ul, self.y_rot, self.y_res,
        ];
        array
    }
}

// Creates and empty raster to store the ouptuts of a BlockWorker
pub fn raster_from_size<T>(
    file_name: &PathBuf,
    geo_transform: GeoTransform,
    epsg_code: u32,
    block_size: BlockSize,
    n_bands: isize,
) where
    T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
    let size_x = block_size.cols;
    let size_y = block_size.rows;
    let driver = Driver::get("GTiff").unwrap();
    let options = [RasterCreationOption {
        key: "COMPRESS",
        value: "LZW",
    }];
    let mut dataset = driver
        .create_with_band_type_with_options::<T>(
            file_name.to_str().unwrap(),
            size_x as isize,
            size_y as isize,
            n_bands,
            &options,
        )
        .unwrap();
    dataset
        .set_geo_transform(&geo_transform.to_array())
        .unwrap();
    let srs = SpatialRef::from_epsg(epsg_code).unwrap();
    dataset.set_spatial_ref(&srs).unwrap();
}

pub fn create_temp_file(ext: &str) -> String {
    let dir = Path::new("/tmp/");
    println!("{}", dir.to_str().unwrap());

    let file_name = format!("{}.{}", Uuid::new_v4(), ext);
    let file_name = dir.join(file_name);
    let file_name_string = file_name.into_os_string().into_string().unwrap();
    file_name_string
}

pub fn rect_view<'a, T>(
    array: &'a ArrayView2<T>,
    rect: Rectangle,
    indices: Index2d,
) -> ArrayView2<'a, T> {
    let max_row = indices.row + rect.bottom + 1;
    let min_row = indices.row - rect.top;

    let max_col = indices.col + rect.right + 1;
    let min_col = indices.col - rect.left;

    array.slice(s![min_row..max_row, min_col..max_col])
}

// transpose a vector of vectors.
pub fn transpose<T>(v: Vec<Vec<T>>) -> Vec<Vec<T>>
where
    T: Clone,
{
    assert!(!v.is_empty());
    (0..v[0].len())
        .map(|i| v.iter().map(|inner| inner[i].clone()).collect::<Vec<T>>())
        .collect()
}

// translates an ArrayView2 to an opencv Mat.
#[cfg(feature = "use_opencv")]
pub fn arrayview2_to_mat<T>(data: ArrayView2<T>) -> Mat
where
    T: gdal::raster::GdalType + opencv::prelude::DataType,
{
    // transforms an Array2 into a opencv Mat data type.
    let rows = data.shape()[0];
    let data_slice = data.as_slice().unwrap();
    let m = Mat::from_slice(data_slice).unwrap();
    let m = m.reshape(0, rows as i32).unwrap();
    m
}