use std::fs::File;
use std::io::{BufRead, BufReader};
use std::str::FromStr;
use std::sync::LazyLock;
use rug::Integer;

static DATA: LazyLock<Vec<u64>> = LazyLock::new(|| {
    let file = File::open("./res/07_weird_assembly_machine.data");
    let lines = BufReader::new(file.unwrap()).lines();

    let mut data = Vec::new();
    for line in lines {
        data.push(u64::from_str(&line.unwrap()).unwrap());
    }
    data
});

const LIMIT: u64 = 1_000_000_000_000_000_000u64;

pub fn solve(n: u64, x: Integer) -> (u64, Integer) {
    let x = x.to_u64_wrapping();
    let (n, x) = h_n(n, x, LIMIT - n);
    (n, Integer::from(x))
}

fn h(x: u64, n: u64) -> u64 {
    g(sub(f(x), n))
}

pub fn h_n(mut n: u64, mut x: u64, k: u64) -> (u64, u64) {
    n += 1;
    x = sub(f(x), n);

    let fast_k = (k - 1) / 192 * 192;
    let sum = fast_fg_sum(n , fast_k);
    x = sub(fg_n(x, fast_k), sum);
    n += fast_k;

    let remaining = (k - 1) - fast_k;
    let sum = fg_sum(n, remaining);
    x = sub(fg_n(x, remaining), sum);
    n += remaining;

    x = g(x);

    (n, x)
}

fn f(x: u64) -> u64 {
    evaluate(x, &DATA[..128])
}

fn g(x: u64) -> u64 {
    evaluate(x, &DATA[128..])
}

/// Quickly computes [`f`](f)([`g`](g)(x)).
fn fg(x: u64) -> u64 {
    !x.rotate_left(17)
}

/// Computes [`fg`](fg)(...[`fg`](fg)(x)...) where [`fg`] is applied `n` times.
fn fg_n(mut x: u64, n: u64) -> u64 {
    for _ in 0..n%64 {
        x = fg(x);
    }
    x
}

/// Computes the sum from `i = 1` to `k` over `fg^(k-i)(n + i)`.
fn fg_sum(n: u64, k: u64) -> u64 {
    let mut sum = 0u64;
    for i in 1..=k {
        sum = add(sum, fg_n(n + i, k - i));
    }
    sum
}

/// Quickly computes `fg_sum` when `k` is divisible by `192`.
fn fast_fg_sum(n: u64, k: u64) -> u64 {
    debug_assert!(k % 192 == 0);
    let mut result = fg_sum(n, 0);
    result = add(result, ((13556435138434861179u128 * (k / 192) as u128) % (u64::MAX as u128)) as u64);
    result
}

/// Evaluates a polynomial over `GF(2)[X] / (X^64 + 1)`
fn evaluate(x: u64, data: &[u64]) -> u64 {
    let mut out = 0;
    for i in 0..data.len() {
        out = xor_mul(out, x) ^ data[i];
    }
    out
}

fn add(a: u64, b: u64) -> u64 {
    let (result, overflow) = a.overflowing_add(b);
    result + overflow as u64
}

fn sub(x: u64, mut n: u64) -> u64 {
    if x <= n {
        n += 1;
    }
    x.wrapping_sub(n)
}

#[cfg(not(all(target_arch = "x86_64", target_feature = "pclmulqdq")))]
fn xor_mul(a: u64, mut b: u64) -> u64 {
    let mut x: u64 = 0;
    while b != 0 {
        let c = b.trailing_zeros();
        b = b ^ 1u64.rotate_left(c);
        x = x ^ a.rotate_left(c);
    }
    x
}

#[cfg(all(target_arch = "x86_64", target_feature = "pclmulqdq"))]
fn xor_mul(a: u64, b: u64) -> u64 {
    use std::arch::x86_64::*;
    unsafe {
        let a_vec = _mm_set_epi64x(0, a as i64);
        let b_vec = _mm_set_epi64x(0, b as i64);
        let result = _mm_clmulepi64_si128(a_vec, b_vec, 0x00);

        // Extract both 64-bit halves and XOR them
        let low = _mm_extract_epi64(result, 0) as u64;
        let high = _mm_extract_epi64(result, 1) as u64;
        low ^ high
    }
}