use rug::Integer;
use std::collections::HashMap;
use tqdm::tqdm;

/// the modul of the discrete logarithm
const P: u64 = 12345679943u64;

pub fn solve(n: u64, x: Integer) -> (u64, Integer) {
    let x = x.modulo(&Integer::from(P - 1)).to_u64().unwrap();

    let mut n = n;
    let mut x = x.clone();
    for _ in tqdm(0..10_000) {
        n += 1;
        x = log(42, 1 + (x + n) % (P - 1), P).unwrap();
    }

    (n, Integer::from(x))
}

/// Computes the discrete logarithm of `b` to the base `g` modulo `m` using the Baby-Step-Giant-Step algorithm.
fn log(g: u64, b: u64, m: u64) -> Option<u64> {
    let g: u128 = g as u128;
    let b: u128 = b as u128;
    let m: u128 = m as u128;

    let s = m.isqrt() + 1;
    let mut r = HashMap::new();

    let mut a = b % m;
    if a == 1 {
        return Some(0);
    }

    for i in 0..s {
        r.insert(a, i);
        a *= g;
        a %= m;
    }

    let gs = pow(g as u64, s as u64, m as u64) as u128;
    a = gs;
    for j in 1..=s {
        if let Some(i) = r.get(&a) {
            return Some((j * s - i) as u64);
        }

        a *= gs;
        a %= m;
    }

    None
}

/// Computes the `k`-th power of `a` modulo `m`.
fn pow(a: u64, mut k: u64, m: u64) -> u64 {
    let mut a: u128 = a as u128;
    let m: u128 = m as u128;
    let mut b = 1u128;
    while k != 0 {
        if k % 2 == 1 {
            b *= a;
            b %= m;
        }
        a *= a;
        a %= m;

        k >>= 1;
    }
    b as u64
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn log_test() {
        assert_eq!(log(3, 12, 23), Some(4));
        assert_eq!(log(2, 3, 5), Some(3));

        log(42, 3816393254, 12345679943);
    }

    #[test]
    fn pow_test() {
        assert_eq!(pow(2, 3, 5), 3);
        assert_eq!(pow(3, 4, 23), 12);
    }
}