mfsk-core 0.5.0

//! Protocol-agnostic soft-decision LLR computation.
//!
//! Extracts complex tone spectra for each data symbol, then computes four
//! log-likelihood ratio variants (llra/b/c/d) matching WSJT-X `ft8b.f90`
//! convention — three `nsym = 1, 2, 3` grouping schemes plus a bit-by-bit
//! normalised variant. Parameterised over any [`Protocol`]: NTONES,
//! BITS_PER_SYMBOL, and the SYNC_BLOCKS layout drive the inner loops.

use alloc::vec;
use alloc::vec::Vec;

use num_complex::Complex;
#[cfg(not(feature = "std"))]
use num_traits::Float;

use super::Protocol;
use crate::core::fft::default_planner;
use crate::core::scalar::{Cmplx, LlrScalar, SpecScalar, complex_slice_as_cmplx_f32};

// ──────────────────────────────────────────────────────────────────────────
// LLR bundle
// ──────────────────────────────────────────────────────────────────────────

/// Four LLR vectors (a, b, c, d) of length `codeword_len()` bits in
/// LDPC bit-index order, generic over the [`LlrScalar`] storage type.
/// `LlrSet<f32>` (default) is the host path; `LlrSet<Q11i16>` is the
/// embedded fixed-point path used together with `bp_decode_generic_nms`.
#[derive(Clone)]
pub struct LlrSet<T: LlrScalar = f32> {
    /// nsym=1 soft metrics, scaled (matches WSJT-X llra).
    pub llra: Vec<T>,
    /// nsym=2 soft metrics, scaled (matches WSJT-X llrb).
    pub llrb: Vec<T>,
    /// nsym=3 soft metrics, scaled (matches WSJT-X llrc).
    pub llrc: Vec<T>,
    /// nsym=1 bit-normalised (matches WSJT-X llrd).
    pub llrd: Vec<T>,
}

/// Default LLR scale factor from WSJT-X ft8b.f90. Individual protocols may
/// override via `ModulationParams::LLR_SCALE`.
pub const LLR_SCALE: f32 = 2.83;

// ──────────────────────────────────────────────────────────────────────────
// Symbol spectra
// ──────────────────────────────────────────────────────────────────────────

/// Extract complex tone spectra for every channel symbol.
///
/// Returns a flat row-major `Vec<Complex<f32>>` of length `N_SYMBOLS × NTONES`;
/// row `k` / column `t` holds the k-th symbol's t-th tone amplitude, scaled
/// by 1/1000 (matching WSJT-X).
///
/// `i_start` is the sample index in `cd0` of the first symbol, from fine sync.
pub fn symbol_spectra<P: Protocol>(cd0: &[Complex<f32>], i_start: usize) -> Vec<Complex<f32>> {
    let ntones = P::NTONES as usize;
    let n_sym = P::N_SYMBOLS as usize;
    let ds_spb = (P::NSPS / P::NDOWN) as usize;

    let mut planner = default_planner();
    let fft = planner.plan_forward(ds_spb);

    let mut cs = vec![Complex::new(0.0f32, 0.0); n_sym * ntones];
    let mut buf = vec![Complex::new(0.0f32, 0.0); ds_spb];

    for k in 0..n_sym {
        let i1 = i_start + k * ds_spb;
        for (j, b) in buf.iter_mut().enumerate() {
            *b = if i1 + j < cd0.len() {
                cd0[i1 + j]
            } else {
                Complex::new(0.0, 0.0)
            };
        }
        fft.process(&mut buf);
        for (t, bin) in buf.iter().take(ntones).enumerate() {
            cs[k * ntones + t] = *bin / 1000.0;
        }
    }
    cs
}

// ──────────────────────────────────────────────────────────────────────────
// Helpers
// ──────────────────────────────────────────────────────────────────────────

// Data-chunk layout (slots between / around sync blocks) is shared
// with the TX side; reuse [`crate::core::tx::data_chunks`] so any
// frame layout the encoder honours is decoded the same way.
use crate::core::tx::data_chunks;

/// Decompose `i` into `nsym` base-`ntones` digits, most significant first.
#[inline]
fn base_digits(mut i: usize, ntones: usize, nsym: usize) -> Vec<usize> {
    let mut out = vec![0usize; nsym];
    for j in (0..nsym).rev() {
        out[j] = i % ntones;
        i /= ntones;
    }
    out
}

