eorst 0.1.0

// use gdal::spatial_ref;
use colored::*;
use gdal::{raster::Buffer, Dataset, DatasetOptions, GdalOpenFlags};
use itertools::Itertools;
use math::round;
use ndarray::{s, Array3};
use rayon::iter::IntoParallelIterator;
use rayon::iter::IntoParallelRefIterator;
use rayon::iter::ParallelIterator;
use std::cmp::{max, min};
use std::collections::BTreeMap;
use std::convert::TryInto;
use std::fmt::Debug;
use std::path::{Path, PathBuf};

// use std::fs;
// use std::env::temp_dir;

use crate::common::*;
// use std::process::Command;

use log::{debug, info};
extern crate env_logger;

#[test]
fn test_n_bands() {
    let source = "data/cemsre_t55jfm_20200614_sub_abam5.tif";
    assert_eq!(6, SingleRasterDatasetBuilder::get_n_bands(source));
}

// SingleRasterDatasets store metadata for single
// gdal compatible raster files.
// .read() method always returns an Array3
#[derive(PartialEq, Debug, Clone)]
pub struct SingleRasterDataset {
    source: String,
    pub(crate) bands: Vec<usize>,
    block_size: BlockSize,
    pub(crate) image_size: ImageSize,
    pub n_blocks: usize,
    pub(crate) geo_transform: GeoTransform,
    pub(crate) epsg_code: u32,
    pub overlap_size: usize,
    pub blocks_attributes: Vec<BlockAttributes>,
    pub(crate) num_threads: usize,
    counter: usize,
}
unsafe impl Send for SingleRasterDataset {}

impl Iterator for SingleRasterDataset {
    type Item = SingleRasterDataset;

    fn next(&mut self) -> Option<SingleRasterDataset> {
        let mut srds = self.clone();
        self.counter += 1;
        // don't overflow the limit
        if self.counter > self.n_blocks {
            None
        }
        // ok we can loop
        else {
            let srds_blocks_attributes = self.blocks_attributes[srds.counter];
            srds.n_blocks = 1;
            srds.blocks_attributes = vec![srds_blocks_attributes];
            srds.geo_transform = srds_blocks_attributes.geo_transform;
            srds.image_size = ImageSize {
                rows: self.block_size.rows,
                cols: self.block_size.cols,
            };
            Some(srds)
        }
    }
}

use rayon::iter::plumbing::UnindexedConsumer;
// use rayon::iter::plumbing::Consumer;
// use rayon::iter::plumbing::par_bridge;

// impl ParallelIterator for SingleRasterDataset {
//     type Item = SingleRasterDataset;

//     fn drive_unindexed<C>(self, consumer: C) -> C::Result
//         where C: UnindexedConsumer<SingleRasterDataset>
//     {
//         bridge_indexed(self.next().unwrap(), consumer)

//     }
// }
impl ParallelIterator for SingleRasterDataset {
    type Item = SingleRasterDataset;
    // type Iter = SingleRasterDataset;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
    where
        C: UnindexedConsumer<Self::Item>,
    {
        self.into_par_iter().drive_unindexed(consumer)
    }
    // fn into_par_iter(self) -> Self::Item {
    //     self.into_iter()
    // }
}

#[derive(Default)]
pub struct SingleRasterDatasetBuilder {
    //like Iter
    // Probably lots of optional fields.
    source: String,
    bands: Vec<usize>,
    image_size: ImageSize,
    block_size: BlockSize,
    n_blocks: usize,
    geo_transform: GeoTransform,
    epsg_code: u32,
    overlap_size: usize,
    blocks_attributes: Vec<BlockAttributes>,
    num_threads: usize,
    counter: usize,
}

impl SingleRasterDataset {
    // This method will help users to discover the builder
    pub fn builder() -> SingleRasterDatasetBuilder {
        SingleRasterDatasetBuilder::default()
    }

