use obikseq::kmer::{Minimizer, hash_kmer};
use obikseq::params;

use crate::encoding::encode_nuc;
use crate::entropy_table::{WS_MAX, emax, entropy_norm_kmer, ln_class_size, log_nwords, n_log_n};

// ── Stack-allocated ring buffer ───────────────────────────────────────────────

/// Fixed-capacity ring buffer backed by a stack array.
/// N must be a power of two; operations are branchless via `% N`.
struct Ring<T: Copy + Default, const N: usize> {
    buf: [T; N],
    head: usize,
    len: usize,
}

impl<T: Copy + Default, const N: usize> Ring<T, N> {
    #[inline]
    fn new() -> Self {
        Self {
            buf: [T::default(); N],
            head: 0,
            len: 0,
        }
    }
    #[inline]
    fn clear(&mut self) {
        self.len = 0;
        self.head = 0;
    }

    #[inline]
    fn push_back(&mut self, val: T) {
        self.buf[(self.head + self.len) % N] = val;
        self.len += 1;
    }
    #[inline]
    fn pop_front(&mut self) -> T {
        let val = self.buf[self.head];
        self.head = (self.head + 1) % N;
        self.len -= 1;
        val
    }
    #[inline]
    fn pop_back(&mut self) -> T {
        self.len -= 1;
        self.buf[(self.head + self.len) % N]
    }
    #[inline]
    fn front(&self) -> Option<&T> {
        if self.len > 0 {
            Some(&self.buf[self.head])
        } else {
            None
        }
    }
    #[inline]
    fn back(&self) -> Option<&T> {
        if self.len > 0 {
            Some(&self.buf[(self.head + self.len - 1) % N])
        } else {
            None
        }
    }

}

// ── MmerItem ──────────────────────────────────────────────────────────────────

#[derive(Clone, Copy, Default)]
struct MmerItem {
    /// 0-based position of this m-mer's first base within the current segment.
    position: usize,
    /// Raw canonical m-mer value (right-aligned).
    canonical: u64,
    /// mix64 hash of the canonical m-mer, used as the random ordering key.
    hash: u64,
}

// ── RollingStat ───────────────────────────────────────────────────────────────

pub struct RollingStat {
    entropy_max_k: usize,
    k: usize,
    m: usize,
    steady: bool,
    rolling_k: u64,
    rolling_rck: u64,
    k_mask: u64,
    m_mask: u64,
    received: usize,

    // Sliding-window queues — stack-allocated, capacity ≤ k ≤ 31.
    k1q: Ring<u64, 32>,
    k2q: Ring<u64, 32>,
    k3q: Ring<u64, 32>,
    k4q: Ring<u64, 32>,
    k5q: Ring<u64, 32>,
    k6q: Ring<u64, 32>,
    minimier: Ring<MmerItem, 32>,

    // Frequency count arrays.
    // Max count per cell ≤ k ≤ 31 → u8 is sufficient.
    k1c: [u8; 4],
    k2c: [u8; 16],
    k3c: [u8; 64],
    k4c: [u8; 256],
    k5c: [u8; 1024],
    k6c: [u8; 4096],

    sum_f_log_f: [f64; WS_MAX + 1],
    sum_f_log_s: [f64; WS_MAX + 1],
}

impl RollingStat {
    pub fn new(entropy_max_k: usize) -> Self {
        let k = params::k();
        let m = params::m();
        Self {
            entropy_max_k,
            k,
            m,
            steady: false,
            rolling_k: 0,
            rolling_rck: 0,
            k_mask: (!0u64) >> (64 - k * 2),
            m_mask: (!0u64) >> (64 - m * 2),
            received: 0,
            k1q: Ring::new(),
            k2q: Ring::new(),
            k3q: Ring::new(),
            k4q: Ring::new(),
            k5q: Ring::new(),
            k6q: Ring::new(),
            minimier: Ring::new(),
            k1c: [0; 4],
            k2c: [0; 16],
            k3c: [0; 64],
            k4c: [0; 256],
            k5c: [0; 1024],
            k6c: [0; 4096],
            sum_f_log_f: [0.0; WS_MAX + 1],
            sum_f_log_s: [0.0; WS_MAX + 1],
        }
    }

    pub fn reset(&mut self) {
        self.rolling_k = 0;
        self.rolling_rck = 0;
        self.received = 0;
        self.steady = false;

        //        for i in self.k1q.iter() { self.k1c[i as usize] = 0; }
        //        for i in self.k2q.iter() { self.k2c[i as usize] = 0; }
        //        for i in self.k3q.iter() { self.k3c[i as usize] = 0; }
        //        for i in self.k4q.iter() { self.k4c[i as usize] = 0; }
        //        for i in self.k5q.iter() { self.k5c[i as usize] = 0; }
        //        for i in self.k6q.iter() { self.k6c[i as usize] = 0; }

        self.k1c.fill(0);
        self.k2c.fill(0);
        self.k3c.fill(0);
        self.k4c.fill(0);
        self.k5c.fill(0);
        self.k6c.fill(0);

        self.k1q.clear();
        self.k2q.clear();
        self.k3q.clear();
        self.k4q.clear();
        self.k5q.clear();
        self.k6q.clear();
        self.minimier.clear();

        self.sum_f_log_f = [0.0; WS_MAX + 1];
        self.sum_f_log_s = [0.0; WS_MAX + 1];
    }

