orion-sdr 0.0.27

// src/codec/psk31_conv.rs
//
// Rate-1/2, constraint-length-5 convolutional encoder and Viterbi decoder
// for QPSK31.
//
// Generator polynomials (octal): G0 = 25, G1 = 23.
//   G0 = 0b10101 — taps at bits {0, 2, 4} of the 5-bit shift register
//   G1 = 0b10011 — taps at bits {0, 1, 4}
//
// Shift register convention:
//   sr is a 4-bit register holding the K-1 = 4 past input bits.
//   For each new input bit b, the 5-bit encoder window is:
//     window = (b << 4) | sr      (b is the newest bit, sr[3] is oldest)
//   After encoding:
//     sr = (sr >> 1) | (b << 3)   (shift right, insert b at MSB of 4-bit sr)
//
// Interleaving: for each input bit the encoder emits [g0_bit, g1_bit].
// Output length = 2 × input.len().
//
// Viterbi decoder:
//   16 states (2^(K-1) = 2^4 = 16).
//   `decisions[t][s]` stores the previous state that reached state s at time t
//   with the minimum accumulated metric.
//   Traceback follows the chain of previous states from the best final state.

// ── Encoder ───────────────────────────────────────────────────────────────────

/// Parity of all set bits in `x` — returns 0 or 1.
#[inline(always)]
fn parity(x: u8) -> u8 {
    let x = x ^ (x >> 4);
    let x = x ^ (x >> 2);
    (x ^ (x >> 1)) & 1
}

/// Rate-1/2, K=5 convolutional encoder.
///
/// G0 = 0b10101 (octal 25), G1 = 0b10011 (octal 23).
///
/// Input: slice of bits (each byte = 0 or 1).
/// Output: interleaved coded bits `[g0_0, g1_0, g0_1, g1_1, …]`; length = 2 × input.len().
pub fn conv_encode(bits: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(bits.len() * 2);
    let mut sr: u8 = 0;
    for &b in bits {
        let window = ((b & 1) << 4) | (sr & 0x0F);
        out.push(parity(window & 0b10101)); // G0
        out.push(parity(window & 0b10011)); // G1
        sr = (sr >> 1) | ((b & 1) << 3);
    }
    out
}

// ── Viterbi decoder ───────────────────────────────────────────────────────────

const NUM_STATES: usize = 16;

/// Compute the two coded bits produced when state `s` receives input `b`.
#[inline]
fn branch_bits(s: u8, b: u8) -> (u8, u8) {
    let window = ((b & 1) << 4) | (s & 0x0F);
    (parity(window & 0b10101), parity(window & 0b10011))
}

/// Next state when current state is `s` and input `b` is received.
#[inline]
fn next_state(s: u8, b: u8) -> u8 {
    (s >> 1) | ((b & 1) << 3)
}

// DQPSK expected phasors indexed by dibit = c0*2 + c1.
// Matches the QPSK31_PHASE_STEP table in src/modulate/psk31.rs:
//   dibit 0 (c0=0,c1=0): step = (+1,  0)
//   dibit 1 (c0=0,c1=1): step = ( 0, -1)
//   dibit 2 (c0=1,c1=0): step = ( 0, +1)
//   dibit 3 (c0=1,c1=1): step = (-1,  0)
pub const DQPSK_EXP: [(f32, f32); 4] = [
    ( 1.0,  0.0), // dibit 0
    ( 0.0, -1.0), // dibit 1
    ( 0.0,  1.0), // dibit 2
    (-1.0,  0.0), // dibit 3
];