#[inline]
fn normalize_bmet(bmet: &mut [f32]) {
    let n = bmet.len() as f32;
    let mean = bmet.iter().sum::<f32>() / n;
    let mean_sq = bmet.iter().map(|&x| x * x).sum::<f32>() / n;
    let var = mean_sq - mean * mean;
    let sig = if var > 0.0 {
        var.sqrt()
    } else {
        mean_sq.sqrt()
    };
    if sig > 0.0 {
        bmet.iter_mut().for_each(|x| *x /= sig);
    }
}

// ──────────────────────────────────────────────────────────────────────────
// LLR computation
// ──────────────────────────────────────────────────────────────────────────

/// Compute soft LLRs from the flat symbol-spectra vector.
///
/// All four LLR variants (llra, llrb, llrc, llrd) are produced.
/// Costs grow with `nsym`: nsym=3 alone is ~80 % of the work
/// because nt = 512 = 8³ tone-combinations. Callers that only need
/// `llra`/`llrd` (the BP-only path) should use [`compute_llr_fast`]
/// instead — it caps at nsym=1 and skips the heavy nsym=2/3 loops.
///
/// The `Complex<f32>` cs API is a layout-compatible wrapper around
/// the generic `compute_llr_generic` — use the generic form when
/// the caller already holds [`Cmplx<S>`] storage (e.g. the
/// embedded fixed-point cs Box on Phase 2+ of the i16 migration).
pub fn compute_llr<P: Protocol, T: LlrScalar>(cs: &[Complex<f32>]) -> LlrSet<T> {
    compute_llr_generic::<P, f32, T>(complex_slice_as_cmplx_f32(cs), 3)
}

/// Same as [`compute_llr`] but stops at nsym=1. `llrb`/`llrc` come
/// back zero-filled. Use when the caller will only ever read
/// `llra` (or `llrd`), e.g. embedded `decode_block` with
/// `DecodeDepth::Bp`. ~5× faster than the full computation.
pub fn compute_llr_fast<P: Protocol, T: LlrScalar>(cs: &[Complex<f32>]) -> LlrSet<T> {
    compute_llr_generic::<P, f32, T>(complex_slice_as_cmplx_f32(cs), 1)
}

/// Compute the unnormalised, unscaled bit-metric arrays for ONE
/// `nsym` level. Internal helper shared by [`compute_llr_generic`]
/// (which calls it 1..=max_nsym times) and the lazy single-nsym
/// [`compute_llr_partial`] used to feed the BP staircase one variant
/// at a time without re-doing the cheaper levels.
///
/// `bmet_primary` receives `max_one - max_zero` (= llra at nsym=1,
/// llrb at nsym=2, llrc at nsym=3). At nsym=1 only, `bmet_norm`
/// receives the bit-normalised variant (= llrd). For nsym ≥ 2 pass
/// `bmet_norm = None`.
fn fill_bmet_for_nsym<P: Protocol, S: SpecScalar>(
    cs: &[Cmplx<S>],
    nsym: usize,
    bmet_primary: &mut [f32],
    bmet_norm: Option<&mut [f32]>,
) {
    let ntones = P::NTONES as usize;
    let bps = P::BITS_PER_SYMBOL as usize;
    let gray_map = P::GRAY_MAP;
    let chunks = data_chunks::<P>();
    let codeword_len = bmet_primary.len();

    let nt = ntones.pow(nsym as u32);
    let ibmax = bps * nsym - 1;
    let mut s2 = vec![0.0f32; nt];

    // Bit-normalised array (llrd) is only produced for nsym=1.
    let mut bmet_norm_holder = bmet_norm;

    let mut chunk_bit_base = 0usize;
    for &(chunk_start_sym, chunk_len) in &chunks {
        let mut k = 0usize;
        while k + nsym <= chunk_len {
            let ks = chunk_start_sym + k;

            // |Σ cs_k[gray[idx_k]]| for each tone-combination.
            for (i, s2_i) in s2.iter_mut().enumerate() {
                let digits = base_digits(i, ntones, nsym);
                let mut sum_re = 0.0f32;
                let mut sum_im = 0.0f32;
                for j in 0..nsym {
                    let entry = cs[(ks + j) * ntones + gray_map[digits[j]] as usize];
                    sum_re += entry.re.to_f32();
                    sum_im += entry.im.to_f32();
                }
                *s2_i = (sum_re * sum_re + sum_im * sum_im).sqrt();
            }

            let i_bit_base = chunk_bit_base + k * bps;
            for ib in 0..=ibmax {
                let bit_idx = i_bit_base + ib;
                if bit_idx >= codeword_len {
                    break;
                }
                let bit_sel = ibmax - ib;
                let max_one = s2
                    .iter()
                    .enumerate()
                    .filter(|&(i, _)| (i >> bit_sel) & 1 == 1)
                    .map(|(_, &v)| v)
                    .fold(f32::NEG_INFINITY, f32::max);
                let max_zero = s2
                    .iter()
                    .enumerate()
                    .filter(|&(i, _)| (i >> bit_sel) & 1 == 0)
                    .map(|(_, &v)| v)
                    .fold(f32::NEG_INFINITY, f32::max);
                let bm = max_one - max_zero;
                bmet_primary[bit_idx] = bm;
                if let Some(b) = bmet_norm_holder.as_deref_mut() {
                    let den = max_one.max(max_zero);
                    b[bit_idx] = if den > 0.0 { bm / den } else { 0.0 };
                }
            }

            k += nsym;
        }
        chunk_bit_base += chunk_len * bps;
    }
}

