use num_complex::Complex;
use std::f64::EPSILON;

pub fn newton_complex(
    func: &dyn Fn(Complex<f64>) -> Complex<f64>,
    x0: Complex<f64>,
    fprime: Option<&dyn Fn(Complex<f64>) -> Complex<f64>>,
    tol: f64,
    maxiter: usize,
    fprime2: Option<&dyn Fn(Complex<f64>) -> Complex<f64>>,
    x1: Option<Complex<f64>>,
    rtol: f64,
    disp: bool,
) -> Result<Complex<f64>, String>
{
    if tol <= 0.0 {
        return Err(format!("tol must be positive, got {}", tol));
    }
    if maxiter < 1 {
        return Err("maxiter must be >= 1".to_string());
    }

    // let mut funcalls = 0;
    // let mut method = "secant".to_string();

    if let Some(df) = fprime {
        // Newton-Raphson or Halley method
        // method = "newton".to_string();
        let mut p0 = x0;

        for i in 0..maxiter {
            let fval = func(p0);
            // funcalls += 1;

            if fval.norm() == 0.0 {
                return Ok(p0);
            }

            let dval = df(p0);
            // funcalls += 1;

            if dval.norm() == 0.0 {
                let msg = format!(
                    "Derivative was zero. Failed to converge after {} iterations, value is {}.",
                    i + 1,
                    p0
                );
                if disp {
                    return Err(msg);
                } else {
                    return Ok(p0);
                }
            }

            let mut step = fval / dval;

            if let Some(df2) = &fprime2 {
                let d2val = df2(p0);
                // funcalls += 1;
                // method = "halley".to_string();

                let adj = step * d2val / dval / Complex::from(2.0);
                if adj.norm() < 1.0 {
                    step = step / (Complex::from(1.0) - adj);
                }
            }

            let p = p0 - step;

            if (p - p0).norm() <= tol + rtol * p.norm() {
                return Ok(p);
            }

            p0 = p;
        }
    } else {
        // Secant method
        // method = "secant".to_string();
        let mut p0 = x0;
        let mut p1 = if let Some(x1_val) = x1 {
            if x1_val == x0 {
                return Err("x1 and x0 must be different".to_string());
            }
            x1_val
        } else {
            let eps = 1e-4;
            let dx = Complex::new(eps, eps);
            let mut guess = x0 * (Complex::from(1.0) + dx);
            if guess == x0 {
                guess += if x0.re >= 0.0 { dx } else { -dx };
            }
            guess
        };

        let mut q0 = func(p0);
        let mut q1 = func(p1);
        // funcalls += 2;

        if q1.norm() < q0.norm() {
            std::mem::swap(&mut p0, &mut p1);
            std::mem::swap(&mut q0, &mut q1);
        }

        for i in 0..maxiter {
            if (q1 - q0).norm() < EPSILON {
                let p = (p1 + p0) / Complex::from(2.0);

                if disp {
                    return Err(format!(
                        "Tolerance reached. Failed to converge after {} iterations, value is {}.",
                        i + 1,
                        p
                    ));
                } else {
                    return Ok(p);
                }
            }

            let p = if q1.norm() > q0.norm() {
                (-q0 / q1 * p1 + p0) / (Complex::from(1.0) - q0 / q1)
            } else {
                (-q1 / q0 * p0 + p1) / (Complex::from(1.0) - q1 / q0)
            };

            if (p - p1).norm() <= tol + rtol * p.norm() {
                return Ok(p);
            }

            p0 = p1;
            q0 = q1;
            p1 = p;
            q1 = func(p1);
            // funcalls += 1;
        }
    }

    if disp {
        Err(format!(
            "Failed to converge after {} iterations.",
            maxiter
        ))
    } else {
        Ok(x0)
    }
}

pub fn newton(
    func: &dyn Fn(f64) -> f64,
    x0: f64,
    fprime: Option<&dyn Fn(f64) -> f64>,
    tol: f64,
    maxiter: usize,
    fprime2: Option<&dyn Fn(f64) -> f64>,
    x1: Option<f64>,
    rtol: f64,
    disp: bool,
) -> Result<f64, String>
{
    if tol <= 0.0 {
        return Err(format!("tol too small ({:.1e} <= 0)", tol));
    }
    if maxiter == 0 {
        return Err("maxiter must be greater than 0".to_string());
    }

    // let mut funcalls = 0;
    // let mut method = "secant".to_string();

    // Newton-Raphson or Halley method
    if let Some(df) = fprime {
        // method = "newton".to_string();
        let mut p0 = x0;

        for i in 0..maxiter {
            let fval = func(p0);
            // funcalls += 1;

            if fval.abs() == 0.0 {
                return Ok(p0);
            }

            let dval = df(p0);
            // funcalls += 1;

            if dval.abs() < EPSILON {
                let msg = format!(
                    "Derivative was zero. Failed to converge after {} iterations, value is {}.",
                    i + 1,
                    p0
                );
                if disp {
                    return Err(msg);
                } else {
                    return Ok(p0);
                }
            }

            let mut step = fval / dval;

            if let Some(df2) = &fprime2 {
                let d2val = df2(p0);
                // funcalls += 1;
                // method = "halley".to_string();

                let adj = 0.5 * step * d2val / dval;
                if adj.abs() < 1.0 {
                    step = step / (1.0 - adj);
                }
            }

            let p = p0 - step;

            if (p - p0).abs() <= tol + rtol * p.abs() {
                return Ok(p);
            }

            p0 = p;
        }
    } else {
        // Secant method
        // method = "secant".to_string();
        let mut p0 = x0;
        let mut p1 = if let Some(x1_val) = x1 {
            if x1_val == x0 {
                return Err("x1 and x0 must be different".to_string());
            }
            x1_val
        } else {
            let eps = 1e-4;
            let mut guess = x0 * (1.0 + eps);
            if guess == x0 {
                guess += if guess >= 0.0 { eps } else { -eps };
            }
            guess
        };

        let mut q0 = func(p0);
        let mut q1 = func(p1);
        // funcalls += 2;

        if q1.abs() < q0.abs() {
            std::mem::swap(&mut p0, &mut p1);
            std::mem::swap(&mut q0, &mut q1);
        }

        for i in 0..maxiter {
            if (q1 - q0).abs() < EPSILON {
                let p = 0.5 * (p1 + p0);
                if disp {
                    return Err(format!(
                        "Tolerance reached. Failed to converge after {} iterations, value is {}.",
                        i + 1,
                        p
                    ));
                } else {
                    return Ok(p);
                }
            }

            let p = if q1.abs() > q0.abs() {
                (-q0 / q1 * p1 + p0) / (1.0 - q0 / q1)
            } else {
                (-q1 / q0 * p0 + p1) / (1.0 - q1 / q0)
            };

            if (p - p1).abs() <= tol + rtol * p.abs() {
                return Ok(p);
            }

            p0 = p1;
            q0 = q1;
            p1 = p;
            q1 = func(p1);
            // funcalls += 1;
        }
    }

    // If not returned, failed to converge
    if disp {
        Err(format!(
            "Failed to converge after {} iterations.",
            maxiter
        ))
    } else {
        Ok(x0)
    }
}

