kriging 0.0.1

use std::path::PathBuf;
use gdal::raster::Buffer;
use gdal::Dataset;
use gdal::DriverManager;
use gdal::raster::RasterBand;
use gdal::GeoTransform;
use gdal::GeoTransformEx;
use nalgebra::DVector;
use std::rc::Rc;

/// Reads the geographical coordinates of all pixels in a raster file.
///
/// Each pixel's center coordinates are computed using the raster's GeoTransform.
///
/// # Arguments
///
/// * `path` - A `PathBuf` pointing to the raster file.
///
/// # Returns
///
/// A `Result` containing a vector of `(longitude, latitude)` tuples, or a GDAL error.
pub fn read_raster_coords(path: &PathBuf) -> gdal::errors::Result<Vec<(f64, f64)>> {
    let dataset = Dataset::open(path)?;
    let geo_transform = dataset.geo_transform()?; // [origin_x, pixel_width, 0, origin_y, 0, pixel_height]

    let (width, height) = (dataset.raster_size().0, dataset.raster_size().1);
    let mut coords = Vec::with_capacity(width * height);

    for y in 0..height {
        for x in 0..width {
            let px = geo_transform[0] + (x as f64 + 0.5) * geo_transform[1];
            let py = geo_transform[3] + (y as f64 + 0.5) * geo_transform[5];
            coords.push((px, py));
        }
    }

    Ok(coords)
}

/// Exports a vector of prediction values to a GeoTIFF raster file.
///
/// The output raster uses the same geotransform and projection as the input template raster.
///
/// # Arguments
///
/// * `template_tif_path` - Path to a reference raster file for spatial metadata.
/// * `output_tif_path` - Path to write the output raster file.
/// * `predictions` - A flat array of prediction values, matching the pixel count of the template raster.
///
/// # Returns
///
/// A `Result` indicating success or a GDAL error.
pub fn export_predictions_to_raster(
    template_tif_path: &PathBuf,
    output_tif_path: &PathBuf,
    predictions: &[f64],
) -> gdal::errors::Result<()> {
    // Open template raster (to get geotransform, projection, etc.)
    let template = Dataset::open(template_tif_path)?;

    let (size_x, size_y) = (template.raster_size().0, template.raster_size().1);
    assert_eq!(
        predictions.len(),
        (size_x * size_y) as usize,
        "Prediction vector length doesn't match raster size"
    );

    // Create a new raster with same dimensions and projection
    let driver = DriverManager::get_driver_by_name("GTiff")?;
    let mut output = driver.create_with_band_type::<f64, &PathBuf>(
        output_tif_path,
        size_x,
        size_y,
        1
    )?;

    // Copy geotransform and projection
    output.set_projection(&template.projection())?;
    output.set_geo_transform(&template.geo_transform()?)?;

    // Create a buffer with the same size as the predictions array
    let mut buffer = Buffer::<f64>::new((size_x, size_y), predictions.to_vec());

    // Write predictions to band 1
    let mut band: RasterBand = output.rasterband(1)?; // Get the first band
    band.write((0, 0), (size_x, size_y), &mut buffer)?;
    band.set_no_data_value(Some(f64::NAN))?;

    Ok(())
}

/// A utility for managing raster-based drift functions used in universal kriging.
///
/// `RasterDrift` loads one or more raster layers (e.g., topographic, vegetation index)
/// and uses them as spatial covariates for modeling the trend component in kriging.
pub struct RasterDrift {
    _datasets: Vec<Rc<Dataset>>,
    bands: Vec<RasterBand<'static>>,
    geo_transform: GeoTransform,
    inv_transform: GeoTransform,
    coefficients: Option<DVector<f64>>
}

