kriging 0.0.1

use std::path::PathBuf;
use plotters::prelude::*;
use haversine_rs::point::Point;
use haversine_rs::units::Unit;
use haversine_rs::distance;
use std::collections::HashMap;

/// Represents an empirical variogram, containing semivariances at different binned distances.
pub struct EmpiricalVariogram {
    pub distances: Vec<f64>,
    pub semivariances: Vec<f64>
}

impl EmpiricalVariogram {
    /// Constructs a new `EmpiricalVariogram` from given coordinates and values.
    ///
    /// # Arguments
    ///
    /// * `longitudes` - Vector of longitude values
    /// * `latitudes` - Vector of latitude values
    /// * `values` - Vector of measured values at the corresponding coordinates
    ///
    /// # Returns
    ///
    /// An `EmpiricalVariogram` instance with binned distances and semivariances
    pub fn new(
        longitudes: &Vec<f64>, 
        latitudes: &Vec<f64>,
        values: &Vec<f64>
    ) -> Self {
        // Store values as points
        let mut points = vec![];
        for i in 0..longitudes.len() {
            if let (Some(lon), Some(lat), Some(value)) = (longitudes.get(i), latitudes.get(i), values.get(i)) {
                points.push((Point::new(*lat, *lon), value));
            }
        }

        // Compute all pairwise distances
        let mut distances = vec![];
        for i in 0..points.len() {
            for j in (i + 1)..points.len() {
                let (p1, _) = &points[i];
                let (p2, _) = &points[j];
                let dist = distance(*p1, *p2, Unit::Meters);
                distances.push(dist);
            }
        }

        // Dynamic bin settings
        let max_distance = distances
            .iter()
            .cloned()
            .fold(0.0 / 0.0, f64::max); // max()
        let max_dist = 0.5 * max_distance;
        let num_bins = 15;
        let bin_size = max_dist / num_bins as f64;

        /*println!("Max pairwise distance: {:.2} m", max_distance);
        println!("Using max_dist = {:.2} m and bin_size = {:.2} m", max_dist, bin_size);*/

        // Compute semivariogram
        let mut semivariance_bins: HashMap<usize, (f64, usize)> = HashMap::new();

        for i in 0..points.len() {
            for j in (i + 1)..points.len() {
                let (p1, t1) = points[i];
                let (p2, t2) = points[j];
                let dist = distance(p1, p2, Unit::Meters);

                if dist <= max_dist {
                    let bin = (dist / bin_size).floor() as usize;
                    let gamma = 0.5 * (t1 - t2).powi(2);
                    let entry = semivariance_bins.entry(bin).or_insert((0.0, 0));
                    entry.0 += gamma;
                    entry.1 += 1;
                }
            }
        }

        // Calculate the average semivariance for each bin
        let mut binned_distances = vec![];
        let mut binned_semivariances = vec![];

        // Sort by distance
        let mut sorted_bins: Vec<_> = semivariance_bins.into_iter().collect();
        sorted_bins.sort_by_key(|&(bin, _)| bin);

        //println!("\nLag_Distance(m), Semivariance");
        for (bin, (sum, count)) in sorted_bins {
            if count > 0 {
                let center = (bin as f64 + 0.5) * bin_size;
                let gamma = sum / count as f64;

                binned_distances.push(center);
                binned_semivariances.push(gamma);
            }
        }

        EmpiricalVariogram { distances: binned_distances, semivariances: binned_semivariances }
    }

    /// Returns the maximum distance used in the variogram
    pub fn get_max_distance(&self) -> f64 {
        self.distances.iter().cloned().fold(f64::NEG_INFINITY, f64::max)
    }

    /// Returns the maximum semivariance value
    pub fn get_max_semivariance(&self) -> f64 {
        self.semivariances.iter().cloned().fold(f64::NEG_INFINITY, f64::max)
    }
}

/// The available types of variogram models.
pub enum VariogramModel {
    Spherical,
    Exponential,
    Gaussian
}

/// A structure representing a theoretical variogram model.
pub struct Variogram {
    /// The nugget effect (semivariance at zero distance).
    pub nugget: f64,
    /// The sill (maximum semivariance reached).
    pub sill: f64,
    /// The range (distance at which the sill is approached).
    pub range: f64,
    /// The type of model used.
    pub model: VariogramModel,
}