    fn n_block_cols(&self) -> usize {
        let image_size = self.image_size;
        let block_size = self.block_size;
        round::ceil(image_size.cols as f64 / block_size.cols as f64, 0) as usize
    }

    fn block_col_row(&self, id: usize) -> (usize, usize) {
        let block_row = id as usize / self.n_block_cols();
        let block_col = id as usize - (block_row * self.n_block_cols());

        (block_col, block_row)
    }

    fn get_block_gt(&self, read_window: ReadWindow) -> GeoTransform {
        let x_ul_image = self.geo_transform.x_ul;
        let y_ul_image = self.geo_transform.y_ul;

        let x_res = self.geo_transform.x_res;
        let y_res = self.geo_transform.y_res;
        let x_pos = read_window.offset.cols;
        let y_pos = read_window.offset.rows;

        let x_ul_block = x_ul_image + x_res * x_pos as f64; // to fix! will break with overlap
        let y_ul_block = y_ul_image + y_res * y_pos as f64; // to fix! will break with overlap

        GeoTransform {
            x_ul: x_ul_block,
            x_res,
            x_rot: self.geo_transform.x_rot,
            y_ul: y_ul_block,
            y_rot: self.geo_transform.x_rot,
            y_res,
        }
    }
    fn block_from_id(&self, id: usize) -> BlockAttributes {
        let mut overlap = Overlap {
            left: 0,
            top: 0,
            right: 0,
            bottom: 0,
        };
        // get ul corner of the block
        let ul_x = (self.block_size.cols * self.block_col_row(id).0) as isize;
        let ul_y = (self.block_size.rows * self.block_col_row(id).1) as isize;

        // now compute with overlap
        let ul_x_overlap = max(0, ul_x - self.overlap_size as isize);
        let ul_y_overlap = max(0, ul_y - self.overlap_size as isize);
        let lr_x_overlap = min(
            self.image_size.cols as isize,
            ul_x + self.block_size.cols as isize + self.overlap_size as isize,
        );
        let lr_y_overlap = min(
            self.image_size.rows as isize,
            ul_y + self.block_size.rows as isize + self.overlap_size as isize,
        );

        let win_width = lr_x_overlap - ul_x_overlap;
        let win_height = lr_y_overlap - ul_y_overlap;

        let arr_width = win_width;
        let arr_height = win_height;

        let read_window = ReadWindow {
            offset: Offset {
                cols: ul_x_overlap,
                rows: ul_y_overlap,
            },
            size: Size {
                cols: arr_width,
                rows: arr_height,
            },
        };

        // add borders
        if (ul_x - self.overlap_size as isize) < 0 {
            // add left
            overlap.left = 1 * self.overlap_size;
        }
        if (ul_y - self.overlap_size as isize) < 0 {
            // add top
            overlap.top = 1 * self.overlap_size;
        }
        if (self.image_size.cols as isize)
            < (ul_x + self.block_size.cols as isize + self.overlap_size as isize)
        {
            // add right
            overlap.right = 1 * self.overlap_size;
        }

        if (self.image_size.rows as isize)
            < (ul_y + self.block_size.rows as isize + self.overlap_size as isize)
        {
            // add bottom
            overlap.bottom = 1 * self.overlap_size;
        }

        let block_geo_transform = self.get_block_gt(read_window);
        let block_attributes = BlockAttributes {
            block_index: id,
            read_window,
            overlap_size: self.overlap_size,
            geo_transform: block_geo_transform,
            overlap,
        };
        block_attributes
    }

    // read the full image as Array3
    pub fn read<T>(&self) -> Array3<T>
    where
        T: gdal::raster::GdalType + Copy,
    {
        let size = self.get_image_size();
        let read_window = ReadWindow {
            offset: Offset { rows: 0, cols: 0 },
            size: Size {
                rows: size.rows as isize,
                cols: size.cols as isize,
            },
        };
        self.read_window(read_window)
    }

