use num::{
    bigint::{BigInt, ToBigInt},
    Signed,
};

/// A simple conversion function from anything that can be converted
/// to a big integer to a big rational number.
fn conversion(n: impl ToBigInt) -> BigInt {
    n.to_bigint().unwrap()
}

fn print_poly(poly: &[BigInt]) {
    struct Monomial(usize);

    impl std::fmt::Display for Monomial {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            match self.0 {
                0 => Ok(()),
                1 => {
                    write!(f, "x")
                }
                _ => {
                    write!(f, "x^{}", self.0)
                }
            }
        }
    }

    struct Coefficient(BigInt);

    impl std::fmt::Display for Coefficient {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            match &self.0 {
                n if n == &conversion(0) || n == &conversion(1) || n == &conversion(-1) => Ok(()),
                _ => {
                    write!(f, "{}", self.0.abs())
                }
            }
        }
    }

    for (n, coeff) in poly.iter().enumerate().rev() {
        if coeff == &conversion(0) {
            continue;
        }

        print!(
            "{}{}{}",
            if n + 1 == poly.len() {
                ""
            } else if coeff >= &conversion(0) {
                " + "
            } else {
                " - "
            },
            Coefficient(coeff.clone()),
            Monomial(n),
        );

        if coeff.abs() == conversion(1) && n == 0 {
            print!("1");
        }
    }

    println!();
}

#[allow(unused)]
macro_rules! print_square_matrix {
    ($s:literal, $g: expr) => {
        print_matrix::<$s, $s, { $s * $s }>($g);
    };
}

fn main() {
    let poly = read_poly(&std::env::args().nth(1).expect("Enter a polynomial"));
    let p = std::env::args()
        .nth(2)
        .expect("Enter a prime")
        .parse::<usize>()
        .unwrap();

    println!("poly is:");
    print_poly(&poly);

    println!("prime is {p}");

    // let poly1 = [12, 12, 0, 1].map(conversion);
    // let poly2 = [-4, 1].map(conversion);

    // let composition = berlekamp::composition(&poly1, &poly2);

    use berlekamp::Poly;

    // println!(
    //     "the composition of {} and {} is {}",
    //     Poly(&poly1),
    //     Poly(&poly2),
    //     Poly(&composition)
    // );

    // println!("poly2 is {factor}");

    // let poly = [-25, 0, 15, 0, -3, 0, 1].map(conversion);

    // let p = 5;

    let factors = berlekamp::factors(&poly, p);

    match factors {
        Ok(factors) => {
            println!("the irreducible factors are:");

            for factor in factors {
                print_poly(&factor);
            }
        }
        Err(gcd) => {
            println!(
                "the gcd of the polynomial with its derivative is {}",
                Poly(&gcd)
            );

            if gcd.len() == poly.len() {
                // the derivative is zero
                println!("the polynomial is a {p}-th power of some polynomial");

                return;
            }

            let (quotient, _) = berlekamp::poly_division(&poly, &gcd, conversion(p));

            println!(
                "the quotient of the polynomial by the above gcd is {}",
                Poly(&quotient)
            );

            let factors =
                berlekamp::factors(&quotient, p).expect("this quotient should be square-free");

            println!("the irreducible factors of the quotient, and hence of the polynomial, are:");

            for factor in factors {
                print_poly(&factor);
            }
        }
    }
}

#[allow(unused)]
fn read_poly(s: &str) -> Vec<BigInt> {
    let mut degree: usize;
    let mut coefficient = 1isize;
    let mut pending_coefficient = false;

    let mut iter = s.chars().peekable();

    let mut degrees_and_cos: Vec<(usize, BigInt)> = Vec::new();

    while let Some(c) = iter.next() {
        match c {
            '\n' | '\t' | ' ' => {}
            '+' => {
                if pending_coefficient {
                    degrees_and_cos.push((0usize, conversion(coefficient)));
                }

                pending_coefficient = false;
                coefficient = 1isize;
            }
            'x' => {
                pending_coefficient = false;

                if !matches!(iter.peek(), Some(c) if *c == '^') {
                    degrees_and_cos.push((1usize, conversion(coefficient)));
                    continue;
                }

                let _ = iter.next();

                degree = 0usize;

                while matches!(iter.peek(), Some(c) if c.is_digit(10)) {
                    degree *= 10;
                    degree += iter.next().unwrap() as usize - '0' as usize;
                }

                degrees_and_cos.push((degree, conversion(coefficient)));
            }
            c if c.is_digit(10) || c == '-' => {
                pending_coefficient = true;

                let mut negative = false;

                coefficient = if c.is_digit(10) {
                    c as isize - '0' as isize
                } else {
                    negative = true;
                    -1isize
                };

                let mut first_negative = negative;

                while matches!(iter.peek(), Some(c) if c.is_digit(10)) {
                    if first_negative {
                        coefficient = '0' as isize - iter.next().unwrap() as isize;

                        first_negative = false;

                        continue;
                    }

                    let factor = if negative {
                        '0' as isize - iter.next().unwrap() as isize
                    } else {
                        iter.next().unwrap() as isize - '0' as isize
                    };

                    coefficient *= 10isize;
                    coefficient += factor;
                }
            }
            _ => {
                panic!("invalid: {c}");
            }
        }
    }

    if pending_coefficient {
        degrees_and_cos.push((0usize, conversion(coefficient)));
    }

    let mut degree_co_map: std::collections::HashMap<usize, BigInt> = Default::default();

    for (degree, coefficient) in degrees_and_cos {
        if let Some(orig) = degree_co_map.get(&degree) {
            degree_co_map.insert(degree, orig.clone() + coefficient);
        } else {
            degree_co_map.insert(degree, coefficient);
        }
    }

    // degree_co_map.extend(degrees_and_cos);

    let mut max_degree = 0usize;
    let mut non_zero_p = false;

    for (d, c) in degree_co_map.iter() {
        if c == &conversion(0) {
            continue;
        }

        non_zero_p = true;

        max_degree = std::cmp::max(*d, max_degree);
    }

    if !non_zero_p {
        Vec::new()
    } else {
        let mut result: Vec<_> = std::iter::repeat_with(|| conversion(0))
            .take(max_degree + 1)
            .collect();

        for (d, c) in degree_co_map {
            if c == conversion(0) {
                continue;
            }

            *result.get_mut(d).unwrap() = c;
        }

        result
    }
}

/// Print an array as a square matrix nicely.
///
/// This has the same requirement as the function `power`.  Similarly,
/// one can use the macro `print_square_matrix` to correctly fill in
/// the dimensions automatically.
#[allow(unused)]
fn print_matrix<const X: usize, const Y: usize, const T: usize>(matrix: &[BigInt; T]) {
    if X * Y != T {
        panic!("dimensions do not match: {X} * {Y} is not {T}");
    }

    println!("[");

    let mut max_lens: [usize; Y] = [0; Y];

    for j in 0..Y {
        for i in 0..X {
            let entry_str = format!("{}", matrix.get(Y * i + j).unwrap());

            *max_lens.get_mut(j).unwrap() =
                std::cmp::max(*max_lens.get(j).unwrap(), entry_str.len());
        }
    }

    for i in 0..X {
        print!("  ");

        for j in 0..Y {
            let entry_str = format!("{}", matrix.get(Y * i + j).unwrap());

            let fill_in_space: String = std::iter::repeat(32 as char)
                .take(*max_lens.get(j).unwrap() - entry_str.len())
                .collect();

            print!("{}{} ", fill_in_space, entry_str);
        }

        println!();
    }
    println!("]");
}