#[inline]
fn scale_bmet<T: LlrScalar>(mut v: Vec<f32>, scale: f32) -> Vec<T> {
    normalize_bmet(&mut v);
    v.into_iter().map(|x| T::from_f32(x * scale)).collect()
}

#[inline]
fn codeword_bit_len<P: Protocol>() -> usize {
    let bps = P::BITS_PER_SYMBOL as usize;
    data_chunks::<P>().iter().map(|&(_, l)| l).sum::<usize>() * bps
}

/// Generic LLR computation accepting any [`Cmplx<S>`] cs storage.
/// Inner `bmet` arithmetic stays in `f32` (norms / max-min are
/// awkward to quantise mid-stream) — `S` only changes how we read
/// each cs entry (`S::to_f32` per component). Final scale-and-round
/// to `T` happens at the bundle boundary.
pub fn compute_llr_generic<P: Protocol, S: SpecScalar, T: LlrScalar>(
    cs: &[Cmplx<S>],
    max_nsym: usize,
) -> LlrSet<T> {
    let codeword_len = codeword_bit_len::<P>();
    let mut bmeta = vec![0.0f32; codeword_len];
    let mut bmetb = vec![0.0f32; codeword_len];
    let mut bmetc = vec![0.0f32; codeword_len];
    let mut bmetd = vec![0.0f32; codeword_len];

    for nsym in 1usize..=max_nsym {
        let primary: &mut [f32] = match nsym {
            1 => &mut bmeta,
            2 => &mut bmetb,
            3 => &mut bmetc,
            _ => unreachable!(),
        };
        if nsym == 1 {
            // Split bmeta/bmetd borrow trick: both come from the same
            // function-local Vecs but we need disjoint &mut. The
            // explicit shadow keeps the borrow checker happy.
            let (bmeta_slice, bmetd_slice) = (&mut bmeta[..], &mut bmetd[..]);
            fill_bmet_for_nsym::<P, S>(cs, 1, bmeta_slice, Some(bmetd_slice));
        } else {
            fill_bmet_for_nsym::<P, S>(cs, nsym, primary, None);
        }
    }

    let s = P::LLR_SCALE;
    LlrSet {
        llra: scale_bmet::<T>(bmeta, s),
        llrb: scale_bmet::<T>(bmetb, s),
        llrc: scale_bmet::<T>(bmetc, s),
        llrd: scale_bmet::<T>(bmetd, s),
    }
}