    // reads a block of data (rows*cols + overlap)
    pub fn read_block<T>(&self, block_attributes: BlockAttributes) -> Array3<T>
    where
        T: gdal::raster::GdalType + Copy + num_traits::identities::Zero + std::fmt::Debug,
    {
        let overlap = block_attributes.overlap;
        let read_window = block_attributes.read_window;
        let data = self.read_window(read_window);
        // we need to add rows and cols if needed (ie shape is not block_size + 2 * overlap_size).
        // that info is stored in block_attributes.read_window.overlap
        let data = add_cols_left(&data, &overlap.left);
        let data = add_cols_right(&data, &overlap.right);
        let data = add_rows_bottom(&data, &overlap.bottom);
        let data = add_rows_top(&data, &overlap.top);

        data
    }

    // read a gdal_window as Array3
    pub fn read_window<T>(&self, read_window: ReadWindow) -> Array3<T>
    where
        T: gdal::raster::GdalType + Copy,
    {
        // todo!: implement bounds check. i.e offset.1 + size.1 can't be > n_rows.
        let raster_path = Path::new(&self.source);
        let bands: Vec<usize> = self.bands.to_vec();
        let n_bands = bands.len();
        let ds = Dataset::open(&raster_path).unwrap();
        let mut data: Vec<T> = Vec::new();
        for b in bands {
            let band = ds.rasterband(b as isize).unwrap();
            let mut band_data = band
                .read_as::<T>(
                    (read_window.offset.cols, read_window.offset.rows),
                    (
                        read_window.size.cols as usize,
                        read_window.size.rows as usize,
                    ),
                    (
                        read_window.size.cols as usize,
                        read_window.size.rows as usize,
                    ),
                    None,
                )
                .unwrap()
                .data;
            data.append(&mut band_data);
        }
        let array = Array3::from_shape_vec(
            (
                n_bands,
                read_window.size.rows as usize,
                read_window.size.cols as usize,
            ),
            data,
        )
        .unwrap();
        array
    }

    pub fn oom_apply<T>(&self, f: fn(&Array3<T>) -> Array3<T>, fn_prefix: &str)
    where
        T: gdal::raster::GdalType + Copy + num_traits::Zero + Debug,
    {
        let pool = rayon::ThreadPoolBuilder::new()
            .num_threads(self.num_threads)
            .build()
            .unwrap();
        pool.install(|| {
            (0..self.n_blocks)
                // .to_owned()
                .into_par_iter()
                .for_each(|id| {
                    let block_attributes = self.blocks_attributes[id];
                    let w = block_attributes.read_window;
                    let block_data = self.read_window::<T>(w);
                    let result = f(&block_data);
                    let out_fn = format!("{}_{}.tif", fn_prefix, id);
                    let out_fn = out_fn.as_str();
                    self.write_block3(id, result, out_fn) // save block to disk
                })
        })
    }

    pub fn write_block3<T>(&self, block_index: usize, data: Array3<T>, file_name: &str)
    where
        T: gdal::raster::GdalType + Copy + num_traits::identities::Zero + Debug,
    {
        // create and empty raster with the right size
        let epsg_code = self.epsg_code;
        let overlap_size = self.blocks_attributes[block_index].overlap_size;
        let trimmed = trimm_array3(&data, overlap_size);
        let size_y = trimmed.shape()[1];
        let size_x = trimmed.shape()[2];
        let dataset_options: DatasetOptions = DatasetOptions {
            open_flags: GdalOpenFlags::GDAL_OF_UPDATE,
            allowed_drivers: None,
            open_options: None,
            sibling_files: None,
        };
        let block_geotransform = self.blocks_attributes[block_index].geo_transform;
        let out_size: BlockSize = BlockSize {
            rows: size_y,
            cols: size_x,
        };
        let n_bands = data.shape()[0];
        raster_from_size::<T>(
            &PathBuf::from(file_name),
            block_geotransform,
            epsg_code,
            out_size,
            n_bands as isize,
        );

        let out_ds = Dataset::open_ex(Path::new(file_name), dataset_options).unwrap();
        for band in 0..trimmed.shape()[0] {
            let b = (band + 1) as isize;
            let mut out_band = out_ds.rasterband(b).unwrap();
            let data_vec: Vec<T> = trimmed
                .slice(s![band, .., ..])
                .into_iter()
                .map(|v| *v)
                .collect();

            let data_buffer = Buffer {
                size: (out_size.cols, out_size.rows),
                data: data_vec,
            };

            out_band
                .write((0, 0), (out_size.cols, out_size.rows), &data_buffer)
                .unwrap();
        }
    }