    #[inline]
    fn update_sums_decrement<const K: usize>(
        sum_f_log_f: &mut [f64; WS_MAX + 1],
        sum_f_log_s: &mut [f64; WS_MAX + 1],
        canonical: u64,
        f: usize,
    ) {
        sum_f_log_f[K] += n_log_n(f - 1) - n_log_n(f);
        sum_f_log_s[K] -= ln_class_size::<false, K>(canonical);
    }

    #[inline]
    fn update_sums_increment<const K: usize>(
        sum_f_log_f: &mut [f64; WS_MAX + 1],
        sum_f_log_s: &mut [f64; WS_MAX + 1],
        canonical: u64,
        g: usize,
    ) {
        sum_f_log_f[K] += n_log_n(g + 1) - n_log_n(g);
        sum_f_log_s[K] += ln_class_size::<false, K>(canonical);
    }

    pub fn push(&mut self, nuc: u8) {
        let k = self.k;
        let m = self.m;

        let bnuc = encode_nuc(nuc);
        let cnuc = bnuc ^ 3;

        self.rolling_k = ((self.rolling_k << 2) | (bnuc as u64)) & self.k_mask;
        self.rolling_rck =
            ((self.rolling_rck >> 2) | ((cnuc as u64) << ((k - 1) * 2))) & self.k_mask;

        let canonical_k1 = entropy_norm_kmer::<false, 1>(self.rolling_k & 3);
        let canonical_k2 = entropy_norm_kmer::<false, 2>(self.rolling_k & 15);
        let canonical_k3 = entropy_norm_kmer::<false, 3>(self.rolling_k & 63);
        let canonical_k4 = entropy_norm_kmer::<false, 4>(self.rolling_k & 255);
        let canonical_k5 = entropy_norm_kmer::<false, 5>(self.rolling_k & 1023);
        let canonical_k6 = entropy_norm_kmer::<false, 6>(self.rolling_k & 4095);

        self.received += 1;

        if self.received >= m {
            let possible_canonical_m =
                (self.rolling_k & self.m_mask).min(self.rolling_rck >> ((k - m) * 2));
            let possible_hash_m = hash_kmer(possible_canonical_m << 64 - m * 2);
            let possible_pos_m = self.received - m;

            while self
                .minimier
                .back()
                .map_or(false, |it| it.hash >= possible_hash_m)
            {
                self.minimier.pop_back();
            }
            self.minimier.push_back(MmerItem {
                position: possible_pos_m,
                canonical: possible_canonical_m,
                hash: possible_hash_m,
            });

            if self.received > k {
                while self
                    .minimier
                    .front()
                    .map_or(false, |it| it.position + k < self.received)
                {
                    self.minimier.pop_front();
                }
            }
        }

        if self.received > k {
            let old1 = self.k1q.pop_front();
            let f1 = self.k1c[old1 as usize] as usize;
            Self::update_sums_decrement::<1>(
                &mut self.sum_f_log_f,
                &mut self.sum_f_log_s,
                old1,
                f1,
            );
            self.k1c[old1 as usize] -= 1;

            let old2 = self.k2q.pop_front();
            let f2 = self.k2c[old2 as usize] as usize;
            Self::update_sums_decrement::<2>(
                &mut self.sum_f_log_f,
                &mut self.sum_f_log_s,
                old2,
                f2,
            );
            self.k2c[old2 as usize] -= 1;

            let old3 = self.k3q.pop_front();
            let f3 = self.k3c[old3 as usize] as usize;
            Self::update_sums_decrement::<3>(
                &mut self.sum_f_log_f,
                &mut self.sum_f_log_s,
                old3,
                f3,
            );
            self.k3c[old3 as usize] -= 1;

            let old4 = self.k4q.pop_front();
            let f4 = self.k4c[old4 as usize] as usize;
            Self::update_sums_decrement::<4>(
                &mut self.sum_f_log_f,
                &mut self.sum_f_log_s,
                old4,
                f4,
            );
            self.k4c[old4 as usize] -= 1;

            let old5 = self.k5q.pop_front();
            let f5 = self.k5c[old5 as usize] as usize;
            Self::update_sums_decrement::<5>(
                &mut self.sum_f_log_f,
                &mut self.sum_f_log_s,
                old5,
                f5,
            );
            self.k5c[old5 as usize] -= 1;

            let old6 = self.k6q.pop_front();
            let f6 = self.k6c[old6 as usize] as usize;
            Self::update_sums_decrement::<6>(
                &mut self.sum_f_log_f,
                &mut self.sum_f_log_s,
                old6,
                f6,
            );
            self.k6c[old6 as usize] -= 1;
        }

        if self.steady {
            let g1 = self.k1c[canonical_k1 as usize] as usize;
            Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1);
            self.k1c[canonical_k1 as usize] += 1;
            self.k1q.push_back(canonical_k1);

            let g2 = self.k2c[canonical_k2 as usize] as usize;
            Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2);
            self.k2c[canonical_k2 as usize] += 1;
            self.k2q.push_back(canonical_k2);

