computorv1 1.0.11

use ft_lib::ft_sqrt::ft_sqrt;
use crate::aux::Polynomial;
use crate::solver::degree_1::degree_1;
use crate::solver::Roots;

// Degree 2 Polynomial (Quadratic):
// a2 * x^2 + a1 * x + a0 = 0, where a2 != 0
// Discriminant: D = a1^2 - 4 * a2 * a0
// Solution:
// x = (-a1 ± sqrt(D)) / (2 * a2)
// - D > 0: two real roots
// - D = 0: one real root (double root)
// - D < 0: two complex roots

pub fn degree_2(poly: Polynomial) -> Result<Roots, String> {
    let a: f64 = poly.terms.get(&2).copied().unwrap_or(0.0);
    let b: f64 = poly.terms.get(&1).copied().unwrap_or(0.0);
    let c: f64 = poly.terms.get(&0).copied().unwrap_or(0.0);

    if a == 0.0 {
        return degree_1(poly);
    }

    let discriminant: f64 = b * b - 4.0 * a * c;
    let two_a: f64 = 2.0 * a;
    if discriminant > 0.0 {
        println!("Discriminant: {} is strictly positive", discriminant);
        let x1: f64 = (-b + ft_sqrt(discriminant)) / two_a;
        let x2: f64 = (-b - ft_sqrt(discriminant)) / two_a;
        Ok(Roots::with_real(vec![x1, x2]))
    } else if discriminant == 0.0 {
        println!("Discriminant is zero, the solution is:");
        let x: f64 = -b / two_a;
        Ok(Roots::with_real(vec![x]))
    } else {
        let sqrt_disc: f64 = ft_sqrt(-discriminant); // note: take sqrt of (-discriminant)
        let real_part: f64 = -b / two_a;
        let imaginary_part: f64 = sqrt_disc / two_a;
        Ok(Roots::with_complex(vec![
            (real_part, imaginary_part),
            (real_part, -imaginary_part),
        ]))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::parse_equation;
    mod degree_2 {
        use super::*;

        #[test]
        fn test_degree_2_no_linear_term() {
            let equation: &str = "1X^2+0X-4 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![2.0, -2.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(
                result.real,
                expected_real
            );
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_degree_2_no_constant_term() {
            let equation: &str = "2X^2-6X+0 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![3.0, 0.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(
                result.real,
                expected_real
            );
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_degree_1() {
            let equation: &str = "0X^2+2X+1 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![-0.5];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }
        #[test]
        fn test_degree_0_invalid() {
            let equation: &str = "0X^2 + 0X^1 + 1X^0 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Result<Roots, String> = degree_2(poly);
            assert!(result.is_err(), "Expected an error got {:?}", result);
        }
        #[test]
        fn test_degree_0_valid() {
            let equation: &str = "0X^2 + 0X^1 + 0X^0 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            assert_eq!(result.real, vec![]);
        }
        #[test]
        fn test_one_real_root() {
            let equation: &str = "X^2 - 2*X^1 + 1 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![1.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_fractional_coefficients() {
            let equation: &str = "0.5 * X^2 - 1.5 * X^1 + 1 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![2.0, 1.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(
                result.real,
                expected_real
            );
            assert_eq!(result.complex, expected_complex);

            let equation: &str = "2*X^2 + 5*X - 3 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![0.5, -3.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(
                result.real,
                expected_real
            );
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_negative_leading_coefficient() {
            let equation: &str = "-1 * X^2 + 3*X^1 - 2 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![1.0, 2.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_reduced_to_degree_1() {
            let equation: &str = "0 * X^2 + 2*X^1 - 4 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![2.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_zero_polynomial() {
            let equation: &str = "0 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_unbalanced_equation() {
            let equation: &str = "X^2 + 2*X^1 + 1 = X^2";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![-0.5];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_two_real_roots() {
            let equation: &str = "X^2 - 5*X + 6 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![3.0, 2.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(
                result.real,
                expected_real
            );
            assert_eq!(result.complex, expected_complex);

            let equation: &str = "X^2 - 4*X + 4 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![2.0];
            let expected_complex: Vec<(f64, f64)> = vec![];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_negative_coefficients() {
            let equation: &str = "-X^2 + X + 2 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![-1.0, 2.0];
            let expected_complex: Vec<(f64, f64)> = vec![];
            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_missing_linear_term() {
            let equation: &str = "X^2 - 9 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![3.0, -3.0];
            assert_eq!(result.real, expected_real);
        }

        #[test]
        fn test_missing_constant_term() {
            let equation: &str = "X^2 + 3*X = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![0.0, -3.0];
            assert_eq!(result.real, expected_real);
        }

        #[test]
        fn test_all_zero_coefficients() {
            let equation: &str = "0*X^2 + 0*X + 0 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![];
            let expected_complex: Vec<(f64, f64)> = vec![];
            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }
        #[test]
        fn test_linear_case_reduced_from_quadratic() {
            let equation: &str = "X^2 + 2*X + 1 = X^2 + 3";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![1.0];
            assert_eq!(result.real, expected_real);
        }

        #[test]
        fn test_valid() {
            let equation: &str = "X^2 + 6X + 5 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![-1.0, -5.0];
            assert_eq!(result.real, expected_real);

            let equation: &str = "X^2 - 2X+ 1 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![1.0];
            assert_eq!(result.real, expected_real);
        }
    }

    mod degree_2_complex {
        use super::*;

        #[test]
        fn test_no_real_roots() {
            let equation: &str = "X^2 + 1 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();
            let expected_real: Vec<f64> = vec![];
            let expected_complex: Vec<(f64, f64)> = vec![(0.0, 1.0), (0.0, -1.0)];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }
        #[test]
        fn test_complex_roots() {
            let equation: &str = "X^2 + 4 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![];
            let expected_complex: Vec<(f64, f64)> = vec![(0.0, 2.0), (0.0, -2.0)];

            assert_eq!(result.real, expected_real);
            assert_eq!(result.complex, expected_complex);
        }

        #[test]
        fn test_more_complex(){
            let equation: &str = "X^2 - 3X + 10  = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![];
            let expected_complex: Vec<(f64, f64)> = vec![(1.5, 2.783882181415011), (1.5, -2.783882181415011)];
            
            //assert_eq!(result.complex, expected_complex);
            assert_eq!(
                result.complex,
                expected_complex
            );
            assert_eq!(result.real, expected_real);
        }
        #[test]
        fn test_no_real_solutions() {
            let equation: &str = "X^2 + X + 1 = 0";
            let poly: Polynomial = parse_equation(equation).expect("Failed to parse equation");
            let result: Roots = degree_2(poly).unwrap();

            let expected_real: Vec<f64> = vec![];
            let expected_complex: Vec<(f64, f64)> = vec![(-0.5, 0.8660254037844386), (-0.5, -0.8660254037844386)];

            assert_eq!(result.real, expected_real);
            assert_eq!(
                result.complex,
                expected_complex
            );
        }

    }
}