use num::{
    bigint::{BigInt, ToBigInt},
    rational::BigRational,
    traits::Euclid,
    Signed,
};

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

/// A simple conversion function from anything that can be converted
/// to a big integer to a big rational number.
#[allow(unused)]
fn ratconversion(n: impl ToBigInt) -> BigRational {
    BigRational::from_integer(n.to_bigint().unwrap())
}

/// Return the inverse of `a` modulo `p`.  Here `p` is assumed to be
/// relatively prime with `a`.
///
/// This is based on Euclidean algorithm.  In fact, it returns a
/// coefficient `x` such that ax + py = gcd(a, p), for some y.
pub fn modulo_inverse(a: BigInt, p: BigInt) -> BigInt {
    // perform Euclidean algorithm of a by p
    let mut quotients: Vec<BigInt> = Vec::new();
    let mut lhss: Vec<BigInt> = Vec::new();

    let mut divider: BigInt = p.clone();
    let mut divided: BigInt = a.clone();

    if divider <= conversion(1) {
        panic!("1 or 0 is not a prime");
    }

    while divider > conversion(1) {
        // The mighty compiler should simplify this to a single
        // instruction on a modern "enough" CPU.
        let div: BigInt = divided.div_euclid(&divider);
        let rem: BigInt = divided.rem_euclid(&divider);

        // for debugging

        // #[cfg(debug_assertions)]
        // println!("{divided} = {div} * {divider} + {rem}");

        quotients.push(div);
        lhss.push(divided);

        divided = divider;
        divider = rem;
    }

    if divider != conversion(1) {
        dbg!(&divider, &a, &p);
        panic!("they are not relatively prime");
    }

    assert!(!quotients.is_empty());
    assert!(!lhss.is_empty());

    let mut rhs = quotients.pop().unwrap();
    rhs *= conversion(-1);

    let mut lhs = conversion(1);

    // #[cfg(debug_assertions)]
    // println!("lhs = {lhs}, rhs = {rhs}");

    while let Some(q) = quotients.pop() {
        // already: a * lhs + b * rhs = 1

        // now: b = x - q * a

        // so 1 = a * lhs + rhs (x - q * a) = rhs * x + (lhs - rhs * q) * a

        let temp = rhs.clone();

        rhs = lhs - rhs * q;

        lhs = temp;

        // #[cfg(debug_assertions)]
        // println!("lhs = {lhs}, rhs = {rhs}");
    }

    // finally lhs * a + rhs * p = 1

    lhs.rem_euclid(&p)
}

/// Return the remainder of the Euclidean division of a by b, over the
/// finite field of p elements, where p is assumed to be a natural
/// prime number.
pub fn poly_rem(a: &[BigInt], b: &[BigInt], p: BigInt) -> Vec<BigInt> {
    if a.len() < b.len() {
        return a.to_vec();
    }

    if b.is_empty() {
        panic!("cannot divide by 0");
    }

    let b_leading_coefficient = b[b.len() - 1].clone();

    let b_leading_inverse = modulo_inverse(b_leading_coefficient, p.clone());

    let mut candidate: Vec<BigInt> = a.iter().map(|c| c.rem_euclid(&p)).collect();

    let mut diff: usize = a.len() - b.len();

    loop {
        let cand_leading = candidate[b.len() + diff - 1].clone();

        let leading_coefficient = b_leading_inverse.clone() * cand_leading;

        if leading_coefficient != conversion(0) {
            for (i, cand_ref) in candidate.iter_mut().skip(diff).take(b.len()).enumerate() {
                // println!(
                //     "bi = {}, leading_coefficient = {leading_coefficient}, cand = {cand_ref}",
                //     b[i]
                // );
                *cand_ref -= b[i].clone() * leading_coefficient.clone();
                // println!("cand = {cand_ref}");
                *cand_ref = cand_ref.rem_euclid(&p);
            }
        }

        if diff > 0 {
            diff -= 1;
        } else {
            break;
        }
    }

    let mut max_non_zero = 0usize;
    let mut non_zero = false;

    // dbg!();
    // print_poly(&candidate);

    for (index, cand) in candidate.iter().enumerate() {
        if *cand != conversion(0) {
            max_non_zero = index;
            non_zero = true;
        }
    }

    if !non_zero {
        return Vec::new();
    }

    candidate.into_iter().take(max_non_zero + 1).collect()
}