/// Compute a single LLR variant by `nsym` level. Returns the bmet
/// vector matching that nsym (llra at 1, llrb at 2, llrc at 3),
/// normalised + scaled exactly the way [`compute_llr_generic`] would
/// produce that array on its own. nsym=1 returns the same llra
/// [`compute_llr_generic`] produces; the bit-normalised variant
/// (llrd) isn't separately exposed because the staircase consumer
/// already obtains llrd from Step 1's [`compute_llr_fast`] output.
///
/// Used by the FT8 stage-3 BP staircase to lazy-compute Step-2
/// variants — Step-1 already did the nsym=1 work, so Step 2 only
/// pays the (cheap) nsym=2 if variant b is tried, the (expensive)
/// nsym=3 only if variant c is needed. Empirically variant a from
/// Step 1 fails identically in Step 2 (same input, same BP), so
/// this path skips it — see `process_candidates_with` in
/// `ft8::decode_block`.
pub fn compute_llr_partial<P: Protocol, S: SpecScalar, T: LlrScalar>(
    cs: &[Cmplx<S>],
    nsym: usize,
) -> Vec<T> {
    debug_assert!((1..=3).contains(&nsym));
    let codeword_len = codeword_bit_len::<P>();
    let mut bmet = vec![0.0f32; codeword_len];
    fill_bmet_for_nsym::<P, S>(cs, nsym, &mut bmet, None);
    scale_bmet::<T>(bmet, P::LLR_SCALE)
}

// ──────────────────────────────────────────────────────────────────────────
// SNR estimation
// ──────────────────────────────────────────────────────────────────────────

/// WSJT-X compatible SNR (dB) estimate from symbol spectra + decoded tones.
///
/// Signal: `Σ |cs[k][itone[k]]|²`. Noise reference: `Σ |cs[k][(itone[k] +
/// NTONES/2) mod NTONES]|²` (tone on the "opposite side" of the comb).
/// SNR_dB = `10·log10(sig/noi − 1) − 27` clamped to −24 dB floor (WSJT-X
/// convention, applied per-tone bandwidth → 2500 Hz reference).
pub fn compute_snr_db<P: Protocol>(cs: &[Complex<f32>], itone: &[u8]) -> f32 {
    compute_snr_db_generic::<P, f32>(complex_slice_as_cmplx_f32(cs), itone)
}

/// Same as [`compute_snr_db`] but generic over the [`SpecScalar`]
/// type. The signal/noise sums use `S::Wide` accumulator and convert
/// to f32 at the boundary, so a `Cmplx<Q14i16>` cs gives a sane SNR
/// without intermediate f32 quantisation.
pub fn compute_snr_db_generic<P: Protocol, S: SpecScalar>(cs: &[Cmplx<S>], itone: &[u8]) -> f32 {
    let ntones = P::NTONES as usize;
    let n_sym = P::N_SYMBOLS as usize;
    let mut xsig = 0.0f32;
    let mut xnoi = 0.0f32;
    let offset = ntones / 2;
    for k in 0..n_sym.min(itone.len()) {
        let t = itone[k] as usize % ntones;
        xsig += cs[k * ntones + t].norm_sqr_f32();
        xnoi += cs[k * ntones + (t + offset) % ntones].norm_sqr_f32();
    }
    if xnoi < f32::EPSILON {
        return -24.0;
    }
    let ratio = xsig / xnoi - 1.0;
    if ratio <= 0.001 {
        return -24.0;
    }
    (10.0 * ratio.log10() - 27.0_f32).max(-24.0)
}

/// Hard-decision sync quality — count sync symbols whose dominant tone
/// matches the protocol's Costas pattern. Range is 0..N_SYNC; callers
/// typically threshold on this.
pub fn sync_quality<P: Protocol>(cs: &[Complex<f32>]) -> u32 {
    sync_quality_generic::<P, f32>(complex_slice_as_cmplx_f32(cs))
}

/// Generic version of [`sync_quality`]. Accepts any [`Cmplx<S>`]
/// slice; the per-tone "is this the dominant magnitude?" comparison
/// uses `S::Wide` (i32 for Q14i16) so no f32 round-trip is needed
/// on the embedded fixed-point path.
pub fn sync_quality_generic<P: Protocol, S: SpecScalar>(cs: &[Cmplx<S>]) -> u32
where
    S::Wide: PartialOrd,
{
    let ntones = P::NTONES as usize;
    let mut count = 0u32;
    for block in P::SYNC_MODE.blocks() {
        let start = block.start_symbol as usize;
        for (t, &expected) in block.pattern.iter().enumerate() {
            let sym = start + t;
            let best = (0..ntones)
                .max_by(|&a, &b| {
                    let na = cs[sym * ntones + a].norm_sqr_wide();
                    let nb = cs[sym * ntones + b].norm_sqr_wide();
                    na.partial_cmp(&nb).unwrap()
                })
                .unwrap_or(0);
            if best == expected as usize {
                count += 1;
            }
        }
    }
    count
}