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

use berlekamp::Poly;

/// 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()
}

#[allow(unused)]
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();

    if std::env::args().len() >= 4 {
        // Search for primes modulo which the polynomial is irreducible instead.

        let primes = generate_primes(p);

        println!("Searching {} primes up to {p}...", primes.len());

        for p in primes {
            let factors = berlekamp::factors(&poly, p);

            let total = factors.iter().fold(0usize, |sum, (deg, _)| sum + deg);

            if total == 1 {
                println!("The polynomial {} is irreducible mod {p}", Poly(&poly));

                return;
            }
        }

        println!("The polynomial is reducible modulo primes up to {p}.");

        return;
    }

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

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

    // let poly1 = [-20, -21, 0, 1].map(conversion);
    // let poly2 = [1, 3].map(conversion);

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

    // 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);

    print!("{} ≡ ", Poly(&poly));

    let mut total = 0usize;

    for (index, (mul, f)) in factors.into_iter().enumerate() {
        total += mul;

        match mul {
            0 => {
                dbg!("zero multiplicity?");
            }
            1 => {
                print!("{}({})", if index > 0 { " * " } else { "" }, Poly(&f));
            }
            _ => {
                print!(
                    "{}({})^{}",
                    if index > 0 { " * " } else { "" },
                    Poly(&f),
                    mul
                );
            }
        }
    }

    println!(" (mod {p})");

    if total == 1 {
        println!("the polynomial is irreducible");
    }
}

#[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!("]");
}

#[allow(unused)]
/// Return a vector of primes <= BOUND.
///
/// This is naïve and slow, so BOUND should not be too large.
fn generate_primes(bound: usize) -> Vec<usize> {
    if bound <= 1 {
        return Vec::new();
    }

    let sqrt = (bound as f64).sqrt() as usize;

    let mut result: Vec<usize> = Vec::with_capacity(bound);

    let mut records: Vec<bool> = std::iter::repeat(true).take(bound).collect();

    records[0] = false;
    records[1] = false;

    for i in 2..=(sqrt + 1) {
        if !matches!(records.get(i), Some(true)) {
            continue;
        }

        let mut multiple = i + i;

        while let Some(rec) = records.get_mut(multiple) {
            *rec = false;
            multiple += i;
        }
    }

    result.extend(
        records
            .iter()
            .enumerate()
            .filter_map(|(n, rec)| rec.then_some(n)),
    );

    result
}