#include "common.h"

extern int parityNum;
extern Mapping erasures, errors;

/* Maximum number of parity bytes */
#define NPAR 20
/* Maximum degree of various polynomials. */
#define MAXDEG (NPAR*2)

#define DEBUG FALSE

/* galois arithmetic tables */
int gexp[512];
int glog[256];

/* generator polynomial */
int genPoly[MAXDEG*2];

/* Encoder parity bytes */
int pBytes[MAXDEG];

/* Decoder syndrome bytes */
int synBytes[MAXDEG];

/* error locations */
int ErrorLocs[256];
int NErrors;

/* The Error Locator Polynomial, Lambda[0] == 1 */
int Lambda[MAXDEG];

/* The Error Evaluator Polynomial */
int Omega[MAXDEG];

/* multiplication using logarithms */
int gmult(int a, int b)
{
	int i,j;
	if (a==0 || b == 0) return (0);
	i = glog[a];
	j = glog[b];
	return (gexp[i+j]);
}

int ginv (int elt) 
{ 
	return (gexp[255-glog[elt]]);
}

void zero_poly (int poly[]) 
{
	int i;
	int n = parityNum*2;
	for (i = 0; i < n; i++) 
		poly[i] = 0;
}

void copy_poly (int dst[], int src[]) 
{
	int i;
	for (i = 0; i < (parityNum*2); i++) dst[i] = src[i];
}

/* polynomial multiplication */
void mult_polys (int dst[], int p1[], int p2[])
{
	int i, j;
	int tmp1[MAXDEG*2];
	
	for (i=0; i < (parityNum*4); i++) dst[i] = 0;
	
	for (i = 0; i < (parityNum*2); i++) {
		for(j=(parityNum*2); j<(parityNum*4); j++) tmp1[j]=0;
		
		for(j=0; j<(parityNum*2); j++) tmp1[j]=gmult(p2[j], p1[i]);
		
		for (j = ((parityNum*4)-1); j >= i; j--) tmp1[j] = tmp1[j-i];
		for (j = 0; j < i; j++) tmp1[j] = 0;
		
		for(j=0; j < (parityNum*4); j++) dst[j] ^= tmp1[j];
	}
}

/* This is one of 14 irreducible polynomials
* of degree 8 and cycle length 255.
* The high order 1 bit is implicit */
/* x^8 + x^4 + x^3 + x^2 + 1 */

static void init_exp_table (void)
{
	int i, z;
	int pinit,p1,p2,p3,p4,p5,p6,p7,p8;
	
	pinit = p2 = p3 = p4 = p5 = p6 = p7 = p8 = 0;
	p1 = 1;
	
	gexp[0] = 1;
	gexp[255] = gexp[0];
	glog[0] = 0;			
	
	for (i = 1; i < 256; i++) {
		pinit = p8;
		p8 = p7;
		p7 = p6;
		p6 = p5;
		p5 = p4 ^ pinit;
		p4 = p3 ^ pinit;
		p3 = p2 ^ pinit;
		p2 = p1;
		p1 = pinit;
		gexp[i] = p1 + p2*2 + p3*4 + p4*8 + p5*16 + p6*32 + p7*64 + p8*128;
		gexp[i+255] = gexp[i];
	}
	
	for (i = 1; i < 256; i++) {
		for (z = 0; z < 256; z++) {
			if (gexp[z] == i) {
				glog[i] = z;
				break;
			}
		}
	}
}

void init_galois_tables (void)
{	
	init_exp_table();
}

/* Create a generator polynomial for an n byte RS code. 
* The coefficients are returned in the genPoly arg.
*/

static void compute_genpoly (int nbytes, int genpoly[])
{
	int i, tp[256], tp1[256];
	
	zero_poly(tp1);
	tp1[0] = 1;
	
	for (i = 1; i <= nbytes; i++) {
		zero_poly(tp);
		tp[0] = gexp[i];		
		tp[1] = 1;
		
		mult_polys(genpoly, tp, tp1);
		copy_poly(tp1, genpoly);
	}
}

/* Initialize lookup tables, polynomials, etc. */
void initialize_ecc ()
{
	/* Initialize the galois field arithmetic tables */
    init_galois_tables();
	
    /* Compute the encoder generator polynomial */
    compute_genpoly(parityNum, genPoly);
}

/* Append the parity bytes */
void build_codeword (unsigned char msg[], int nbytes, unsigned char dst[])
{
	int i;
	
	for (i = 0; i < nbytes; i++) dst[i] = msg[i];
	
	for (i = 0; i < parityNum; i++) {
		dst[i+nbytes] = pBytes[parityNum-1-i];
	}
}


void encode_data (unsigned char msg[], int nbytes, unsigned char dst[])
{
	int i, LFSR[NPAR+1],dbyte, j;
	
	for(i=0; i < parityNum+1; i++) LFSR[i]=0;
	
	for (i = 0; i < nbytes; i++) {
		dbyte = msg[i] ^ LFSR[parityNum-1];
		for (j = parityNum-1; j > 0; j--) {
			LFSR[j] = LFSR[j-1] ^ gmult(genPoly[j], dbyte);
		}
		LFSR[0] = gmult(genPoly[0], dbyte);
	}
	
	for (i = 0; i < parityNum; i++) 
		pBytes[i] = LFSR[i];
	
	build_codeword(msg, nbytes, dst);
}

/* Computes the syndrome of a codeword. Puts the results
 * into the synBytes[] array.
 */