    fn get_image_size(&self) -> ImageSize {
        let source: &Path = Path::new(&self.source);
        let ds: Dataset = Dataset::open(source).unwrap();
        let size = ds.raster_size(); // x, y
        ImageSize {
            rows: size.1,
            cols: size.0,
        }
    }

    fn geo_to_global_rc(&self, point: Coordinates) -> Index2d {
        let gt = self.geo_transform.to_array();
        let row: usize = ((point.y - gt[3]) / gt[5]) as usize;
        let col: usize = ((point.x - gt[0]) / gt[1]) as usize;
        Index2d { col, row }
    }

    fn geoms_to_global_indices(
        &self,
        geoms: BTreeMap<i64, Vec<(f64, f64, f64)>>,
    ) -> BTreeMap<i64, Index2d> {
        debug!("Transforming geographic points to array indices.");

        let idx_global: BTreeMap<_, _> = geoms
            .par_iter()
            .map(|(pid, p)| {
                let point: Coordinates = Coordinates {
                    x: p[0].0,
                    y: p[0].1,
                };
                (*pid, self.geo_to_global_rc(point))
            })
            .collect();
        idx_global
    }

    fn id_from_indices(&self, index: Index2d) -> usize {
        let n_block_cols = self.n_block_cols();
        (index.col / self.block_size.cols) + (index.row / self.block_size.rows) * n_block_cols
    }

    fn global_rc_to_block_rc(&self, global_index: Index2d) -> Index2d {
        let mut block_col = global_index.col as isize % self.block_size.cols as isize;
        let mut block_row = global_index.row as isize % self.block_size.rows as isize;

        let block_col_ov = block_col + self.overlap_size as isize;
        let block_row_ov = block_row + self.overlap_size as isize;

        if (global_index.col as isize - block_col_ov) > 0 {
            block_col = block_col_ov;
        } else {
        };

        if global_index.row as isize - block_row_ov > 0 {
            block_row = block_row_ov;
        } else {
        };

        Index2d {
            col: block_col as usize,
            row: block_row as usize,
        }
    }

    fn block_id_rowcol(&self, pid: i64, index: Index2d) -> (usize, (i64, Index2d)) {
        // find the block id for a global index. Returns block_id, (point_id, row_col)
        let id = self.id_from_indices(index);
        let row_col = self.global_rc_to_block_rc(index);
        (id, (pid, row_col))
    }

