greeners-ols 2.0.1

use crate::ols::OLS;
use greeners_core::{CovarianceType, DataFrame, Formula, GreenersError};
use ndarray::{Array1, Array2, Axis};
use rand::distributions::Distribution;
use rand_distr::Uniform;
use std::fmt;

#[derive(Debug)]
pub struct QuantileResult {
    pub tau: f64,
    pub params: Array1<f64>,
    pub std_errors: Array1<f64>,
    pub t_values: Array1<f64>,
    pub p_values: Array1<f64>,
    pub r_squared: f64, // Pseudo-R2 (Koenker & Machado)
    pub iterations: usize,
    pub variable_names: Option<Vec<String>>,
}

impl fmt::Display for QuantileResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(
            f,
            "\n{:=^78}",
            format!(" Quantile Regression (tau={:.2}) ", self.tau)
        )?;
        writeln!(f, "{:<20} {:>15.4} (Pseudo)", "R-squared:", self.r_squared)?;
        writeln!(f, "{:<20} {:>15}", "Iterations:", self.iterations)?;

        writeln!(f, "\n{:-^78}", "")?;
        writeln!(
            f,
            "{:<10} | {:>10} | {:>10} | {:>8} | {:>8}",
            "Variable", "Coef", "Std Err", "t", "P>|t|"
        )?;
        writeln!(f, "{:-^78}", "")?;

        for i in 0..self.params.len() {
            let var_name = if let Some(ref names) = self.variable_names {
                if i < names.len() {
                    names[i].clone()
                } else {
                    format!("x{}", i)
                }
            } else {
                format!("x{}", i)
            };

            writeln!(
                f,
                "{:<10} | {:>10.4} | {:>10.4} | {:>8.3} | {:>8.3}",
                var_name, self.params[i], self.std_errors[i], self.t_values[i], self.p_values[i]
            )?;
        }
        writeln!(f, "{:=^78}", "")
    }
}

pub struct QuantileReg;

impl QuantileReg {
    /// Estimates Quantile Regression using a formula and DataFrame.
    pub fn from_formula(
        formula: &Formula,
        data: &DataFrame,
        tau: f64,
        n_boot: usize,
    ) -> Result<QuantileResult, GreenersError> {
        let (y, x) = data.to_design_matrix(formula)?;
        let var_names = data.formula_var_names(formula)?;
        Self::fit_with_names(&y, &x, tau, n_boot, Some(var_names))
    }

    /// Estimate Quantyl Regression via IRLS.
    ///
    /// # Arguments
    /// * `tau` - The desired quantil (e.g. 0.5 for median, 0.9 for upper decile).
    /// * `n_boot` - Número de repetições de Bootstrap para erro padrão (rec: 200+).
    pub fn fit(
        y: &Array1<f64>,
        x: &Array2<f64>,
        tau: f64,
        n_boot: usize,
    ) -> Result<QuantileResult, GreenersError> {
        Self::fit_with_names(y, x, tau, n_boot, None)
    }

    pub fn fit_with_names(
        y: &Array1<f64>,
        x: &Array2<f64>,
        tau: f64,
        n_boot: usize,
        variable_names: Option<Vec<String>>,
    ) -> Result<QuantileResult, GreenersError> {
        if tau <= 0.0 || tau >= 1.0 {
            return Err(GreenersError::OptimizationFailed); // "Tau must be in (0, 1)"
        }

        // Check for NaN/Inf in input data
        if y.iter().any(|v| !v.is_finite()) || x.iter().any(|v| !v.is_finite()) {
            return Err(GreenersError::InvalidOperation(
                "Input data contains NaN or Inf values".into(),
            ));
        }

        let k = x.ncols();

        //1. Time Estimation (IRLS)
        let (params, iter) = Self::irls_solver(y, x, tau)?;

        //2. Default Errors via Bootstrap, if requested
        let std_errors = if n_boot == 0 {
            Array1::from_elem(k, f64::NAN)
        } else {
            Self::bootstrap_se(y, x, tau, n_boot, &params)?
        };

        //3. Final Statistics
        let t_values = &params / &std_errors;
        let normal = statrs::distribution::Normal::standard();
        use statrs::distribution::ContinuousCDF;
        let p_values = t_values.mapv(|t| 2.0 * (1.0 - normal.cdf(t.abs())));

        // 4. Pseudo-R2 (Goodness of Fit)
        // R1 = 1 - (Loss(Model) / Loss(Null))
        //Null Model: Intercept only (quantile of raw variable Y)
        let loss_model = Self::check_loss(y, &x.dot(&params), tau);

        //Null model (constant only)
        let n = y.len();
        let ones = Array2::<f64>::ones((n, 1));
        let (params_null, _) = Self::irls_solver(y, &ones, tau)?;
        let loss_null = Self::check_loss(y, &ones.dot(&params_null), tau);

        let pseudo_r2 = 1.0 - (loss_model / loss_null);

        Ok(QuantileResult {
            tau,
            params,
            std_errors,
            t_values,
            p_values,
            r_squared: pseudo_r2,
            iterations: iter,
            variable_names,
        })
    }