impl Variogram {
    /// Creates a new `Variogram` with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `nugget` - The nugget effect.
    /// * `sill` - The sill value.
    /// * `range` - The range value.
    /// * `model` - The variogram model to use.
    pub fn new(nugget: f64, sill: f64, range: f64, model: VariogramModel) -> Self {
        Variogram { nugget, sill, range, model }
    }

    /// Evaluates the semivariance at a given distance `h`.
    ///
    /// # Arguments
    ///
    /// * `h` - The distance at which to evaluate the semivariance.
    ///
    /// # Returns
    ///
    /// The computed semivariance based on the current model and parameters.
    pub fn semivariance(&self, h: f64) -> f64 {
        match self.model {
            VariogramModel::Spherical => {
                if h <= self.range {
                    let hr = h / self.range;
                    self.nugget + (self.sill - self.nugget) * (1.5 * hr - 0.5 * hr.powi(3))
                } else {
                    self.sill
                }
            }

            VariogramModel::Exponential => {
                self.nugget + (self.sill - self.nugget) * (1.0 - (-h / self.range).exp())
            }

            VariogramModel::Gaussian => {
                self.nugget + (self.sill - self.nugget) * (1.0 - (- (h / self.range).powi(2)).exp())
            }
        }
    }

    /// Computes the cost (sum of squared differences) between the
    /// model semivariance and the empirical semivariances.
    ///
    /// # Arguments
    ///
    /// * `x_data` - Distances from the empirical variogram.
    /// * `y_data` - Semivariances from the empirical variogram.
    ///
    /// # Returns
    ///
    /// The total cost as a `f64` value.
    fn compute_cost(&self, x_data: &[f64], y_data: &[f64]) -> f64 {
        // Calculate the sum of squared differences
        x_data.iter().zip(y_data.iter())
            .map(|(&h, &y)| {
                let model = self.semivariance(h);
                (y - model).powi(2)
            })
            .sum()
    }

    /// Computes the gradients of the cost function with respect to the
    /// variogram parameters: nugget, sill, and range.
    ///
    /// # Arguments
    ///
    /// * `x_data` - Distances from the empirical variogram.
    /// * `y_data` - Semivariances from the empirical variogram.
    ///
    /// # Returns
    ///
    /// A tuple of gradients `(d_nugget, d_sill, d_range)`.
    fn compute_gradient(&mut self, x_data: &[f64], y_data: &[f64]) -> (f64, f64, f64) {
        let mut grad_nugget = 0.0;
        let mut grad_sill = 0.0;
        let mut grad_range = 0.0;
    
        for (&h, &y) in x_data.iter().zip(y_data.iter()) {
            let model_value = self.semivariance(h);
            let diff = model_value - y;

            match self.model {
                VariogramModel::Spherical => {
                    if h <= self.range {
                        let hr = h / self.range;
                        let common = 1.5 * hr - 0.5 * hr.powi(3);
    
                        // Derivatives
                        grad_nugget += diff * 1.0;
                        grad_sill += diff * common;
    
                        // d(common)/d(range)
                        let d_common_d_range = (-1.5 * h / self.range.powi(2)) + (1.5 * h.powi(3) / self.range.powi(4));
                        grad_range += diff * (self.sill - self.nugget) * d_common_d_range;
                    } else {
                        // For h > range, γ(h) = sill → derivative wrt sill is 1, others are 0
                        grad_nugget += 0.0;
                        grad_sill += diff * 1.0;
                        grad_range += 0.0;
                    }
                }
    
                VariogramModel::Exponential => {
                    let exp_term = (-h / self.range).exp();
                    let factor = self.sill - self.nugget;
    
                    grad_nugget += diff * 1.0;
                    grad_sill += diff * (1.0 - exp_term);
                    grad_range += diff * factor * (h / self.range.powi(2)) * exp_term;
                }
    
                VariogramModel::Gaussian => {
                    let exp_term = (- (h / self.range).powi(2)).exp();
    
                    grad_nugget += diff;
                    grad_sill += diff * (1.0 - exp_term);
                    grad_range += diff * (self.sill - self.nugget) * (-2.0 * h / self.range.powi(3)) * exp_term;
                }
            }
        }
    
        (grad_nugget, grad_sill, grad_range)
    }

