#ifndef LEANSDR_RS_H
#define LEANSDR_RS_H

#include "leansdr/math.h"

#define DEBUG_RS 0

namespace leansdr {

  // Finite group GF(2^N).

  // GF(2) is the ring ({0,1},+,*).
  // GF(2)[X] is the ring of polynomials with coefficients in GF(2).
  // P(X) is an irreducible polynomial of GF(2)[X].
  // N is the degree of P(x).
  // P is the bitfield representation of P(X), with degree 0 at LSB.
  // GF(2)[X]/(P) is GF(2)[X] modulo P(X).
  // (GF(2)[X]/(P), +) is a group with 2^N elements.
  // (GF(2)[X]/(P)*, *) is a group with 2^N-1 elements.
  // (GF(2)[X]/(P), +, *) is a field with 2^N elements, noted GF(2^N).
  // Te is a C++ integer type for representing elements of GF(2^N).
  //   "0" is 0
  //   "1" is 1
  //   "2" is X
  //   "3" is X+1
  //   "4" is X^2
  // Tp is a C++ integer type for representing P(X) (1 bit larger than Te).
  // ALPHA is a primitive element of GF(2^N). Usually "2"=[X] is chosen.

  template<typename Te, typename Tp, Tp P, int N, Te ALPHA>
  struct gf2x_p {
    gf2x_p() {
      if ( ALPHA != 2 ) fail("alpha!=2 not implemented");
      // Precompute log and exp tables.
      Tp alpha_i = 1;
      for ( Tp i=0; i<(1<<N); ++i ) {
	lut_exp[i] = alpha_i;
	lut_exp[((1<<N)-1)+i] = alpha_i;
	lut_log[alpha_i] = i;      
	alpha_i <<= 1;  // Multiply by alpha=[X] i.e. increase degrees
	if ( alpha_i & (1<<N) ) alpha_i ^= P;  // Modulo P iteratively
      }
    }
    static const Te alpha = ALPHA;
    inline Te add(Te x, Te y) { return x ^ y; }  // Addition modulo 2
    inline Te sub(Te x, Te y) { return x ^ y; }  // Subtraction modulo 2
    inline Te mul(Te x, Te y) {
      if ( !x || !y ) return 0;
      return lut_exp[lut_log[x] + lut_log[y]];
    }
    inline Te div(Te x, Te y) {
      //if ( ! y ) fail("div"); // TODO
      if ( ! x ) return 0;
      return lut_exp[lut_log[x] + ((1<<N)-1) - lut_log[y]];
    }
    inline Te inv(Te x) {
      //    if ( ! x ) fail("inv");
      return lut_exp[((1<<N)-1) - lut_log[x]];
    }
    inline Te exp(Te x) { return lut_exp[x]; }
    inline Te log(Te x) { return lut_log[x]; }
  private:
    Te lut_exp[(1<<N)*2];  // Wrap to avoid indexing modulo 2^N-1 
    Te lut_log[1<<N];
  };

  // Reed-Solomon for RS(204,188) shortened from RS(255,239).

  struct rs_engine {
    // EN 300 421, section 4.4.2, Field Generator Polynomial
    // p(X) = X^8 + X^4 + X^3 + X^2 + 1
    gf2x_p<unsigned char, unsigned short, 0x11d, 8, 2> gf;

    u8 G[17];  // { G_16, ..., G_0 }

    rs_engine() {
      // EN 300 421, section 4.4.2, Code Generator Polynomial
      // G(X) = (X-alpha^0)*...*(X-alpha^15)
      for ( int i=0; i<=16; ++i ) G[i] = (i==16) ? 1 : 0;  // Init G=1
      for ( int d=0; d<16; ++d ) {
	// Multiply by (X-alpha^d)
	// G := X*G - alpha^d*G
	for ( int i=0; i<=16; ++i )
	  G[i] = gf.sub((i==16)?0:G[i+1], gf.mul(gf.exp(d),G[i]));
      }
#if DEBUG_RS
      fprintf(stderr, "RS generator:");
      for ( int i=0; i<=16; ++i ) fprintf(stderr, " %02x", G[i]);
      fprintf(stderr, "\n");
#endif
    }

    // RS-encoded messages are interpreted as coefficients in
    // GF(256) of a polynomial P.
    // The syndromes are synd[i] = P(alpha^i).
    // By convention coefficients are listed by decreasing degree here,
    // so we can evaluate syndromes of the shortened code without
    // prepending with 51 zeroes.
    bool syndromes(const u8 *poly, u8 *synd) {
      bool corrupted = false;
      for ( int i=0; i<16; ++i ) {
	synd[i] = eval_poly_rev(poly, 204, gf.exp(i));
	if ( synd[i] ) corrupted = true;
      }
      return corrupted;
    }
    u8 eval_poly_rev(const u8 *poly, int n, u8 x) {
      // poly[0]*x^(n-1) + .. + poly[n-1]*x^0 with Hörner method.
      u8 acc = 0;
      for ( int i=0; i<n; ++i ) acc = gf.add(gf.mul(acc,x), poly[i]);
      return acc;
    }

    // Evaluation with coefficients listed by increasing degree.
    u8 eval_poly(const u8 *poly, int deg, u8 x) {
      // poly[0]*x^0 + .. + poly[deg]*x^deg with Hörner method.
      u8 acc = 0;
      for ( ; deg>=0; --deg ) acc = gf.add(gf.mul(acc,x), poly[deg]);
      return acc;
    }