void decode_data(unsigned char data[], int nbytes)
{
  int i, j, sum;
  for (j = 0; j < parityNum;  j++) {
    sum	= 0;
    for (i = 0; i < nbytes; i++) {
      sum = data[i] ^ gmult(gexp[j+1], sum);
    }
    synBytes[j]  = sum;
  }
}

/* Check if the syndrome is zero */
int check_syndrome (void)
{
	int i, nz = 0;
	for (i =0 ; i < parityNum; i++) {
		if (synBytes[i] != 0) nz = 1;
	}
	return nz;
}

void scale_poly (int k, int poly[]) 
{	
	int i;
	for (i = 0; i < parityNum*2; i++) 
		poly[i] = gmult(k, poly[i]);
}

void add_polys (int dst[], int src[]) 
{
	int i;
	for (i = 0; i < parityNum*2; i++) dst[i] ^= src[i];
}

void mul_z_poly (int src[])
{
	int i;
	for (i = parityNum*2-1; i > 0; i--) src[i] = src[i-1];
	src[0] = 0;
}

/* gamma = product (1-z*a^Ij) for erasure locs Ij */
void init_gamma (int gamma[])
{
	int e, tmp[MAXDEG];
	
	zero_poly(gamma);
	zero_poly(tmp);
	gamma[0] = 1;
	
	int size = erasures.Count();
	for (e = 0; e < size; e++) {
		copy_poly(tmp, gamma);
		scale_poly(gexp[erasures.getKey(e)], tmp);
		mul_z_poly(tmp);
		add_polys(gamma, tmp);
	}
}

int compute_discrepancy (int lambda[], int S[], int L, int n)
{
	int i, sum=0;
	
	for (i = 0; i <= L; i++) 
		sum ^= gmult(lambda[i], S[n-i]);
	return (sum);
}

void compute_modified_omega ()
{
	int i;
	int product[MAXDEG*2];
	
	mult_polys(product, Lambda, synBytes);	
	zero_poly(Omega);
	for(i = 0; i < parityNum; i++) 
		Omega[i] = product[i];	
}

void Find_Position (void)
{	
	int n, L, L2, k, d, i;
	int psi[MAXDEG], psi2[MAXDEG], D[MAXDEG];
	int gamma[MAXDEG];
	
	/* initialize Gamma, the erasure locator polynomial */
	init_gamma(gamma);
	
	copy_poly(D, gamma);
	mul_z_poly(D);
	
	copy_poly(psi, gamma);	
	k = -1; L = erasures.Count();
	
	for (n = erasures.Count(); n < parityNum; n++) {
		
		d = compute_discrepancy(psi, synBytes, L, n);
		
		if (d != 0) {
			
			for (i = 0; i < parityNum*2; i++) psi2[i] = psi[i] ^ gmult(d, D[i]);
			
			if (L < (n-k)) {
				L2 = n-k;
				k = n-L;
				
				for (i = 0; i < parityNum*2; i++) D[i] = gmult(psi[i], ginv(d));
				L = L2;
			}
			
			for (i = 0; i < parityNum*2; i++) psi[i] = psi2[i];
		}
		
		mul_z_poly(D);
	}
	
	for(i = 0; i < parityNum*2; i++) 
		Lambda[i] = psi[i];
	compute_modified_omega();	
}

void Find_Roots (void)
{
	int sum, r, k;	
	NErrors = 0;
	
	for (r = 1; r < 256; r++) {
		sum = 0;
		
		for (k = 0; k < parityNum+1; k++) {
			sum ^= gmult(gexp[(k*r)%255], Lambda[k]);
		}
		if (sum == 0) 
		{ 
			ErrorLocs[NErrors] = (255-r); 
			NErrors++; 
			if (DEBUG) {
				CString stemp;
				stemp.Format("Root found at r = %d, (255-r) = %d\n", r, (255-r));
				AfxMessageBox(stemp);
			}
		}
	}
}

/* 
* Pass in the codeword, its size in bytes, as well as
* an array of any known erasure locations, along the number
* of these erasures.
* 
* Returns 1 if everything ok, or 0 if an out-of-bounds error is found
*
*/
int correct_errors_erasures (unsigned char codeword[], int csize)
{
	int r, i, j, err;
	
	Find_Position();
	Find_Roots();
	
	if ((NErrors <= parityNum) && NErrors > 0) { 
		
		/* first check for illegal error locs */
		for (r = 0; r < NErrors; r++) {
			if (ErrorLocs[r] >= csize) {
				if (DEBUG) {
					CString stemp;
					stemp.Format("Error loc i=%d outside of codeword length %d\n", ErrorLocs[r], csize);
					AfxMessageBox(stemp);
				}
				return(0);
			}
		}
		
		for (r = 0; r < NErrors; r++) {
			int num, denom;
			i = ErrorLocs[r];
						
			num = 0;
			for (j = 0; j < parityNum*2; j++) 
				num ^= gmult(Omega[j], gexp[((255-i)*j)%255]);
			
			denom = 0;
			for (j = 1; j < parityNum*2; j += 2) {
				denom ^= gmult(Lambda[j], gexp[((255-i)*(j-1)) % 255]);
			}
			
			err = gmult(num, ginv(denom));
			if (DEBUG) {
				CString stemp;
				stemp.Format("Error magnitude %#x at loc %d\n", err, csize-i);
				AfxMessageBox(stemp);
			}
			
			codeword[csize-i-1] ^= err;
		}
		return(1);
	}
	else {
		if (DEBUG && NErrors) {
			CString stemp;
			stemp.Format("Uncorrectable codeword");
			AfxMessageBox(stemp);
		}
		return(0);
	}
}