/// Soft Viterbi decoder for the rate-1/2, K=5 code.
///
/// `soft` is an interleaved slice `[re_0, im_0, re_1, im_1, …]` of DQPSK
/// differential-detection outputs.  The branch metric uses the actual DQPSK
/// constellation phasors as expected values, matching the modulator's
/// `QPSK31_PHASE_STEP` table.
///
/// Returns a decoded bit slice of length `soft.len() / 2`.
pub fn viterbi_decode(soft: &[f32]) -> Vec<u8> {
    let n_syms = soft.len() / 2;
    if n_syms == 0 {
        return Vec::new();
    }

    // Path metrics — lower is better (minimising Euclidean distance).
    let inf = f32::MAX / 2.0;
    let mut pm = [inf; NUM_STATES];
    pm[0] = 0.0;

    // decisions[t][s] = previous state that led to state `s` at time `t`.
    // Stored as u8 (state index 0–15).
    let mut prev_state_table: Vec<[u8; NUM_STATES]> = vec![[0u8; NUM_STATES]; n_syms];

    for t in 0..n_syms {
        let s0 = soft[t * 2];
        let s1 = soft[t * 2 + 1];
        let mut new_pm = [inf; NUM_STATES];

        for prev in 0..NUM_STATES {
            if pm[prev] >= inf { continue; }
            for &bit in &[0u8, 1u8] {
                let (c0, c1) = branch_bits(prev as u8, bit);
                // Expected soft values are the DQPSK phasor for this dibit,
                // not ±1 on both axes — the DQPSK constellation places all
                // energy on a single axis per symbol.
                let dibit = (c0 & 1) * 2 + (c1 & 1);
                let (exp0, exp1) = DQPSK_EXP[dibit as usize];
                let bm = (s0 - exp0) * (s0 - exp0) + (s1 - exp1) * (s1 - exp1);
                let ns = next_state(prev as u8, bit) as usize;
                let cand = pm[prev] + bm;
                if cand < new_pm[ns] {
                    new_pm[ns] = cand;
                    prev_state_table[t][ns] = prev as u8;
                }
            }
        }
        pm = new_pm;
    }

    // Best final state.
    let mut state = pm
        .iter()
        .enumerate()
        .min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i)
        .unwrap_or(0);

    // Traceback: recover the sequence of decoded bits.
    let mut bits_out = vec![0u8; n_syms];
    for t in (0..n_syms).rev() {
        let prev = prev_state_table[t][state] as usize;
        // The input bit that caused the transition prev → state:
        // next_state(prev, b) == state
        // => (prev >> 1) | (b << 3) == state
        // => b = (state >> 3) & 1  (MSB of state came from input bit b)
        let b = (state >> 3) as u8 & 1;
        bits_out[t] = b;
        state = prev;
    }

    bits_out
}