    /// Fits the variogram model to an empirical variogram using gradient descent.
    ///
    /// # Arguments
    ///
    /// * `empirical_variogram` - Reference to an `EmpiricalVariogram` containing the data.
    /// * `learning_rate` - Learning rate for gradient descent.
    /// * `max_iter` - Maximum number of iterations.
    /// * `tol` - Tolerance for convergence based on cost improvement.
    pub fn fit(&mut self, empirical_variogram: &EmpiricalVariogram, learning_rate: f64, max_iter: usize, tol: f64) {
        let x_data = &empirical_variogram.distances;
        let y_data = &empirical_variogram.semivariances;
        let mut prev_cost = f64::INFINITY;
    
        for iter in 0..max_iter {
            // Compute the cost
            let cost = self.compute_cost(x_data, y_data);
            
            // Check for convergence
            if (prev_cost - cost).abs() < tol {
                println!("Converged after {} iterations", iter);
                break;
            }
            
            // Compute the gradients
            let (nugget_grad, sill_grad, range_grad) = self.compute_gradient(x_data, y_data);

            // Update parameters
            self.nugget -= learning_rate * nugget_grad;
            self.sill -= learning_rate * sill_grad;
            self.range -= learning_rate * range_grad;
            
            // Update previous cost
            prev_cost = cost;
            
            // Print progress every 10 iterations
            if iter % 10 == 0 {
                println!("Iteration {}: Cost = {:.6}", iter, cost);
            }
        }
    }
}

/// Plots the empirical and fitted variogram, saving the result to an image file.
///
/// This function creates a 2D chart showing:
/// - The empirical semivariogram points (as blue dots).
/// - The fitted variogram model curve (in red).
///
/// The chart is rendered to a PNG image at the given file path.
///
/// # Arguments
///
/// * `variogram` - A reference to the fitted `Variogram` model to plot.
/// * `empirical_variogram` - A reference to the `EmpiricalVariogram` containing observed data.
/// * `file_path` - A reference to the file path where the plot will be saved.
///
/// # Panics
///
/// This function may panic if:
/// - The drawing area cannot be initialized or written to.
/// - The chart configuration fails.
/// - The plotting operations (drawing series) fail.
///
/// # Example
///
/// ```rust
/// let variogram = Variogram::new(0.1, 1.0, 10.0, VariogramModel::Gaussian);
/// let empirical_variogram = EmpiricalVariogram::from_points(&points);
/// let file_path = PathBuf::from("output.png");
/// plot(&variogram, &empirical_variogram, &file_path);
/// ```
pub fn plot(variogram: &Variogram, empirical_variogram: &EmpiricalVariogram, file_path: &PathBuf) {
    let max_x_data = empirical_variogram.get_max_distance();
    let max_y_data = empirical_variogram.get_max_semivariance();

    let root = BitMapBackend::new(&file_path, (640, 480)).into_drawing_area();
    root.fill(&WHITE).unwrap();

    let mut chart = ChartBuilder::on(&root)
        .caption("Semivariogram", ("Arial", 20))
        .x_label_area_size(50)
        .y_label_area_size(50)
        .build_cartesian_2d(0.0..max_x_data, 0.0..max_y_data)
        .unwrap();

    // Configure the mesh (axes and grid)
    chart.configure_mesh()
        .x_desc("Lag Distance (h)")  // X-axis label
        .y_desc("Semivariance")  // Label font and size
        .draw()
        .unwrap();

    let points: Vec<(f64, f64)> = empirical_variogram.distances.iter()
        .copied()
        .zip(empirical_variogram.semivariances.iter().copied())
        .collect();
    
    // Plot data points (dots)
    chart.draw_series(
        points.iter().map(|&(x, y)| {
            Circle::new((x, y), 5, 
                ShapeStyle {
                    color: BLUE.to_rgba(),
                    filled: true,
                    stroke_width: 2,
                })
        })
    ).unwrap();

    // Plot fitted Gaussian curve
    chart.draw_series(LineSeries::new(
        (0..max_x_data as usize)
            .map(|i| (i as f64, variogram.semivariance(i as f64))),
        &RED,
    )).unwrap();

    root.present().unwrap();
}