    pub fn extract_blockwise(
        &self,
        vector_path: &str,
        id_col_name: &str,
        method: Option<&str>,
        buffer_size: Option<usize>,
    ) -> BTreeMap<i16, Vec<i16>> {
        let method: &str = method.unwrap_or("value");
        let buffer_size = buffer_size.unwrap_or(0);
        assert!(
            buffer_size <= self.overlap_size,
            "Buffer size > overlap size"
        );
        info!("Strarting extract parallel");
        info!("Method: {:?}", method);

        debug!("Opening vector dataset");
        let vector_dataset = Dataset::open(Path::new(vector_path)).unwrap();
        let mut layer = vector_dataset.layer(0).unwrap();
        debug!("Reading geometries as a BTreeMap (k: pid, v: geometries)");
        let mut geoms = BTreeMap::new();

        let fields_defn = layer
            .defn()
            .fields()
            .map(|field| (field.name(), field.field_type(), field.width()))
            .collect::<Vec<_>>();
        debug!("{:?}", fields_defn);

        for feature in layer.features() {
            let geom = feature.geometry().get_point_vec();
            let pid = feature
                .field(id_col_name)
                .unwrap()
                .unwrap()
                .into_int64()
                .unwrap();
            geoms.insert(pid, geom);
        }
        debug!("geoms: {:?}", geoms);
        // Find global index for each point. BtreeMap<point_id, global_index>
        let idx_global = self.geoms_to_global_indices(geoms);
        debug!("idx_global: {:?}", idx_global);

        debug!("Addding block id to the array indices.");
        let id_indices: Vec<(usize, (i64, Index2d))> = idx_global
            .par_iter()
            .map(|(pid, index)| self.block_id_rowcol(*pid, *index))
            .collect();
        debug!("Block id, local point coords: {:?}.", id_indices);

        debug!("Find what blocks need to be processed.");
        let block_ids: Vec<_> = idx_global
            .par_iter()
            .map(|(_, index)| self.id_from_indices(*index))
            .collect();
        debug!("block_ids: {:?}.", block_ids);

        let blocks_to_process: Vec<_> = block_ids.iter().unique().collect();
        debug!("Will process blocks: {:?}.", blocks_to_process);

        let pool = rayon::ThreadPoolBuilder::new().build().unwrap();
        let handle = pool.install(|| {
            blocks_to_process
                .into_par_iter()
                .map(|id| -> (Vec<usize>, Vec<usize>, Vec<Vec<i16>>) {
                    //id is block id, not point id
                    // let id = *id;
                    // print!("block id: {:?}.", id);
                    let mut pos: Vec<Index2d> = Vec::new();
                    let mut idx: Vec<usize> = Vec::new();
                    let mut pids: Vec<usize> = Vec::new();
                    for (pid, p) in id_indices.iter().enumerate() {
                        if p.0 == *id {
                            pos.push(Index2d {
                                col: p.1 .1.col,
                                row: p.1 .1.row,
                            });
                            pids.push(pid);
                            idx.push(p.1 .0 as usize);
                        } else {
                        };
                    }
                    debug!("Pos: {:?}.", pos);
                    debug!("Pids: {:?}.", pids);

                    let mut res = Vec::new();
                    let rect: Rectangle = Rectangle {
                        left: buffer_size,
                        top: buffer_size,
                        right: buffer_size,
                        bottom: buffer_size,
                    };

                    let block_attributes = self.block_from_id(*id);
                    let data = self.read_block::<i16>(block_attributes);
                    let bands = data.shape()[0];
                    for band_n in 0..bands {
                        let mut res_band = Vec::new();
                        let band_data = data.slice(s![band_n, .., ..]);
                        for (_, point) in pos.iter().enumerate() {
                            let point_shifted = Index2d {
                                col: point.col + self.blocks_attributes[*id].overlap.left,
                                row: point.row + self.blocks_attributes[*id].overlap.top,
                            };
                            match method {
                                "mode" => {
                                    let mut window_data: Vec<_> =
                                        rect_view(&band_data, rect, point_shifted)
                                            .iter()
                                            .map(|x| x.clone())
                                            .collect();
                                    window_data.sort_by(|a, b| a.partial_cmp(b).unwrap());
                                    let median = window_data[window_data.len() / 2];
                                    res_band.push(median);
                                }
                                "value" => {
                                    let d = band_data[(point_shifted.row, point_shifted.col)];
                                    res_band.push(d);
                                }
                                _ => println!("ERROR HERE!"),
                            }
                        }
                        res.push(res_band);
                    }
                    (pids, idx, res)
                })
        });

        let collected: Vec<_> = handle.collect();
        let pids: Vec<_> = collected.iter().map(|(pid, _, _)| pid).collect();
        let vals: Vec<_> = collected.iter().map(|(_, _, vals)| vals).collect();
        let idxs: Vec<_> = collected.iter().map(|(_, idx, _)| idx).collect();
        let mut results = BTreeMap::new();
        info!("idx len -> {:?}", idxs.len());
        debug!("pids -> {:?}", pids);
        debug!("vals -> {:?}", vals);

        let num_bands = vals[0].len();
        let num_blocks = pids.len();
        for block in 0..num_blocks {
            for i in 0..pids[block].len() {
                let mut vals_point: Vec<i16> = Vec::new();
                let id = idxs[block][i];
                for band in 0..num_bands {
                    vals_point.push(vals[block][band][i]);
                }
                let mut res_point = BTreeMap::new();
                res_point.insert(id as i16, vals_point);
                results.append(&mut res_point);
            }
        }

        results
    }
}