    /// Iteratively Reweighted Least Squares solver
    fn irls_solver(
        y: &Array1<f64>,
        x: &Array2<f64>,
        tau: f64,
    ) -> Result<(Array1<f64>, usize), GreenersError> {
        let n = y.len();
        let max_iter = 1000;
        let tol = 1e-6;
        let epsilon = 1e-6; //To avoid division by zero

        // Chute inicial: OLS
        let ols = OLS::fit(y, x, CovarianceType::NonRobust)?;
        let mut beta = ols.params;
        let mut diff = 1.0;
        let mut iter = 0;

        while diff > tol && iter < max_iter {
            let old_beta = beta.clone();

            //Calculate Waste
            let pred = x.dot(&beta);
            let residuals = y - &pred;

            //Calculate IRLS Weights for Quantil
            // w_i = tau / |u| se u > 0
            // w_i = (1-tau) / |u| se u < 0
            // Aproximação suave: w_i = 1 / max(epsilon, |u|) * (tau if u>0 else 1-tau)

            let mut weights = Array1::<f64>::zeros(n);
            for i in 0..n {
                let u = residuals[i];
                let abs_u = u.abs().max(epsilon); // Evitar 0

                let w = if u >= 0.0 {
                    tau / abs_u
                } else {
                    (1.0 - tau) / abs_u
                };
                weights[i] = w;
            }

            // Rodar WLS (Weighted Least Squares)
            // Beta = (X' W X)^-1 X' W y
            //Algebraic trick: Multiply X and y by sqrt(w) and run OLS
            let sqrt_w = weights.mapv(f64::sqrt);
            let y_w = y * &sqrt_w;
            let mut x_w = x.clone();
            for (i, mut row) in x_w.axis_iter_mut(Axis(0)).enumerate() {
                row *= sqrt_w[i];
            }

            let wls_res = OLS::fit(&y_w, &x_w, CovarianceType::NonRobust);

            //If WLS fails (singular matrix due to extreme weights), regressor error or stop
            if let Ok(res) = wls_res {
                beta = res.params;
            } else {
                return Err(GreenersError::OptimizationFailed);
            }

            //Convergence based on coefficients change
            diff = (&beta - &old_beta).mapv(|v| v.abs()).sum();
            iter += 1;
        }

        Ok((beta, iter))
    }

    //Bootstrap / Pairs for Default Error
    fn bootstrap_se(
        y: &Array1<f64>,
        x: &Array2<f64>,
        tau: f64,
        n_boot: usize,
        point_estimate: &Array1<f64>,
    ) -> Result<Array1<f64>, GreenersError> {
        let n = y.len();
        let k = x.ncols();
        let mut boot_betas = Array2::<f64>::zeros((n_boot, k));
        for mut row in boot_betas.rows_mut() {
            row.assign(point_estimate);
        }

        let mut rng = rand::thread_rng();

        for b in 0..n_boot {
            //Resampling with replacement
            let indices = Array1::from_iter((0..n).map(|_| Uniform::new(0, n).sample(&mut rng)));

            let mut y_boot_vec = Vec::with_capacity(n);
            let mut x_boot_vec = Vec::with_capacity(n * k);

            for &idx in &indices {
                y_boot_vec.push(y[idx]);
                for val in x.row(idx) {
                    x_boot_vec.push(*val);
                }
            }

            let y_boot = Array1::from(y_boot_vec);
            let x_boot = Array2::from_shape_vec((n, k), x_boot_vec)
                .map_err(|_| GreenersError::ShapeMismatch("Bootstrap shape error".into()))?;

            // Estimar no bootstrap
            //If it fails (rare), we just ignore this iteration by copying the previous (simplification)
            if let Ok((beta_b, _)) = Self::irls_solver(&y_boot, &x_boot, tau) {
                boot_betas.row_mut(b).assign(&beta_b);
            }
        }

        //Calculate Default deviation of bootstrap estimates
        let mut std_errs = Array1::<f64>::zeros(k);
        for j in 0..k {
            let col = boot_betas.column(j);
            let mean = col.mean().unwrap_or(0.0);
            let var = col.mapv(|v| (v - mean).powi(2)).sum() / ((n_boot - 1) as f64);
            std_errs[j] = var.sqrt();
        }

        Ok(std_errs)
    }

    /// Loss Check Function
    fn check_loss(y: &Array1<f64>, y_pred: &Array1<f64>, tau: f64) -> f64 {
        let res = y - y_pred;
        res.mapv(|u| if u >= 0.0 { tau * u } else { (tau - 1.0) * u })
            .sum()
    }
}