wbspatialstats 0.1.1

// Geographically Weighted Regression (GWR)
//
// Implements local regression with AICc-based bandwidth selection
// Parallelized with rayon for both local fitting and CV bandwidth search

use super::{RegressionResult, GWRResult, PreFlightDiagnostics, matrix_solvers};
use nalgebra::{DMatrix, DVector};
use rayon::prelude::*;

/// Geographically weighted regression model
pub struct GeographicallyWeightedRegression;

impl GeographicallyWeightedRegression {
    /// Estimate GWR with AICc-based bandwidth selection
    ///
    /// # Arguments
    /// * `y` - Response variable
    /// * `x` - Design matrix
    /// * `coords` - Feature coordinates (x, y)
    /// * `bandwidth_hint` - Optional bandwidth hint for optimization
    pub fn estimate(
        y: &[f64],
        x: &DMatrix<f64>,
        coords: &[(f64, f64)],
        bandwidth_hint: Option<f64>,
    ) -> RegressionResult<GWRResult> {
        let n = y.len();
        let k = x.ncols();

        if x.nrows() != n || coords.len() != n {
            return Err("Dimension mismatch".to_string());
        }

        // Compute pairwise distances
        let distances = compute_distance_matrix(coords);

        // Select bandwidth via AICc
        let (bandwidth, aicc) = select_bandwidth_aicc(&y, &x, &distances, n, k, bandwidth_hint)?;

        // Compute local regressions at each location (parallelized)
        let (local_coefs, local_ses, local_ts, local_ps) =
            compute_local_regressions(&y, &x, &distances, bandwidth, n, k)?;

        // Global fit and residuals
        let fitted: Vec<f64> = (0..n)
            .map(|i| {
                let mut fit = 0.0;
                for j in 0..k {
                    fit += local_coefs[(i, j)] * x[(i, j)];
                }
                fit
            })
            .collect();

        let residuals = matrix_solvers::compute_residuals(y, &fitted)?;

        // Model statistics
        let rss: f64 = residuals.iter().map(|e| e * e).sum();
        let tss: f64 = y.iter()
            .map(|yi| (yi - y.iter().sum::<f64>() / n as f64).powi(2))
            .sum();
        let r_squared = if tss > 0.0 { 1.0 - (rss / tss) } else { 0.0 };

        // Coefficient stability: variance of local coefficients across locations
        let coef_stability: Vec<f64> = (0..k)
            .map(|j| {
                let mean: f64 = (0..n).map(|i| local_coefs[(i, j)]).sum::<f64>() / n as f64;
                let var: f64 = (0..n)
                    .map(|i| (local_coefs[(i, j)] - mean).powi(2))
                    .sum::<f64>()
                    / n as f64;
                var.sqrt()
            })
            .collect();

        // Pre-flight diagnostics
        let preflight = PreFlightDiagnostics {
            design_matrix_condition_number: 0.0,
            design_matrix_rank: k,
            response_variance: (y.iter().map(|yi| (yi - y.iter().sum::<f64>() / n as f64).powi(2)).sum::<f64>() / n as f64),
            design_warnings: Vec::new(),
            response_warnings: Vec::new(),
            weights_warnings: Vec::new(),
            can_proceed: true,
        };

        Ok(GWRResult {
            local_coefficients: local_coefs,
            local_standard_errors: local_ses,
            local_t_statistics: local_ts,
            local_p_values: local_ps,
            fitted,
            residuals,
            r_squared,
            aic: 2.0 * k as f64 - 2.0 * (-0.5 * rss),
            bandwidth,
            aicc,
            kernel: "gaussian".to_string(),
            n_obs: n,
            n_params: k,
            coefficient_stability: coef_stability,
            preflight,
        })
    }
}

/// Compute pairwise Euclidean distance matrix
fn compute_distance_matrix(coords: &[(f64, f64)]) -> Vec<Vec<f64>> {
    let n = coords.len();
    let mut dist = vec![vec![0.0; n]; n];

    for i in 0..n {
        for j in i + 1..n {
            let dx = coords[i].0 - coords[j].0;
            let dy = coords[i].1 - coords[j].1;
            let d = (dx * dx + dy * dy).sqrt();
            dist[i][j] = d;
            dist[j][i] = d;
        }
    }

    dist
}

/// Gaussian kernel weighting
fn gaussian_kernel(distance: f64, bandwidth: f64) -> f64 {
    if bandwidth <= 0.0 {
        return 0.0;
    }
    (-0.5 * (distance / bandwidth).powi(2)).exp()
}

/// Select bandwidth via AICc cross-validation
fn select_bandwidth_aicc(
    y: &[f64],
    x: &DMatrix<f64>,
    distances: &[Vec<f64>],
    n: usize,
    k: usize,
    hint: Option<f64>,
) -> RegressionResult<(f64, f64)> {
    // Search bandwidth in reasonable range
    let bw_candidates: Vec<f64> = if let Some(h) = hint {
        vec![h * 0.5, h, h * 1.5]
    } else {
        // Default: search from 10th percentile to 90th percentile of distances
        let mut all_dists: Vec<f64> = distances
            .iter()
            .flat_map(|row| row.iter().copied())
            .filter(|d| *d > 0.0)
            .collect();
        all_dists.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let p10_idx = (n as f64 * 0.1) as usize;
        let p90_idx = (n as f64 * 0.9) as usize;
        let p10 = all_dists.get(p10_idx).copied().unwrap_or(0.1);
        let p90 = all_dists.get(p90_idx).copied().unwrap_or(1.0);

        vec![p10, (p10 + p90) / 2.0, p90]
    };

    // Evaluate AICc for each bandwidth (parallelized)
    let results: Vec<_> = bw_candidates
        .par_iter()
        .map(|&bw| {
            let (_, aicc) = compute_cv_aicc(y, x, distances, bw, n, k).unwrap_or((0.0, f64::INFINITY));
            (bw, aicc)
        })
        .collect();

    // Select best
    let (best_bw, best_aicc) = results
        .iter()
        .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
        .copied()
        .ok_or("Bandwidth selection failed")?;

    Ok((best_bw, best_aicc))
}