/// Coherent soft Viterbi decoder for the rate-1/2, K=5 code.
///
/// `soft` is an interleaved slice `[re_0, im_0, re_1, im_1, …]` of phase-
/// corrected absolute symbol estimates from the DFM.  Each state tracks a
/// hypothesised absolute phasor; the branch metric is the squared Euclidean
/// distance between the received symbol and the hypothesis, eliminating the
/// ~3 dB noise-product penalty of differential detection.
///
/// `phase_steps` is the DQPSK phase-step table indexed by dibit = c0*2+c1,
/// matching the modulator's `QPSK31_PHASE_STEP`.
///
/// Returns a decoded bit slice of length `soft.len() / 2`.
pub fn viterbi_decode_coherent(soft: &[f32], phase_steps: &[(f32, f32); 4]) -> Vec<u8> {
    let n_syms = soft.len() / 2;
    if n_syms == 0 {
        return Vec::new();
    }

    let inf = f32::MAX / 2.0;
    let mut pm = [inf; NUM_STATES];
    pm[0] = 0.0;
    // Hypothesised absolute phasor per state.  Initial phasor (1,0) matches
    // Qpsk31Mod starting phase.
    let mut hyp = [(1.0f32, 0.0f32); NUM_STATES];

    let mut prev_state_table: Vec<[u8; NUM_STATES]> = vec![[0u8; NUM_STATES]; n_syms];
    // Store hypothesised phasor that won each transition, for propagation.
    let mut hyp_table: Vec<[(f32, f32); NUM_STATES]> = vec![[(0.0, 0.0); NUM_STATES]; n_syms];

    for t in 0..n_syms {
        let s_re = soft[t * 2];
        let s_im = soft[t * 2 + 1];
        let mut new_pm  = [inf; NUM_STATES];
        let mut new_hyp = [(0.0f32, 0.0f32); NUM_STATES];

        for prev in 0..NUM_STATES {
            if pm[prev] >= inf { continue; }
            let (h_re, h_im) = hyp[prev];
            for &bit in &[0u8, 1u8] {
                let (c0, c1) = branch_bits(prev as u8, bit);
                let dibit = (c0 & 1) * 2 + (c1 & 1);
                let (step_re, step_im) = phase_steps[dibit as usize];
                // Hypothesised phasor after this transition: h * step.
                let nh_re = h_re * step_re - h_im * step_im;
                let nh_im = h_im * step_re + h_re * step_im;
                // Coherent branch metric: |sym_c - hyp|².
                let bm = (s_re - nh_re) * (s_re - nh_re)
                       + (s_im - nh_im) * (s_im - nh_im);
                let ns = next_state(prev as u8, bit) as usize;
                let cand = pm[prev] + bm;
                if cand < new_pm[ns] {
                    new_pm[ns] = cand;
                    new_hyp[ns] = (nh_re, nh_im);
                    prev_state_table[t][ns] = prev as u8;
                    hyp_table[t][ns] = (nh_re, nh_im);
                }
            }
        }
        pm  = new_pm;
        hyp = new_hyp;
    }

    // Best final state.
    let mut state = pm
        .iter()
        .enumerate()
        .min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i)
        .unwrap_or(0);

    // Traceback.
    let mut bits_out = vec![0u8; n_syms];
    for t in (0..n_syms).rev() {
        let prev = prev_state_table[t][state] as usize;
        let b = (state >> 3) as u8 & 1;
        bits_out[t] = b;
        state = prev;
    }

    bits_out
}

// ── Streaming coherent Viterbi decoder ────────────────────────────────────────

/// Fixed-lag sliding-window Viterbi decoder for coherent QPSK31.
///
/// Processes one QPSK symbol at a time and emits decoded bits with a fixed
/// latency of `TRACEBACK_DEPTH` symbols.  Uses the same coherent branch
/// metric as `viterbi_decode_coherent` (hypothesised absolute phasor per
/// trellis state).
///
/// Based on the fldigi approach: after each ACS step, if enough history has
/// accumulated, traceback from the best-metric state and emit the oldest
/// undecoded bit.
pub struct StreamingViterbi {
    pm:      [f32; NUM_STATES],
    history: Vec<[u8; NUM_STATES]>,  // circular buffer of prev_state
    ptr:     usize,                  // write pointer into circular buffer
    count:   usize,                  // total symbols processed
    phase_steps: [(f32, f32); 4],
}

/// Traceback depth.  Textbook = 5×(K-1) = 20 for rate-1/2 K=5.
/// Use 32 for extra convergence margin with differential detection.
const TRACEBACK_DEPTH: usize = 32;
/// Circular buffer size (must be > TRACEBACK_DEPTH).
const PATHMEM: usize = 128;

impl StreamingViterbi {
    /// Create a new streaming Viterbi decoder.
    /// `phase_steps` is the DQPSK phase-step table (same as for batch decoder).
    pub fn new(phase_steps: &[(f32, f32); 4]) -> Self {
        let inf = f32::MAX / 2.0;
        let mut pm = [inf; NUM_STATES];
        pm[0] = 0.0;
        Self {
            pm,
            history: vec![[0u8; NUM_STATES]; PATHMEM],
            ptr:     0,
            count:   0,
            phase_steps: *phase_steps,
        }
    }

