use num_complex::{Complex, Complex64};

use crate::filters::iir::lib::kve;

/// Aberth-Ehrlich method for simultaneous root-finding
pub fn aberth<F, Fp>(
    f: &F,
    fp: &Fp,
    x0: &[Complex<f64>],
    tol: f64,
    max_iter: usize,
) -> Result<Vec<Complex<f64>>, String>
where
    F: Fn(&[Complex<f64>]) -> Vec<Complex<f64>>,
    Fp: Fn(&[Complex<f64>]) -> Vec<Complex<f64>>,
{
    let n = x0.len();
    let mut x = x0.to_vec();
    let mut beta = vec![Complex::new(0.0, 0.0); n];

    for _ in 0..max_iter {
        let fx = f(&x);
        let fpx = fp(&x);

        let alpha: Vec<Complex<f64>> = fx
            .iter()
            .zip(fpx.iter())
            .map(|(f, fp)| -f / fp)
            .collect();

        // Calculate beta terms (repulsion)
        for k in 0..n {
            let mut sum = Complex::new(0.0, 0.0);
            for j in 0..n {
                if j != k {
                    sum += Complex::new(1.0, 0.0) / (x[k] - x[j]);
                }
            }
            beta[k] = sum;
        }

        // Update x
        for k in 0..n {
            let denom = Complex::new(1.0, 0.0) + alpha[k] * beta[k];
            x[k] += alpha[k] / denom;
        }

        // Check for NaN / Inf
        if x.iter().any(|z| !z.re.is_finite() || !z.im.is_finite()) {
            return Err("Root-finding failed: non-finite values".to_string());
        }

        // Convergence check
        if alpha.iter().all(|a| a.norm() <= tol) {
            return Ok(x);
        }
    }

    Err("Aberth method failed to converge".to_string())
}
pub fn make_fp(n: usize) -> impl Fn(&[Complex<f64>]) -> Vec<Complex<f64>> {
    move |x: &[Complex<f64>] | {
        x.iter()
            .map(|&z| {
                let inv = Complex64::new(1.0, 0.0) / z;
                let x2 = z * z;
                let term1 = kve(n as f64 - 0.5, inv, 50) / (2.0 * x2);
                let term2 = kve(n as f64 + 0.5, inv, 50) / x2;
                let term3 = kve(n as f64 + 1.5, inv, 50) / (2.0 * x2);
                term1 - term2 + term3
            })
            .collect()
    }
}

pub fn make_f(n: usize) -> impl Fn(&[Complex<f64>]) -> Vec<Complex<f64>> {
    move |x: &[Complex<f64>]| {
        x.iter()
            .map(|&z| {
                let inv = Complex64::new(1.0, 0.0) / z; // real part only for placeholder
                kve(n as f64 + 0.5, inv, 50)
            })
            .collect()
    }
}