/// Return both the quotient and the remainder of the Euclidean
/// division of a by b, over the finite field of p elements, where p
/// is assumed to be a natural prime number.
pub fn poly_division(a: &[BigInt], b: &[BigInt], p: BigInt) -> (Vec<BigInt>, Vec<BigInt>) {
    if a.len() < b.len() {
        return (Vec::new(), a.to_vec());
    }

    if b.is_empty() {
        panic!("cannot divide by 0");
    }

    let b_leading_coefficient = b[b.len() - 1].clone();

    let b_leading_inverse = modulo_inverse(b_leading_coefficient, p.clone());

    let mut candidate: Vec<BigInt> = a.iter().map(|c| c.rem_euclid(&p)).collect();

    let mut diff: usize = a.len() - b.len();

    let mut quotient: Vec<BigInt> = Vec::with_capacity(diff);

    loop {
        let cand_leading = candidate[b.len() + diff - 1].clone();

        let leading_coefficient = b_leading_inverse.clone() * cand_leading;

        quotient.push(leading_coefficient.clone());

        if leading_coefficient != conversion(0) {
            for (i, cand_ref) in candidate.iter_mut().skip(diff).take(b.len()).enumerate() {
                // println!(
                //     "bi = {}, leading_coefficient = {leading_coefficient}, cand = {cand_ref}",
                //     b[i]
                // );
                *cand_ref -= b[i].clone() * leading_coefficient.clone();
                // println!("cand = {cand_ref}");
                *cand_ref = cand_ref.rem_euclid(&p);
            }
        }

        if diff > 0 {
            diff -= 1;
        } else {
            break;
        }
    }

    quotient = quotient.into_iter().rev().collect();

    let mut max_non_zero = 0usize;
    let mut non_zero = false;

    // dbg!();
    // print_poly(&candidate);

    for (index, cand) in candidate.iter().enumerate() {
        if *cand != conversion(0) {
            max_non_zero = index;
            non_zero = true;
        }
    }

    if !non_zero {
        (quotient, Vec::new())
    } else {
        (
            quotient,
            candidate.into_iter().take(max_non_zero + 1).collect(),
        )
    }
}

pub fn poly_gcd(a: &[BigInt], b: &[BigInt], p: BigInt) -> Vec<BigInt> {
    let mut a = a.to_vec();
    let mut b = b.to_vec();

    let mut rem = poly_rem(&a, &b, p.clone());

    while !rem.is_empty() {
        a = b;
        b = rem;

        rem = poly_rem(&a, &b, p.clone());
    }

    b
}