    // Append parity symbols

    void encode(u8 msg[204]) {
      // TBD Avoid copying
      u8 p[204];
      memcpy(p, msg, 188);
      memset(p+188, 0, 16);
      // p = msg*X^16
#if DEBUG_RS
      fprintf(stderr, "uncoded:");
      for ( int i=0; i<204; ++i ) fprintf(stderr, " %d", p[i]);
      fprintf(stderr, "\n");
#endif
      // Compute remainder modulo G
      for ( int d=0; d<188; ++d ) {
	// Clear monomial of degree d
	if ( ! p[d] ) continue;
	u8 k = gf.div(p[d], G[0]);
	// p(X) := p(X) - k*G(X)*X^(188-d)
	for ( int i=0; i<=16; ++i )
	  p[d+i] = gf.sub(p[d+i], gf.mul(k,G[i]));
      }
#if DEBUG_RS
      fprintf(stderr, "coded:");
      for ( int i=0; i<204; ++i ) fprintf(stderr, " %d", p[i]);
      fprintf(stderr, "\n");
#endif
      memcpy(msg+188, p+188, 16);
    }

    // Try to fix errors in pout[].
    // If pin[] is provided, errors will be fixed in the original
    // message too and syndromes will be updated.

    bool correct(u8 synd[16], u8 pout[188],
		 u8 pin[204]=NULL, int *bits_corrected=NULL) {
      // Berlekamp - Massey
      // http://en.wikipedia.org/wiki/Berlekamp%E2%80%93Massey_algorithm#Code_sample
      u8 C[16] = { 1,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0 };  // Max degree is L
      u8 B[16] = { 1,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0 };
      int L = 0;
      int m = 1;
      u8 b = 1;
      for ( int n=0; n<16; ++n ) {
	u8 d = synd[n];
	for ( int i=1; i<=L; ++i ) d ^= gf.mul(C[i], synd[n-i]);
	if ( ! d ) {
	  ++m;
	} else if ( 2*L <= n ) {
	  u8 T[16];
	  memcpy(T, C, sizeof(T));
	  for ( int i=0; i<16-m; ++i )
	    C[m+i] ^= gf.mul(d, gf.mul(gf.inv(b),B[i]));
	  L = n + 1 - L;
	  memcpy(B, T, sizeof(B));
	  b = d;
	  m = 1;
	} else {
	  for ( int i=0; i<16-m; ++i )
	    C[m+i] ^= gf.mul(d, gf.mul(gf.inv(b),B[i]));
	  ++m;
	}
      }
      // L is the number of errors
      // C of degree L is now the error locator polynomial (Lambda)
#if DEBUG_RS
      fprintf(stderr, "[L=%d  C=",L);
      for ( int i=0; i<16; ++i ) fprintf(stderr, " %d", C[i]);
      fprintf(stderr, "]\n");
      fprintf(stderr, "[S=");
      for ( int i=0; i<16; ++i ) fprintf(stderr, " %d", synd[i]);
      fprintf(stderr, "]\n");
#endif
    
      // Forney
      // http://en.wikipedia.org/wiki/Forney_algorithm  (2t=16)
    
      // Compute Omega
      u8 omega[16];
      memset(omega, 0, sizeof(omega));
      // TODO loops
      for ( int i=0; i<16; ++i )
	for ( int j=0; j<16; ++j )
	  if ( i+j < 16 ) omega[i+j] ^= gf.mul(synd[i], C[j]);
#if DEBUG_RS
      fprintf(stderr, "omega=");
      for ( int i=0; i<16; ++i ) fprintf(stderr, " %d", omega[i]);
      fprintf(stderr, "\n");
#endif
    
      // Compute Lambda'
      u8 Cprime[15];
      for ( int i=0; i<15; ++i )
	Cprime[i] = (i&1) ? 0 : C[i+1];
#if DEBUG_RS
      fprintf(stderr, "Cprime=");
      for ( int i=0; i<15; ++i ) fprintf(stderr, " %d", Cprime[i]);
      fprintf(stderr, "\n");
#endif
    
      // Find zeroes of C by exhaustive search?
      // TODO Chien method
      int roots_found = 0;
      for ( int i=0; i<255; ++i ) {
	u8 r = gf.exp(i);  // Candidate root alpha^0..alpha^254
	u8 v = eval_poly(C, L, r);
	if ( ! v ) {
	  // r is a root X_k^-1 of the error locator polynomial.
	  u8 xk = gf.inv(r);
	  int loc = (255-i) % 255;  // == log(xk)
#if DEBUG_RS
	  fprintf(stderr, "found root=%d, inv=%d, loc=%d\n", r, xk, loc);
#endif
	  if ( loc < 204 ) {
	    // Evaluate e_k
	    u8 num = gf.mul(xk, eval_poly(omega, L, r));
	    u8 den = eval_poly(Cprime, 14, r);
	    u8 e = gf.div(num, den);
	    // Subtract e from coefficient of degree loc.
	    // Note: Coeffients listed by decreasing degree in pin[] and pout[].
	    if ( bits_corrected ) *bits_corrected += hamming_weight(e);
	    if ( loc >= 16 ) pout[203-loc] ^= e;
	    if ( pin ) pin[203-loc] ^= e;
	  }
	  if ( ++roots_found == L ) break;
	}
      }
      if ( pin )
	return syndromes(pin, synd);
      else
	return false;
    }

  };

}  // namespace

#endif  // LEANSDR_RS_H