    /// Feed one QPSK symbol.
    ///
    /// `s_re, s_im` are the DQPSK differential-detection outputs (or
    /// phase-corrected absolute phasors from the coherent demod).  The branch
    /// metric compares the received symbol against the expected DQPSK step
    /// phasors — this is a non-coherent metric that doesn't require tracking
    /// absolute phase, making it robust to PLL lock-in transients.
    pub fn feed_symbol(&mut self, s_re: f32, s_im: f32) -> Option<u8> {
        let inf = f32::MAX / 2.0;
        let mut new_pm = [inf; NUM_STATES];

        // ACS step — non-coherent DQPSK branch metric.
        for prev in 0..NUM_STATES {
            if self.pm[prev] >= inf { continue; }
            for &bit in &[0u8, 1u8] {
                let (c0, c1) = branch_bits(prev as u8, bit);
                let dibit = (c0 & 1) * 2 + (c1 & 1);
                let (exp_re, exp_im) = self.phase_steps[dibit as usize];
                let bm = (s_re - exp_re) * (s_re - exp_re)
                       + (s_im - exp_im) * (s_im - exp_im);
                let ns = next_state(prev as u8, bit) as usize;
                let cand = self.pm[prev] + bm;
                if cand < new_pm[ns] {
                    new_pm[ns] = cand;
                    self.history[self.ptr][ns] = prev as u8;
                }
            }
        }
        self.pm = new_pm;

        // Periodic metric normalisation to prevent f32 overflow.
        if self.count % 256 == 255 {
            let min_pm = self.pm.iter().copied()
                .filter(|&v| v < inf)
                .fold(inf, f32::min);
            if min_pm > 0.0 {
                for p in &mut self.pm {
                    if *p < inf { *p -= min_pm; }
                }
            }
        }

        self.ptr = (self.ptr + 1) % PATHMEM;
        self.count += 1;

        // Not enough history for traceback yet.
        if self.count <= TRACEBACK_DEPTH {
            return None;
        }

        // Fixed-lag traceback: find best state now, trace back TRACEBACK_DEPTH
        // steps, emit the decoded bit at the traceback endpoint.
        let mut state = self.pm
            .iter()
            .enumerate()
            .min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);

        let mut p = (self.ptr + PATHMEM - 1) % PATHMEM; // current position
        for _ in 0..TRACEBACK_DEPTH {
            state = self.history[p][state] as usize;
            p = (p + PATHMEM - 1) % PATHMEM;
        }

        // The decoded bit is the MSB of the state at the traceback endpoint
        // (same as batch: b = (state >> 3) & 1).
        Some(((state >> 3) & 1) as u8)
    }

    /// Flush remaining bits after the last symbol.  Returns up to
    /// `TRACEBACK_DEPTH` final decoded bits by tracing back from the best
    /// state at progressively shorter depths.
    pub fn flush(&mut self) -> Vec<u8> {
        let mut out = Vec::new();
        // Feed TRACEBACK_DEPTH zero-energy symbols to push out the tail.
        for _ in 0..TRACEBACK_DEPTH {
            if let Some(b) = self.feed_symbol(0.0, 0.0) {
                out.push(b);
            }
        }
        out
    }
}

/// Hard-decision Viterbi decoder (for testing with noiseless hard bits).
///
/// Input: interleaved coded bits `[c0_0, c1_0, c0_1, c1_1, …]`.
/// Converts each (c0, c1) pair to the corresponding DQPSK phasor before
/// calling the soft decoder, matching the branch metric convention.
pub fn viterbi_decode_hard(bits: &[u8]) -> Vec<u8> {
    let n_syms = bits.len() / 2;
    let mut soft = Vec::with_capacity(n_syms * 2);
    for i in 0..n_syms {
        let c0 = bits[i * 2] & 1;
        let c1 = bits[i * 2 + 1] & 1;
        let dibit = c0 * 2 + c1;
        let (re, im) = DQPSK_EXP[dibit as usize];
        soft.push(re);
        soft.push(im);
    }
    viterbi_decode(&soft)
}