impl SingleRasterDatasetBuilder {
    pub fn from_file(source: &str) -> SingleRasterDatasetBuilder {
        // we initialize all the cheap bits here. The rest, we do in build()
        // Is important to note, that if a different projection is assigned to
        // the dataset, then most of this values are going to be is invalid.
        // Not sure if this was the best choice.

        let overlap_size: usize = 0;
        let n_bands = SingleRasterDatasetBuilder::get_n_bands(source);
        let bands: Vec<usize> = (0..n_bands).into_iter().map(|v| v + 1).collect();
        let block_size: BlockSize = BlockSize {
            cols: 1024,
            rows: 1024,
        };
        let geo_transform = SingleRasterDatasetBuilder::get_geotransform(source);
        let image_size = SingleRasterDatasetBuilder::get_image_size(source);
        let epsg_code = SingleRasterDatasetBuilder::get_epsg_code(source);
        let n_blocks = 0;
        let blocks_attributes: Vec<BlockAttributes> = Vec::new();
        //
        SingleRasterDatasetBuilder {
            source: String::from(source),
            bands,
            image_size,
            block_size,
            n_blocks,
            geo_transform,
            epsg_code,
            blocks_attributes,
            overlap_size,
            num_threads: 4,
            counter: 0,
        }
    }
    pub fn epsg(mut self, epsg_code: u32) -> SingleRasterDatasetBuilder {
        self.epsg_code = epsg_code;
        self
    }

    pub fn bands(mut self, bands: Vec<usize>) -> SingleRasterDatasetBuilder {
        // Set the name on the builder itself, and return the builder by value.
        self.bands = bands;
        self
    }

    pub fn block_size(mut self, block_size: BlockSize) -> SingleRasterDatasetBuilder {
        // Set the name on the builder itself, and return the builder by value.
        self.block_size = block_size;
        self
    }

    pub fn overlap_size(mut self, overlap_size: usize) -> SingleRasterDatasetBuilder {
        // Set the name on the builder itself, and return the builder by value.
        self.overlap_size = overlap_size;
        self
    }

    pub fn num_threads(mut self, num_threads: usize) -> SingleRasterDatasetBuilder {
        // Set the name on the builder itself, and return the builder by value.
        self.num_threads = num_threads;
        self
    }

    // If we can get away with not consuming the Builder here, that is an
    // advantage. It means we can use the FooBuilder as a template for constructing
    // many Foos.
    pub fn build(mut self) -> SingleRasterDataset {
        // We need to re-init the parameters related to geometry in case the source
        // in now a vrt in a different projection!.
        let current_epsg = SingleRasterDatasetBuilder::get_epsg_code(&self.source);
        let target_epsg = self.epsg_code;
        let rebuild = current_epsg != target_epsg;
        if rebuild {
            let new_source = create_temp_file("vrt");
            // make a vrt with the new epsg. I also need to think about the resampling algorithm.
            let argv = &[
                "gdalwarp",
                "-t_srs",
                &format!("EPSG:{}", target_epsg),
                &self.source,
                &new_source,
            ];

            std::process::Command::new(&argv[0])
                .args(&argv[1..])
                .spawn()
                .expect("failed to start creating vrt")
                .wait()
                .expect("failed to wait for the vrt");

            self.source = new_source.to_string();
            self.geo_transform = SingleRasterDatasetBuilder::get_geotransform(&self.source);
            self.image_size = SingleRasterDatasetBuilder::get_image_size(&self.source);
        }

        let (n_blocks, blocks_attributes) = self.get_blocks_attributes();
        self.n_blocks = n_blocks;
        self.blocks_attributes = blocks_attributes;

        SingleRasterDataset {
            source: self.source,
            bands: self.bands,
            image_size: self.image_size,
            block_size: self.block_size,
            n_blocks: self.n_blocks,
            geo_transform: self.geo_transform,
            epsg_code: self.epsg_code,
            blocks_attributes: self.blocks_attributes,
            overlap_size: self.overlap_size,
            num_threads: self.num_threads,
            counter: self.counter,
        }
    }

