pub(crate) const NORMK1: [u64; 4] = build_normalized_kmer::<4>();
pub(crate) const NORMK2: [u64; 16] = build_normalized_kmer::<16>();
pub(crate) const NORMK3: [u64; 64] = build_normalized_kmer::<64>();
pub(crate) const NORMK4: [u64; 256] = build_normalized_kmer::<256>();
pub(crate) const NORMK5: [u64; 1024] = build_normalized_kmer::<1024>();
pub(crate) const NORMK6: [u64; 4096] = build_normalized_kmer::<4096>();

include!(concat!(env!("OUT_DIR"), "/ln_class_tables.rs"));

const fn normalize_circular(kmer: u64, ws: usize) -> u64 {
    let mask = (1u64 << (ws * 2)) - 1;
    let mut canonical = kmer & mask;
    let mut current = canonical;
    let mut i = 0;
    while i < (ws - 1) {
        let top = (current >> ((ws - 1) * 2)) & 3;
        current = ((current << 2) | top) & mask;
        if current < canonical {
            canonical = current;
        }
        i += 1;
    }
    canonical
}

const fn build_normalized_kmer<const N: usize>() -> [u64; N] {
    let mut result = [0u64; N];
    let k = k_from_n::<N>();
    let shift = 64 - k * 2;
    let mut i = 0;
    while i < N {
        let la = (i as u64) << shift;
        let ra = i as u64;
        let rc_ra = revcomp_raw(la, k) >> shift;
        let circ = normalize_circular(ra, k);
        let circ_rc = normalize_circular(rc_ra, k);
        result[i] = if circ < circ_rc { circ } else { circ_rc };
        i += 1;
    }
    result
}

const fn revcomp_raw(x: u64, k: usize) -> u64 {
    let x = !x;
    let x = x.swap_bytes();
    let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
    let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
    x << (64 - 2 * k)
}

const fn k_from_n<const N: usize>() -> usize {
    match N {
        4 => 1,
        16 => 2,
        64 => 3,
        256 => 4,
        1024 => 5,
        4096 => 6,
        _ => panic!("N must be a power of 4"),
    }
}

pub(crate) const WS_MAX: usize = 6;

#[inline(always)]
pub(crate) const fn n_log_n(n: usize) -> f64 {
    N_LOG_N[n]
}

#[inline(always)]
pub(crate) const fn emax(k: usize, ws: usize) -> f64 {
    EMAX[k][ws]
}

#[inline(always)]
pub(crate) const fn log_nwords(k: usize, ws: usize) -> f64 {
    LOG_NWORDS[k][ws]
}

#[inline(always)]
pub(crate) const fn entropy_norm_kmer<const LEFT: bool, const K: usize>(kmer: u64) -> u64 {
    const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
    const NORM: [&[u64]; 7] = [&[], &NORMK1, &NORMK2, &NORMK3, &NORMK4, &NORMK5, &NORMK6];

    let shift = SHIFT[K];
    let ra = if LEFT { kmer >> shift } else { kmer };
    let canonical_ra = NORM[K][ra as usize];
    if LEFT {
        canonical_ra << shift
    } else {
        canonical_ra
    }
}

#[inline(always)]
pub(crate) const fn ln_class_size<const LEFT: bool, const K: usize>(kmer: u64) -> f64 {
    const SHIFT: [usize; 7] = [0, 62, 60, 58, 56, 54, 52];
    let ra = if LEFT { kmer >> SHIFT[K] } else { kmer };
    match K {
        1 => LN_CLASS1[ra as usize],
        2 => LN_CLASS2[ra as usize],
        3 => LN_CLASS3[ra as usize],
        4 => LN_CLASS4[ra as usize],
        5 => LN_CLASS5[ra as usize],
        6 => LN_CLASS6[ra as usize],
        _ => panic!("k must be 1..=6"),
    }
}
