sdrangel/plugins/channelrx/demoddatv/leansdr/viterbi.h

#ifndef LEANSDR_VITERBI_H
#define LEANSDR_VITERBI_H

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// This is a generic implementation of Viterbi with explicit
// representation of the trellis.  There is special support for
// convolutional coding, but the code can handle other schemes.
// TBD This is very inefficient. For a specific trellis all loops
// can be be unrolled.

namespace leansdr
{

// TS is an integer type for a least NSTATES+1 states.
// NSTATES is the number of states (e.g. 2^(K-1)).
// TUS is an integer type for uncoded symbols (branch identifiers).
// NUS is the number of uncoded symbols.
// TCS is an integer type for coded symbols (branch labels).
// NCS is the number of coded symbols.
// TP is a type for representing paths.
// TPM, TBM are unsigned integer types for path/branch metrics.
// TPM is at least as wide as TBM.

template<typename TS, int NSTATES, typename TUS, int NUS, int NCS>
struct trellis
{
    static const int NOSTATE = NSTATES + 1;

    struct state
    {
        struct branch
        {
            TS pred;         // Predecessor state or NOSTATE
            TUS us;          // Uncoded symbol
        } branches[NCS];   // Incoming branches indexed by coded symbol
    } states[NSTATES];

    trellis()
    {
        for (TS s = 0; s < NSTATES; ++s)
            for (int cs = 0; cs < NCS; ++cs)
                states[s].branches[cs].pred = NOSTATE;
    }

    // TBD Polynomial width should be a template parameter ?
    void init_convolutional(const uint16_t G[])
    {
        if (NCS & (NCS - 1))
        {
            fprintf(stderr, "leansdr::trellis::init_convolutional: NCS must be a power of 2\n");
            return;
        }
        // Derive number of polynomials from NCS.
        int nG = log2i(NCS);

        for (TS s = 0; s < NSTATES; ++s)
        {
            for (TUS us = 0; us < NUS; ++us)
            {
                // Run the convolutional encoder from state s with input us
                uint64_t shiftreg = s;  // TBD type
                // Reverse bits
                TUS us_rev = 0;
                for (int b = 1; b < NUS; b *= 2)
                    if (us & b)
                        us_rev |= (NUS / 2 / b);
                shiftreg |= us_rev * NSTATES;
                uint32_t cs = 0;  // TBD type
                for (int g = 0; g < nG; ++g)
                    cs = (cs << 1) | parity(shiftreg & G[g]);
                shiftreg /= NUS;  // Shift bits for 1 uncoded symbol
                // [us] at state [s] emits [cs] and leads to state [shiftreg].
                typename state::branch *b = &states[shiftreg].branches[cs];
                if (b->pred != NOSTATE)
                {
                    fprintf(stderr, "leansdr::trellis::init_convolutional: Invalid convolutional code\n");
                    return;
                }
                b->pred = s;
                b->us = us;
            }
        }

    }

    void dump()
    {
        for (int s = 0; s < NSTATES; ++s)
        {
            fprintf(stderr, "State %02x:", s);
            for (int cs = 0; cs < NCS; ++cs)
            {
                typename state::branch *b = &states[s].branches[cs];
                if (b->pred == NOSTATE)
                    fprintf(stderr, "     - ");
                else
                    fprintf(stderr, "   %02x+%x", b->pred, b->us);
            }
            fprintf(stderr, "\n");
        }
    }

};

// Interface that hides the templated internals.
template<typename TUS, typename TCS, typename TBM, typename TPM>
struct viterbi_dec_interface
{
    virtual TUS update(TBM costs[], TPM *quality = NULL)=0;
    virtual TUS update(TCS s, TBM cost, TPM *quality = NULL)=0;
    virtual TUS update(int nm, TCS cs[], TBM costs[], TPM *quality = NULL)=0;
};

template<typename TS, int NSTATES, typename TUS, int NUS, typename TCS, int NCS, typename TBM, typename TPM, typename TP>
struct viterbi_dec: viterbi_dec_interface<TUS, TCS, TBM, TPM>
{

    trellis<TS, NSTATES, TUS, NUS, NCS> *trell;

    struct state
    {
        TPM cost;    // Metric of best path leading to this state
        TP path;     // Best path leading to this state
    };
    typedef state statebank[NSTATES];
    state statebanks[2][NSTATES];
    statebank *states, *newstates;  // Alternate between banks

    viterbi_dec(trellis<TS, NSTATES, TUS, NUS, NCS> *_trellis) :
            trell(_trellis)
    {
        states = &statebanks[0];
        newstates = &statebanks[1];
        for (TS s = 0; s < NSTATES; ++s)
            (*states)[s].cost = 0;
        // Determine max value that can fit in TPM
        max_tpm = (TPM) 0 - 1;
        if (max_tpm < 0)
        {
            // TPM is signed
            for (max_tpm = 0; max_tpm * 2 + 1 > max_tpm; max_tpm = max_tpm * 2 + 1)
                ;
        }
    }

    // Update with full metric