    fn get_n_bands(source: &str) -> usize {
        let source: &Path = Path::new(source);
        let ds: Dataset = Dataset::open(source).unwrap();
        ds.raster_count() as usize
    }

    fn get_geotransform(source: &str) -> GeoTransform {
        let source: &Path = Path::new(&source);
        let ds: Dataset = Dataset::open(source).unwrap();
        let geo_transform = ds.geo_transform().unwrap();
        GeoTransform {
            x_ul: geo_transform[0],
            x_res: geo_transform[1],
            x_rot: geo_transform[2],
            y_ul: geo_transform[3],
            y_rot: geo_transform[4],
            y_res: geo_transform[5],
        }
    }
    fn get_image_size(source: &str) -> ImageSize {
        let source: &Path = Path::new(source);
        let ds: Dataset = Dataset::open(source).unwrap();
        let size = ds.raster_size(); // x, y
        ImageSize {
            rows: size.1,
            cols: size.0,
        }
    }
    fn get_epsg_code(source: &str) -> u32 {
        let source: &Path = Path::new(source);
        let ds: Dataset = Dataset::open(source).unwrap();
        let spatial_ref = ds.spatial_ref().unwrap();
        let epsg_code = spatial_ref.auth_code().expect("No auth_code");
        epsg_code.try_into().unwrap()
    }

    fn n_block_cols(&self) -> usize {
        let image_size = self.image_size;
        let block_size = self.block_size;
        round::ceil(image_size.cols as f64 / block_size.cols as f64, 0) as usize
    }

    fn n_block_rows(&self) -> usize {
        let image_size = self.image_size;
        let block_size = self.block_size;
        round::ceil(image_size.rows as f64 / block_size.rows as f64, 0) as usize
    }

    fn n_blocks(&self) -> usize {
        self.n_block_cols() * self.n_block_rows()
    }

    fn block_col_row(&self, id: usize) -> (usize, usize) {
        let block_row = id as usize / self.n_block_cols();
        let block_col = id as usize - (block_row * self.n_block_cols());
        (block_col, block_row)
    }

    fn get_block_gt(&self, read_window: ReadWindow) -> GeoTransform {
        let x_ul_image = self.geo_transform.x_ul;
        let y_ul_image = self.geo_transform.y_ul;

        let x_res = self.geo_transform.x_res;
        let y_res = self.geo_transform.y_res;
        let x_pos = read_window.offset.cols;
        let y_pos = read_window.offset.rows;

        let x_ul_block = x_ul_image + x_res * x_pos as f64; // to fix! will break with overlap
        let y_ul_block = y_ul_image + y_res * y_pos as f64; // to fix! will break with overlap

        GeoTransform {
            x_ul: x_ul_block,
            x_res,
            x_rot: self.geo_transform.x_rot,
            y_ul: y_ul_block,
            y_rot: self.geo_transform.x_rot,
            y_res,
        }
    }

