// Reed-Solomon encoder and decoder in C++ using galois class libraries
// Implements a (255,223) RS encoder and errors-and-erasures decoder

// February 17, 1999
// Copyright 1999 Phil Karn
// May be used under the terms of the GNU Public License (GPL)

// User functions:
// void rs_init(void) - always call this first!

// void encode_rs(dtype data[KK], dtype bb[NN-KK]);
//  data = message data in, bb = parity symbols out

// int eras_dec_rs(dtype data[NN], int eras_pos[NN-KK], int no_eras);
//  data = codeword in
//  eras_pos = list of erasure locations, if any; may be NULL if no_eras == 0
//  no_eras = number of entries in eras_pos[]
//  returns number of corrected errors, or -1 if failure (too many errors)

// Still to do: support for nonconsecutive generator roots, e.g., like
// those used in the CCSDS standard

#include <stdio.h>
#include "galois.h"
#include "rs.h"

int KK(239); // Number of data symbols in block
const int B0(1);   // First consecutive root of generator polynomial



// RS generator polynomial
polynomial g;

// form Reed-Solomon generator polynomial g
// Must call this before calling encoder or decoder
void rs_init(void)
{
    int i;

    g = 1;
    for(i=B0;i<B0+NN-KK;i++){
        g *= X + (ALPHA^i);
    }
}

// RS encoder
//void encode_rs(dtype data[KK], dtype bb[NN-KK])
void encode_rs(dtype * data, dtype * bb)
{
    int i;
    polynomial message(254);
    polynomial parities;

    for(i=0;i<NN-KK;i++)
        message[i] = 0;
    for(;i<NN;i++)
        message[i] = data[NN-i-1];
    parities = message % g;
    for(i=0;i<NN-KK;i++)
        bb[i] = parities[NN-KK-i-1].poly();
}

// Find roots of polynomial by Chien search (a fancy term for
// just trying all possible values!)
static int roots(polynomial& poly,galois results[])
{
    int i;
    int rootcount = 0;

    for(i=0;i<=NN;i++){
        if(poly(i) == 0){
            results[rootcount++] = i;
            if(rootcount == poly.degree())
                break; // found them all, no need to keep going
        }
    }
    return rootcount;
}

// Massey-Berlekamp algorithm
// Find minimal length shift register that will generate the sequence s[]
static polynomial massey
(
    polynomial& s, // syndrome polynomial
    const int no_eras, // number of elements in erasure locations
    const int epos[] // list of erasure locations
){
    int r,i;
    polynomial lambda,t,b;
    galois discr_r;

    lambda = 1;
    for(i=0;i<no_eras;i++)
        lambda *= X + galois_index(NN-1-epos[i]);

    b = lambda;

    for(r=no_eras;r < NN-KK;r++){
        // compute r-th discrepancy
        discr_r = 0;
        for(i=0;i<=lambda.degree();i++)
            discr_r += lambda[i] * s[r-i];
        if(discr_r == 0){
            b <<= 1;
        } else {
            t = lambda - (discr_r * (b << 1));
            if(2 * lambda.degree() <= r + 1 + no_eras){
                b = lambda / discr_r;
            } else {
                b <<= 1;
            }
            lambda = t;
        }
    }
    return lambda;
}

// Errors and erasures RS decoder
//int eras_dec_rs(dtype data[NN], int eras_pos[NN-KK], int no_eras)
int eras_dec_rs(dtype * data, int * eras_pos, int no_eras)
{
    int i;
    polynomial received(NN-1);
    for(i=0;i<NN;i++)
        received[NN-i-1] = data[i];

    // compute syndromes
    polynomial s(NN-KK-1);
    int errflag = 0;
    for(i=0;i<NN-KK;i++){
        s[i] = received(galois_index(B0+i));
        errflag |= s[i].poly();
    }
    if(!errflag)
        return 0; // no errors, no need to proceed

    // find error+erasure locator polynomial
    polynomial lfsr;
    lfsr = massey(s,no_eras,eras_pos);

    // find roots of error locator polynomial
    galois lfsr_roots[256];
    int lfsr_rootcount;
    lfsr_rootcount = roots(lfsr,lfsr_roots);
    if(lfsr_rootcount != lfsr.degree())
        return -1; // unrecoverable error

    // error evaluator polynomial
    polynomial z;
    z = s * lfsr;
    z.set_degree(NN-KK-1); // modulo X^(NN-KK)

    // find error values
    for(i=0;i<lfsr_rootcount;i++){
        galois num,den;
        int j;

        num = z(lfsr_roots[i]) * (lfsr_roots[i] ^ (B0-1));
        // evaluate formal derivative of lfsr at each lfsr_root
        den = 0;
        for(j=min(lfsr.degree()-1,NN-KK-1) & ~1; j>=0; j-= 2){
            den += lfsr[j+1] * (lfsr_roots[i]^j);
        }
        if(den == 0)
            return -1;
        received[lfsr_roots[i].inv().index()] -=  num/den; // apply correction
    }


    // JR: overwrite with corrected data
    for(i=0;i<NN;i++)
        data[i]=received[NN-i-1].poly();

    return lfsr_rootcount;
}