    TUS update(TBM costs[NCS], TPM *quality = NULL)
    {
        TPM best_tpm = max_tpm, best2_tpm = max_tpm;
        TS best_state = 0;
        // Update all states
        for (int s = 0; s < NSTATES; ++s)
        {
            TPM best_m = max_tpm;
            typename trellis<TS, NSTATES, TUS, NUS, NCS>::state::branch *best_b = NULL;
            // Select best branch
            for (int cs = 0; cs < NCS; ++cs)
            {
                typename trellis<TS, NSTATES, TUS, NUS, NCS>::state::branch *b = &trell->states[s].branches[cs];
                if (b->pred == trell->NOSTATE)
                    continue;
                TPM m = (*states)[b->pred].cost + costs[cs];
                if (m <= best_m)
                {  // <= guarantees one match
                    best_m = m;
                    best_b = b;
                }
            }
            (*newstates)[s].path = (*states)[best_b->pred].path;
            (*newstates)[s].path.append(best_b->us);
            (*newstates)[s].cost = best_m;
            // Select best and second-best states
            if (best_m < best_tpm)
            {
                best_state = s;
                best2_tpm = best_tpm;
                best_tpm = best_m;
            }
            else if (best_m < best2_tpm)
                best2_tpm = best_m;
        }
        // Swap banks
        {
            statebank *tmp = states;
            states = newstates;
            newstates = tmp;
        }
        // Prevent overflow of path metrics
        for (TS s = 0; s < NSTATES; ++s)
            (*states)[s].cost -= best_tpm;
#if 0
        // Observe that the min-max range remains bounded
        fprintf(stderr,"-%2d = [", best_tpm);
        for ( TS s=0; s<NSTATES; ++s ) fprintf(stderr," %d", (*states)[s].cost);
        fprintf(stderr," ]\n");
#endif
        // Return difference between best and second-best as quality metric.
        if (quality)
            *quality = best2_tpm - best_tpm;
        // Return uncoded symbol of best path
        return (*states)[best_state].path.read();
    }

    // Update with partial metrics.
    // The costs provided must be negative.
    // The other symbols will be assigned a cost of 0.

    TUS update(int nm, TCS cs[], TBM costs[], TPM *quality = NULL)
    {
        TPM best_tpm = max_tpm, best2_tpm = max_tpm;
        TS best_state = 0;
        // Update all states
        for (int s = 0; s < NSTATES; ++s)
        {
            // Select best branch among those for with metrics are provided
            TPM best_m = max_tpm;
            typename trellis<TS, NSTATES, TUS, NUS, NCS>::state::branch *best_b = NULL;
            for (int im = 0; im < nm; ++im)
            {
                typename trellis<TS, NSTATES, TUS, NUS, NCS>::state::branch *b = &trell->states[s].branches[cs[im]];
                if (b->pred == trell->NOSTATE)
                    continue;
                TPM m = (*states)[b->pred].cost + costs[im];
                if (m <= best_m)
                {  // <= guarantees one match
                    best_m = m;
                    best_b = b;
                }
            }
            if (nm != NCS)
            {
                // Also scan the other branches.
                // We actually rescan the branches with metrics.
                // This works because costs are negative.
                for (int cs = 0; cs < NCS; ++cs)
                {
                    typename trellis<TS, NSTATES, TUS, NUS, NCS>::state::branch *b = &trell->states[s].branches[cs];
                    if (b->pred == trell->NOSTATE)
                        continue;
                    TPM m = (*states)[b->pred].cost;
                    if (m <= best_m)
                    {
                        best_m = m;
                        best_b = b;
                    }
                }
            }
            (*newstates)[s].path = (*states)[best_b->pred].path;
            (*newstates)[s].path.append(best_b->us);
            (*newstates)[s].cost = best_m;
            // Select best states
            if (best_m < best_tpm)
            {
                best_state = s;
                best2_tpm = best_tpm;
                best_tpm = best_m;
            }
            else if (best_m < best2_tpm)
                best2_tpm = best_m;
        }
        // Swap banks
        {
            statebank *tmp = states;
            states = newstates;
            newstates = tmp;
        }
        // Prevent overflow of path metrics
        for (TS s = 0; s < NSTATES; ++s)
            (*states)[s].cost -= best_tpm;
#if 0
        // Observe that the min-max range remains bounded
        fprintf(stderr,"-%2d = [", best_tpm);
        for ( TS s=0; s<NSTATES; ++s ) fprintf(stderr," %d", (*states)[s].cost);
        fprintf(stderr," ]\n");
#endif
        // Return difference between best and second-best as quality metric.
        if (quality)
            *quality = best2_tpm - best_tpm;
        // Return uncoded symbol of best path
        return (*states)[best_state].path.read();
    }

    // Update with single-symbol metric.
    // cost must be negative.

    TUS update(TCS cs, TBM cost, TPM *quality = NULL)
    {
        return update(1, &cs, &cost, quality);
    }

    void dump()
    {
        fprintf(stderr, "[");
        for (TS s = 0; s < NSTATES; ++s)
            if (states[s].cost)
                fprintf(stderr, " %02x:%d", s, states[s].cost);
        fprintf(stderr, "\n");
    }

private:
    TPM max_tpm;
};

// Paths (sequences of uncoded symbols) represented as bitstreams.
// NBITS is the number of bits per symbol.
// DEPTH is the number of symbols stored in the path.
// T is an unsigned integer type wider than NBITS*DEPTH.

template<typename T, typename TUS, int NBITS, int DEPTH>
struct bitpath
{
    T val;
    bitpath() :
            val(0)
    {
    }
    void append(TUS us)
    {
        val = (val << NBITS) | us;
    }
    TUS read()
    {
        return (val >> (DEPTH - 1) * NBITS) & ((1 << NBITS) - 1);
    }
};

}  // namespace

#endif  // LEANSDR_VITERBI_H