    fn block_from_id(&self, id: usize) -> BlockAttributes {
        let mut overlap = Overlap {
            left: 0,
            top: 0,
            right: 0,
            bottom: 0,
        };
        // get ul corner of the block
        let ul_x = (self.block_size.cols * self.block_col_row(id).0) as isize;
        let ul_y = (self.block_size.rows * self.block_col_row(id).1) as isize;

        // now compute with overlap
        let ul_x_overlap = max(0, ul_x - self.overlap_size as isize);
        let ul_y_overlap = max(0, ul_y - self.overlap_size as isize);
        let lr_x_overlap = min(
            self.image_size.cols as isize,
            ul_x + self.block_size.cols as isize + self.overlap_size as isize,
        );
        let lr_y_overlap = min(
            self.image_size.rows as isize,
            ul_y + self.block_size.rows as isize + self.overlap_size as isize,
        );

        let win_width = lr_x_overlap - ul_x_overlap;
        let win_height = lr_y_overlap - ul_y_overlap;

        let arr_width = win_width;
        let arr_height = win_height;

        let read_window = ReadWindow {
            offset: Offset {
                cols: ul_x_overlap,
                rows: ul_y_overlap,
            },
            size: Size {
                cols: arr_width,
                rows: arr_height,
            },
        };

        // add borders
        if (ul_x - self.overlap_size as isize) < 0 {
            // add left
            overlap.left = 1 * self.overlap_size;
        }
        if (ul_y - self.overlap_size as isize) < 0 {
            // add top
            overlap.top = 1 * self.overlap_size;
        }
        if (self.image_size.cols as isize)
            < (ul_x + self.block_size.cols as isize + self.overlap_size as isize)
        {
            // add right
            overlap.right = 1 * self.overlap_size;
        }

        if (self.image_size.rows as isize)
            < (ul_y + self.block_size.rows as isize + self.overlap_size as isize)
        {
            // add bottom
            overlap.bottom = 1 * self.overlap_size;
        }

        let block_geo_transform = self.get_block_gt(read_window);
        let block_attributes = BlockAttributes {
            block_index: id,
            read_window,
            overlap_size: self.overlap_size,
            geo_transform: block_geo_transform,
            overlap,
        };
        block_attributes
    }

    pub fn get_blocks_attributes(&self) -> (usize, Vec<BlockAttributes>) {
        let n_blocks = self.n_blocks();
        let mut blocks_attributes: Vec<BlockAttributes> = Vec::new();
        (0..n_blocks).into_iter().for_each(|i| {
            let block_attributes = self.block_from_id(i);
            blocks_attributes.push(block_attributes)
        });
        (n_blocks, blocks_attributes)
    }
}

impl std::fmt::Display for SingleRasterDataset {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::result::Result<(), std::fmt::Error> {
        // writeln!(fmt, "{} \t {}", "Source: ".bold(), self.source).unwrap();

        writeln!(fmt, "{}", "Coordinates: ".bold()).unwrap();

        writeln!(fmt, "\t {}", "x ".bold()).unwrap();
        writeln!(fmt, "\t {}", "y ".bold()).unwrap();
        writeln!(fmt, "\t {} \t \t \t {:?}", "Bands".bold(), self.bands).unwrap();

        writeln!(fmt, "{}", "Attributes: ".bold()).unwrap();
        writeln!(
            fmt,
            "\t {} \t \t ( r: {:?}, c: {:?} )",
            "image_size: ".bold(),
            self.image_size.rows,
            self.image_size.cols
        )
        .unwrap();
        writeln!(
            fmt,
            "\t {} \t \t ( {:?}, {:?} )",
            "block_size: ".bold(),
            self.block_size.rows,
            self.block_size.cols
        )
        .unwrap();
        writeln!(fmt, "\t {} \t \t {:?}", "n_blocks: ".bold(), self.n_blocks,).unwrap();
        writeln!(
            fmt,
            "\t {} \t {:?}",
            "geo_transform: ".bold(),
            self.geo_transform
        )
        .unwrap();
        writeln!(
            fmt,
            "\t {} \t \t {:?} ",
            "epsg_code: ".bold(),
            self.epsg_code
        )
    }
}