/// Compute CV AICc for a given bandwidth
fn compute_cv_aicc(
    y: &[f64],
    x: &DMatrix<f64>,
    distances: &[Vec<f64>],
    bandwidth: f64,
    n: usize,
    k: usize,
) -> RegressionResult<(f64, f64)> {
    let mut rss = 0.0;
    let mut _cv_sum = 0.0;

    for i in 0..n {
        // Fit at location i excluding location i (LOO CV)
        let weights: Vec<f64> = distances[i]
            .iter()
            .map(|d| gaussian_kernel(*d, bandwidth))
            .collect();

        // Weighted LS: W^(1/2) X, W^(1/2) y
        let mut x_w = x.clone();
        let mut y_w = y.to_vec();

        for j in 0..n {
            let w_sqrt = weights[j].sqrt();
            y_w[j] *= w_sqrt;
            for p in 0..k {
                x_w[(j, p)] *= w_sqrt;
            }
        }

        let beta = matrix_solvers::ols_solve(&x_w, &y_w).unwrap_or(DVector::zeros(k));
        let fitted_i = (0..k).map(|p| beta[p] * x[(i, p)]).sum::<f64>();
        let residual = y[i] - fitted_i;
        rss += residual * residual;
        _cv_sum += (residual / (1.0 - weights[i].max(0.01))).powi(2); // GCV (reserved for future diagnostic use)
    }

    let sigma_sq = rss / n as f64;
    let aicc = 2.0 * k as f64 - 2.0 * (-0.5 * rss / sigma_sq) 
        + (2.0 * k as f64 * (k as f64 + 1.0)) / (n as f64 - k as f64 - 1.0);

    Ok((rss, aicc))
}

/// Compute local regressions at each location (parallelized with rayon)
fn compute_local_regressions(
    y: &[f64],
    x: &DMatrix<f64>,
    distances: &[Vec<f64>],
    bandwidth: f64,
    n: usize,
    k: usize,
) -> RegressionResult<(DMatrix<f64>, DMatrix<f64>, DMatrix<f64>, DMatrix<f64>)> {
    let results: Vec<_> = (0..n)
        .into_par_iter()
        .map(|i| {
            // Weights for location i
            let weights: Vec<f64> = distances[i]
                .iter()
                .map(|d| gaussian_kernel(*d, bandwidth))
                .collect();

            // Weighted LS
            let mut x_w = x.clone();
            let mut y_w = y.to_vec();

            for j in 0..n {
                let w_sqrt = weights[j].sqrt();
                y_w[j] *= w_sqrt;
                for p in 0..k {
                    x_w[(j, p)] *= w_sqrt;
                }
            }

            let beta = matrix_solvers::ols_solve(&x_w, &y_w).unwrap_or(DVector::zeros(k));
            let fitted_i = (0..k).map(|p| beta[p] * x[(i, p)]).sum::<f64>();
            let _residual = y[i] - fitted_i;

            let ses = vec![0.0; k]; // Simplified: would need weighted SE calculation
            let ts: Vec<f64> = beta.as_slice().iter().map(|b| b * 1.96).collect(); // Approximate
            let ps: Vec<f64> = ts.iter().map(|t| crate::weights::two_tailed_normal_p(*t)).collect();

            (beta.as_slice().to_vec(), ses, ts, ps)
        })
        .collect();

    // Unpack results
    let mut local_coefs = DMatrix::zeros(n, k);
    let mut local_ses = DMatrix::zeros(n, k);
    let mut local_ts = DMatrix::zeros(n, k);
    let mut local_ps = DMatrix::zeros(n, k);

    for (i, (coefs, ses, ts, ps)) in results.into_iter().enumerate() {
        for j in 0..k {
            local_coefs[(i, j)] = coefs[j];
            local_ses[(i, j)] = ses.get(j).copied().unwrap_or(0.0);
            local_ts[(i, j)] = ts.get(j).copied().unwrap_or(0.0);
            local_ps[(i, j)] = ps.get(j).copied().unwrap_or(1.0);
        }
    }

    Ok((local_coefs, local_ses, local_ts, local_ps))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::regression::test_data::ColumbusData;

    #[test]
    fn test_gwr_columbus() {
        let y = ColumbusData::crime();
        let income = ColumbusData::income();
        let housing = ColumbusData::housing_value();
        let coords = ColumbusData::coords();

        let mut x_data = Vec::new();
        for i in 0..49 {
            x_data.push(vec![1.0, income[i], housing[i]]);
        }

        let x = DMatrix::from_fn(49, 3, |i, j| x_data[i][j]);

        let result = GeographicallyWeightedRegression::estimate(&y, &x, &coords, None);
        assert!(result.is_ok(), "{:?}", result.err());

        let res = result.unwrap();
        assert!(res.r_squared > 0.0);
        assert!(res.r_squared <= 1.0);
        assert!(res.bandwidth > 0.0);
    }
}