            let g3 = self.k3c[canonical_k3 as usize] as usize;
            Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3);
            self.k3c[canonical_k3 as usize] += 1;
            self.k3q.push_back(canonical_k3);

            let g4 = self.k4c[canonical_k4 as usize] as usize;
            Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4);
            self.k4c[canonical_k4 as usize] += 1;
            self.k4q.push_back(canonical_k4);

            let g5 = self.k5c[canonical_k5 as usize] as usize;
            Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5);
            self.k5c[canonical_k5 as usize] += 1;
            self.k5q.push_back(canonical_k5);

            let g6 = self.k6c[canonical_k6 as usize] as usize;
            Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6);
            self.k6c[canonical_k6 as usize] += 1;
            self.k6q.push_back(canonical_k6);
        } else {
            self.push_warmup_increments(
                canonical_k1, canonical_k2, canonical_k3,
                canonical_k4, canonical_k5, canonical_k6,
            );
        }
    }

    #[cold]
    #[inline(never)]
    fn push_warmup_increments(
        &mut self,
        canonical_k1: u64, canonical_k2: u64, canonical_k3: u64,
        canonical_k4: u64, canonical_k5: u64, canonical_k6: u64,
    ) {
        let g1 = self.k1c[canonical_k1 as usize] as usize;
        Self::update_sums_increment::<1>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k1, g1);
        self.k1c[canonical_k1 as usize] += 1;
        self.k1q.push_back(canonical_k1);

        if self.received >= 2 {
            let g2 = self.k2c[canonical_k2 as usize] as usize;
            Self::update_sums_increment::<2>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k2, g2);
            self.k2c[canonical_k2 as usize] += 1;
            self.k2q.push_back(canonical_k2);

            if self.received >= 3 {
                let g3 = self.k3c[canonical_k3 as usize] as usize;
                Self::update_sums_increment::<3>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k3, g3);
                self.k3c[canonical_k3 as usize] += 1;
                self.k3q.push_back(canonical_k3);

                if self.received >= 4 {
                    let g4 = self.k4c[canonical_k4 as usize] as usize;
                    Self::update_sums_increment::<4>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k4, g4);
                    self.k4c[canonical_k4 as usize] += 1;
                    self.k4q.push_back(canonical_k4);

                    if self.received >= 5 {
                        let g5 = self.k5c[canonical_k5 as usize] as usize;
                        Self::update_sums_increment::<5>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k5, g5);
                        self.k5c[canonical_k5 as usize] += 1;
                        self.k5q.push_back(canonical_k5);

                        if self.received >= 6 {
                            let g6 = self.k6c[canonical_k6 as usize] as usize;
                            Self::update_sums_increment::<6>(&mut self.sum_f_log_f, &mut self.sum_f_log_s, canonical_k6, g6);
                            self.k6c[canonical_k6 as usize] += 1;
                            self.k6q.push_back(canonical_k6);
                            self.steady = true;
                        }
                    }
                }
            }
        }
    }

    pub fn ready(&self) -> bool {
        self.received >= self.k
    }

    pub fn minimizer_position(&self) -> Option<usize> {
        if self.ready() {
            self.minimier.front().map(|it| it.position)
        } else {
            None
        }
    }

    pub fn canonical_minimizer_raw(&self) -> Option<u64> {
        if self.ready() {
            self.minimier.front().map(|it| it.canonical)
        } else {
            None
        }
    }

    pub fn canonical_minimizer(&self) -> Option<Minimizer> {
        self.canonical_minimizer_raw()
            .map(|raw| Minimizer::from_raw_unchecked(raw << (64 - self.m * 2)))
    }

    pub fn entropy(&self, order: usize) -> Option<f64> {
        if !self.ready() {
            return None;
        }
        let k = self.k;
        let em = emax(k, order);
        if em <= 0.0 {
            return Some(1.0);
        }
        let nwords = k - order + 1;
        let log_nw = log_nwords(k, order);
        let nw_f = nwords as f64;
        let h_corr = log_nw + (self.sum_f_log_s[order] - self.sum_f_log_f[order]) / nw_f;
        Some((h_corr / em).max(0.0))
    }

    pub fn normalized_entropy(&self) -> Option<f64> {
        if !self.ready() {
            return None;
        }
        let min_e = (1..=self.entropy_max_k)
            .filter_map(|ws| self.entropy(ws))
            .fold(f64::MAX, f64::min);
        Some(if min_e == f64::MAX { 1.0 } else { min_e })
    }
}