impl RasterDrift {
    /// Creates a new `RasterDrift` from a list of raster file paths.
    ///
    /// Each raster must contain exactly one band and share the same spatial extent.
    ///
    /// # Arguments
    ///
    /// * `raster_paths` - A list of string paths to raster files.
    ///
    /// # Panics
    ///
    /// Panics if no rasters are provided or if any file fails to open.
    pub fn new(raster_paths: &[&str]) -> Self {
        assert!(!raster_paths.is_empty(), "At least one raster is required");

        let mut datasets = Vec::new();
        let mut bands = Vec::new();

        for path in raster_paths {
            let ds = Rc::new(Dataset::open(path).expect("Failed to open raster"));
            //let band = ds.rasterband(1).expect("Missing raster band");

            // Convert band to 'static by leaking Rc<Dataset>
            let leaked_ds: &'static Dataset = Box::leak(Box::new(Rc::clone(&ds)));
            let static_band = leaked_ds.rasterband(1).unwrap();

            datasets.push(ds);
            bands.push(static_band);
        }

        let geo_transform = datasets[0].geo_transform().expect("Missing geo transform");
        let inv_transform = geo_transform.invert().expect("Failed to invert transform");
        //let inv_transform = inverse_geotransform(&geo_transform).expect("Failed to invert transform");

        Self {
            _datasets: datasets, // Hold on to the datasets
            bands,
            geo_transform,
            inv_transform,
            coefficients: None
        }
    }

    /// Returns the current drift coefficients if they have been set.
    ///
    /// # Returns
    ///
    /// An `Option` containing a vector of drift coefficients.
    pub fn get_coefficients(&self) -> Option<DVector<f64>> {
        self.coefficients.clone()
    }

    /// Sets the drift coefficients for predicting the trend component.
    ///
    /// # Arguments
    ///
    /// * `coeffs` - A vector of coefficients to be applied to drift functions.
    pub fn set_coefficients(&mut self, coeffs: DVector<f64>) {
        self.coefficients = Some(coeffs);
    }

    /// Predicts the trend component at a given geographic coordinate.
    ///
    /// Requires that coefficients have already been set.
    ///
    /// # Arguments
    ///
    /// * `point` - A `(longitude, latitude)` tuple.
    ///
    /// # Returns
    ///
    /// The predicted trend value at the given location.
    ///
    /// # Panics
    ///
    /// Panics if coefficients are not set.
    pub fn predict(&self, point: (f64, f64)) -> f64 {
        let drift_vec = self.drift_functions(point);
        let coeffs = self.coefficients.as_ref()
            .expect("Drift coefficients not set");
        drift_vec.iter().zip(coeffs.iter()).map(|(d, c)| d * c).sum()
    }

    /// Evaluates all drift functions at a given coordinate.
    ///
    /// The drift vector includes a constant intercept term (1.0) followed by values
    /// from each raster at the specified location.
    ///
    /// Out-of-bounds accesses are handled gracefully by returning zero for those values.
    ///
    /// # Arguments
    ///
    /// * `point` - A `(longitude, latitude)` tuple.
    ///
    /// # Returns
    ///
    /// A vector representing the drift features at the point.
    pub fn drift_functions(&self, point: (f64, f64)) -> Vec<f64> {
        let (x , y) = &self.inv_transform.apply(point.0, point.1);
        let px = *x as isize;
        let py = *y as isize;
    
        // Check if pixel coordinates are within bounds
        if px < 0 || py < 0 || px >= self.bands[0].x_size() as isize || py >= self.bands[0].y_size() as isize {
            // If out of bounds, return NaN for this raster
            //println!("Out of bounds in pixel {:?} {:?}", px, py);
            //return vec![f64::NAN, f64::NAN, f64::NAN];
            return vec![0.0, 0.0, 0.0];
        }
    
        let mut drift = vec![1.0];
        for band in &self.bands {
            let value = band
                .read_as::<f64>((px, py), (1, 1), (1, 1), None)
                .expect("Failed to read raster value")
                .into_iter()
                .next()
                .expect("No value found for the pixel");
            drift.push(value);
        }
        drift
    }

    /// Returns the number of drift functions, including the intercept.
    ///
    /// This equals 1 (for the intercept) plus the number of raster bands used.
    ///
    /// # Returns
    ///
    /// The dimensionality of the drift vector.
    pub fn drift_dimension(&self) -> usize {
        self.bands.len() + 1
    }
}