/// Return a vector of left-null-vectors over the finite field of p
/// elements, where p is assumed to be a natural prime number.
///
/// It is assumed that T is the square of X.
pub fn find_left_null_space(x: usize, t: usize, matrix: &[BigInt], p: usize) -> Vec<Vec<BigInt>> {
    assert_eq!(t, x * x);
    assert_eq!(matrix.len(), t);

    let mut matrix = matrix.to_vec();
    let mut result: Vec<Vec<BigInt>> = Vec::new();

    let mut used_pivots_rows: Vec<Option<usize>> =
        std::iter::repeat(None).take(x).collect::<Vec<_>>();

    // also perform row-operations on the identity matrix

    let mut identity: Vec<BigInt> = std::iter::repeat(conversion(0)).take(t).collect::<Vec<_>>();

    for i in 0..x {
        *identity.get_mut(i * x + i).unwrap() = conversion(1);
    }

    for i in 0..x {
        for j in 0..x {
            if used_pivots_rows.get(j).unwrap().is_some() {
                continue;
            }

            let mat_ij = matrix.get(i * x + j).unwrap().clone();

            if mat_ij == conversion(0) {
                continue;
            }

            *used_pivots_rows.get_mut(j).unwrap() = Some(i);

            let mat_ij_inverse = modulo_inverse(mat_ij.clone(), conversion(p));

            for k in 0..x {
                if k == i || matrix.get(k * x + j).unwrap() == &conversion(0) {
                    continue;
                }

                let mat_kj = matrix.get(k * x + j).unwrap().clone();

                for ell in 0..x {
                    let ident_iell = identity.get(i * x + ell).unwrap().clone();
                    let mat_iell = matrix.get(i * x + ell).unwrap().clone();

                    *identity.get_mut(k * x + ell).unwrap() -=
                        mat_ij_inverse.clone() * mat_kj.clone() * ident_iell;

                    *identity.get_mut(k * x + ell).unwrap() = identity
                        .get(k * x + ell)
                        .unwrap()
                        .rem_euclid(&conversion(p));

                    *matrix.get_mut(k * x + ell).unwrap() -=
                        mat_ij_inverse.clone() * mat_kj.clone() * mat_iell;

                    *matrix.get_mut(k * x + ell).unwrap() =
                        matrix.get(k * x + ell).unwrap().rem_euclid(&conversion(p));
                }
            }

            break;
        }
    }

    // println!("pivot: {used_pivots_rows:?}");

    let mut used_rows: Vec<bool> = std::iter::repeat(false)
        .take(x)
        .collect::<Vec<_>>()
        .as_slice()
        .try_into()
        .unwrap();

    for i in used_pivots_rows.iter().filter_map(|x| *x) {
        *used_rows.get_mut(i).unwrap() = true;
    }

    for (index, used) in used_rows.iter().copied().enumerate() {
        if !used {
            let mut newrow: Vec<BigInt> = Vec::with_capacity(x);

            for j in 0..x {
                newrow.push(identity.get(index * x + j).unwrap().clone());
            }

            result.push(newrow);
        }
    }

    result
}

fn derivative(poly: &[BigInt], p: usize) -> Vec<BigInt> {
    let big_p = conversion(p);

    let untrimmed = poly
        .iter()
        .enumerate()
        .skip(1)
        .map(|(n, a)| (conversion(n) * a.clone()).rem_euclid(&big_p));

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

    for (i, c) in untrimmed.clone().enumerate() {
        if c != conversion(0) {
            non_zero_p = true;
            max_non_zero = i;
        }
    }

    if !non_zero_p {
        Vec::new()
    } else {
        untrimmed.take(max_non_zero + 1).collect()
    }
}

/// Return true if and only if the polynomial has no repeated factors
/// modulo the prime p.
pub fn is_square_free(poly: &[BigInt], p: usize) -> Result<(), Vec<BigInt>> {
    let derivative = derivative(poly, p);

    if derivative.is_empty() {
        return Err(poly.to_vec());
    }

    let gcd = poly_gcd(poly, &derivative, conversion(p));

    if gcd.len() <= 1 {
        Ok(())
    } else {
        Err(gcd)
    }
}

#[derive(Debug, Clone, Default)]
pub enum BeRe {
    Prod(Vec<BigInt>, Vec<BigInt>),
    #[default]
    Irreducible,
    NonSquareFree,
}

impl std::fmt::Display for BeRe {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BeRe::Prod(a, b) => write!(f, "product of {} and {}", Poly(a), Poly(b)),
            BeRe::Irreducible => write!(f, "irreducible"),
            BeRe::NonSquareFree => write!(f, "not square-free"),
        }
    }
}

