/*  Decoding Reed-Solomon Codes  with the 
    Berlekamp-Massey algorithm.
*/

//  An example field.  
q := 9;
n := q-1;		// Length of the code.
k := 4; 		// Dimension of the code.
Fq<a> := GF(q);        //Creates the field.

// Let C be the Reed-Solomon code of dimension k and length n=k-1
// constructed in the standard form identified in Sec. 5 
// of the course notes.
// The following generate the generator matrix, G, and the check 
// matrix, H, for C.

Vec_n := VectorSpace(Fq, n); 		// n dim vec space.
Vec_k := VectorSpace(Fq, k);		// k dim vec space.

MatH := KMatrixSpace(Fq, n, n-k);	// Space for check matrix.
MatG := KMatrixSpace(Fq, k, n);		// Space for gen matrix.

G := MatG ! [ a^((i-1)*(j-1)) :  j in [1..n], i in [1..k]];
H := MatH ! [ a^((i-1)*(j)) :  j in [1..n-k], i in [1..n]];
G*H;

// ********************************************
//  Creation of a message vector.
//  Systematic encoding to a code vector.
//  Corruption by an error vector of bounded weight.
//  Computation of the syndrome.

P<x>  := PolynomialRing(Fq);
gx := &*[x-a^i : i in [1..n-k] ] ;  // Generator poly for C.
m := Random(Vec_k);
mx := P ! [m[i] : i in [1..k] ];
mxp := mx*x^(n-k);
qx, rx := Quotrem(mxp,  gx);
cx := mxp- rx;
c := Vec_n ! [ Coefficient(cx,i) : i in [0.. n-1] ];

// The maximum numbers of errors that can be corrected is (d-1)/2.
// You might want to alter to be t of any weight less than d/2.
// Or you could simulate a channel by corrupting a symbol with
// a given  probability.

t := Floor((n-k)/2); 		
locs := Random(Subsets({1..n},t)); 	// A random set of t locations.

e := Vec_n !0;	 	// Random values at the given locations.
for l in locs do
   e[l] := Random(Fq);	// Possibly 0.  This can be altered.
end for;

// The vector recieved.
v := c+e;
// The syndrome vector.
syn := v*H;

//   ************************************************
// Some functions introduced in Sec. 5 of the notes.

Syn := function(syn, fx)
 print("here"); print Degree(fx), Degree(syn);
 if Degree(fx) ge Degree(syn) then
   print "Polynomial is of larger degree than the known syndromes";
   return 0;
 end if;
 fvec := Parent(syn) ! [ Coefficient(fx,i) : i in [0.. n-k-1] ];
 return InnerProduct(syn, fvec);
end function;

Spandiscfail := function(syn, fx)
 d := Fq ! 0;  i := -1;
 while d eq 0 and i lt Degree(syn)-Degree(fx)-1 do
    i:= i+1;
    d := Syn(syn, x^i*fx);
 end while;
 if d eq 0 then 
   print "Span appears to be infinite";
   return -1, 0, -1;
 else  
   return i, d, i + Degree(fx);
end if;
end function;
   
Disc := func< syn, fx, shft | Syn(syn, x^shft*fx) >;
 
// ************************************************
// The Berlekamp-Massey algorithm

// The Update procedure:

Update := procedure(~info, m , syn)
 fx := info[1,1];
 c := Degree(fx) - 1;
 r := m - c -1;
 mu := Syn(syn, x^r*fx);
 print "m is", m, "f is ",  fx;
 print "shift of g is ", c, "shift of f is ", r, "discrepancy of f is", mu;
 if (mu eq 0) then
     Q := RMatrixSpace(Parent(fx),2,2) ! [1,0, 0,1];
 elif (r le c ) then
    Q := RMatrixSpace(Parent(fx),2,2) ! [1, -mu*x^(c-r), 0,1];
 else
    Q := RMatrixSpace(Parent(fx),2,2) ! [x^(r-c), -mu,  mu^(-1), 0];
 end if;
print "Update matrix is ", Q;
 info := Q*info;
print "New info is", info;
end procedure;

// Initialization:
MatPoly := RMatrixSpace(P,2,1);
info := MatPoly ![1,0];

// Run the algorithm: 
for m in [0..n-k-1] do
//   print "info is ", info;
   Update(~info, m, syn);
end for;

// **********************************************
// Find the error locations

FindLocations := function(fx, vec_n)
elocs := vec_n !0;
for  i in [0..Degree(vec_n)-1 ] do
  if (Evaluate(fx,a ^i) eq 0) then
    elocs[i+1] := 1;
  end if;
end for;
return elocs;
end function;

FindVals := function(fx, gx, erlocs)
evals  := erlocs;
fpr := Derivative(fx);
// Degree(evals) is its length, n.
for  i in [0..Degree(evals) -1 ] do
  if (evals[i+1] eq (Fq ! 1)) then
    evals[i+1] := -(Evaluate(fpr,a^i)*Evaluate(gx,a^i))^(-1) ;
  end if;
end for;
return evals;
end function;

e_locs := FindLocations(info[1,1], Vec_n);
print "The locations of the errors are the nonzero entries in \n", e_locs;
e_out := FindVals(info[1,1], info[2,1], e_locs);
print "The error vector as determined by the algorithm is \n ", e_out;
print "The original error vector was \n ", e_out;