/// Return an irreducible polynomial that divides the polynomial, over
/// the finite field specified by p, which is assumed to be a positive
/// prime number.
///
/// Also the polynomial should be square-free.
pub fn berlekamp(poly: &[BigInt], p: usize) -> BeRe {
    // reject non-square-free polynomials
    if is_square_free(poly, p).is_err() {
        return BeRe::NonSquareFree;
    }

    // compute the basis matrix
    let n = poly.len();

    if n == 0 {
        return BeRe::Prod(Vec::default(), Vec::default());
    }

    // print!("poly: ");
    // print_poly(poly);
    // println!();

    let mut xqi_mat: Vec<_> = std::iter::repeat_with(|| conversion(0))
        .take(n * n)
        .collect();

    let mut xq = vec![conversion(1)];

    let rem = poly_rem(&xq, poly, conversion(p));

    for (co, result_mut) in rem.iter().zip(xqi_mat.iter_mut()) {
        *result_mut = co.clone();
    }

    xq = rem;

    for i in 1..=(n - 1) {
        let mut temp: Vec<_> = std::iter::repeat_with(|| conversion(0)).take(p).collect();

        temp.append(&mut xq);

        xq = temp;

        let rem = poly_rem(&xq, poly, conversion(p));

        // print!("i = {i}, rem = ");
        // print_poly(&rem);

        for (co, result_mut) in rem.iter().zip(xqi_mat.iter_mut().skip(n * i)) {
            *result_mut = co.clone();
        }
    }

    // now subtract the identity matrix

    for i in 0..n {
        *xqi_mat.get_mut(i * n + i).unwrap() -= conversion(1);
    }

    // println!("haha");
    // print_matrix::<3, 3, 9>(xqi_mat.as_slice().try_into().unwrap());

    // now find a basis for the left-null-space
    let basis = find_left_null_space(n, n * n, &xqi_mat, p);

    assert!(!basis.is_empty(), "the kernel should not be zero");

    // for (i, v) in basis.clone().into_iter().enumerate() {
    //     print!("{i}: ");
    //     print_matrix::<1, 3, 3>(v.as_slice().try_into().unwrap());
    // }

    if basis.len() == 1 {
        return BeRe::default();
    }

    // now find the corresponding polynomials

    // The first element correponds always to the constant polynomial
    // and serves no purpose to us, so we omit it here.

    let _irr_factors_num = basis.len();

    // dbg!(irr_factors_num);

    for v in basis.clone().into_iter().skip(1) {
        let mut max_non_zero = 0usize;
        let mut non_zero = false;

        for (index, cand) in v.iter().enumerate() {
            if *cand != conversion(0) {
                max_non_zero = index;
                non_zero = true;
            }
        }

        if !non_zero {
            panic!("this should not happen");
        }

        let mut v: Vec<_> = v.iter().take(max_non_zero + 1).cloned().collect();

        for _ in 0..p {
            *v.get_mut(0).unwrap() -= 1;

            // print_poly(&v);

            // Now take the gcd of the polynomial with v

            let gcd = poly_gcd(poly, &v, conversion(p));

            if (2..poly.len()).contains(&gcd.len()) {
                let (quotient, _) = poly_division(poly, &gcd, conversion(p));

                return BeRe::Prod(gcd, quotient);
            }
        }
    }

    println!("weird poly = {}", Poly(poly));

    for (n, v) in basis.clone().into_iter().enumerate() {
        println!("{n}-th basis: {}", Poly(&v));
    }

    panic!("this should not happen");
}

/// Return all irreducible factors of the polynomial modulo the prime
/// p.
///
/// If the polynomial has repeated factors, return an error.
#[allow(clippy::result_unit_err)]
pub fn factors(poly: &[BigInt], p: usize) -> Result<Vec<Vec<BigInt>>, Vec<BigInt>> {
    if let Err(gcd) = is_square_free(poly, p) {
        return Err(gcd);
    }

    let mut result = Vec::new();
    let mut candidates: Vec<Vec<BigInt>> = vec![poly.to_vec()];

    while let Some(poly) = candidates.pop() {
        let bere = berlekamp(&poly, p);

        match bere {
            BeRe::Prod(x, y) => {
                candidates.push(x);
                candidates.push(y);
            }
            BeRe::Irreducible => result.push(poly),
            BeRe::NonSquareFree => return Err(Vec::new()),
        }
    }

    Ok(result)
}

/// Return a sub-slice that trims the leading zeroes.
pub fn trim(p: &[BigInt]) -> &[BigInt] {
    let mut max_non_zero = 0usize;
    let mut non_zero_p = false;

    for (n, c) in p.iter().enumerate() {
        if c != &conversion(0) {
            non_zero_p = true;
            max_non_zero = n;
        }
    }

    if !non_zero_p {
        &p[0..0]
    } else {
        &p[0..=max_non_zero]
    }
}

/// Return the product of two polynomials.
pub fn poly_prod(p1: &[BigInt], p2: &[BigInt]) -> Vec<BigInt> {
    let p1 = trim(p1);

    let p2 = trim(p2);

    // degrees sum up, so the lengths have to subtract one
    let mut result: Vec<BigInt> = std::iter::repeat_with(|| conversion(0))
        .take(p1.len() + p2.len() - 1)
        .collect();

    let mut degree_map: std::collections::HashMap<usize, BigInt> =
        std::collections::HashMap::with_capacity(p1.len() + p2.len());

    for (n1, c1) in p1.iter().enumerate() {
        for (n2, c2) in p2.iter().enumerate() {
            let new_degree = n1 + n2;

            if let Some(orig) = degree_map.get(&new_degree) {
                degree_map.insert(new_degree, orig + c1 * c2);
            } else {
                degree_map.insert(new_degree, c1 * c2);
            }
        }
    }

    for (n, c) in degree_map {
        *result.get_mut(n).unwrap() = c;
    }

    result
}

pub fn power(p: &[BigInt], n: usize) -> Vec<BigInt> {
    if n == 0 {
        return Vec::new();
    }

    let mut result = p.to_vec();

    for _ in 0..(n - 1) {
        result = poly_prod(&result, p);
    }

    result
}

/// Return c * p + q.
pub fn poly_sum(p: &[BigInt], q: &[BigInt], cp: BigInt) -> Vec<BigInt> {
    let p = trim(p);
    let q = trim(q);

    if p.is_empty() {
        return q.to_vec();
    }

    if q.is_empty() {
        return p.iter().map(|c| c * &cp).collect();
    }

    if p.len() >= q.len() {
        let mut result: Vec<_> = p.iter().map(|c| c * &cp).collect();

        for (n, c) in q.iter().enumerate() {
            *result.get_mut(n).unwrap() += c.clone();
        }

        result
    } else {
        let mut result = q.to_vec();

        for (n, c) in p.iter().enumerate() {
            *result.get_mut(n).unwrap() += c.clone() * cp.clone();
        }

        result
    }
}

/// Return the composition of p and q, i.e. p ( q ( x ) ) as a
/// polynomial in x.
pub fn composition(p: &[BigInt], q: &[BigInt]) -> Vec<BigInt> {
    let p = trim(p);
    let q = trim(q);

    let mut power: Vec<_> = vec![conversion(1)];

    let mut result: Vec<BigInt> = Vec::new();

    for c in p.iter() {
        result = poly_sum(&power, &result, c.clone());

        power = poly_prod(&power, q);
    }

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

pub struct Poly<'a>(pub &'a [BigInt]);

impl<'a> std::fmt::Display for Poly<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let poly = self.0;

        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;
            }

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

            if coeff.abs() == conversion(1) && n == 0 {
                write!(f, "1")?;
            }
        }

        Ok(())
    }
}

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