/*  
	This is a simulation of the decoding of generalized LRPC codes
	under the following assumptions
	- the subspace B contains a basis [b_1,...,b_d] such that T_{*,bj,*} is invertible
	- the syndrome space has dimension r*d 
	This corresponds to Algorithm 1 in the paper 
	Generalized LRPC codes at https://arxiv.org/abs/2305.02053
*/
System("clear"); //clear screen 
printf "================================\n";
printf "Start of Program ...\n";
System("date");
printf "================================\n";

/*================================================
		     Setup of Parameters
=================================================*/  


// Base field.
q := 2; Fq := GF(q);
// Error rank.
r := 5;
// Density of the LRPC.
d := 3;
// dimnsion of the GLRPC code will be k*m


failure_rate_theo := [];
failure_rate_sim  := [];

for eta in [ 1 .. 5 ] do

m := r*(d+1) + eta;
k := r*d + (d-1)*eta;
n := 2*k;


// Vector space where supports lives.
Fqm := VectorSpace(Fq, m);


/*================================================
		     Utility Functions
=================================================*/ 
// Define the T-product.
// Given a,b,T returns T_{a,b,*}.
Prod_T := function(T, a, b)
	T_b := Matrix(Fq, T[1]*Transpose(Matrix(b)));
	for i in [2..m] do
		T_b := HorizontalJoin(T_b, T[i]*Transpose(Matrix(b)));
	end for;
	return a*T_b;
end function;
// Trace product of two matrices (sum(a_ij b_ij)).
TraceProduct := function(A, B)
	if (NumberOfRows(A) eq NumberOfRows(B)) and (NumberOfColumns(A) eq NumberOfColumns(B)) then
		return &+[A[i,j] * B[i,j]: i in [1..NumberOfRows(A)], j in [1..NumberOfColumns(A)]];
	end if;
	//print("Matrices have not the same size");
	return 0;
end function;

// Given A,B two subspaces and T a tensor.
// Return T-product(A,B) the product subspace using T.
ProductSpace := function(A,B,T)
	c := [Prod_T(T, a, b): a in Basis(A), b in Basis(B)];
	return sub<Fqm | c>;
end function;


// Let Supp_C = T-product(Supp_A, Supp_B)
// Try to recover Supp_A from the knwoledge of Supp_C and Supp_B.

// Given a vector space V and a 3-tensor T.
// Find a base {v_1, ..., v_k} of V.
// Return {T_{*,v_1,*}^-1, ..., T_{*,v_k,*}^-1}. 
Get_T_Base := function(V, T)
	T_V := [];
	for v in Basis(V) do
		T_v := Matrix(Fq, T[1]*Transpose(Matrix(v)));
		for i in [2..m] do
			T_v := HorizontalJoin(T_v, T[i]*Transpose(Matrix(v)));
		end for;
		T_V := T_V cat [T_v];
	end for;
	return T_V;
end function;

// Invert Product Space.
// From S = T-product(A,B), T  and B get the subspace A.
InvertProductSpace := function(S, B, T)
	// For each element of a base of B compute T(*,b,*)
	T_B := Get_T_Base(B, T);
	Tinv_S := [];
	for T_b in T_B do
		// Intersect the space S with the image of the matrix T_b.
		Z := S meet Image(T_b);
		// Compute all the solutions over a base of Z.
		X := {};
		for z in Basis(Z) do
			X := X join {Solution(T_b, z)};
		end for;
		// Take the subspace generated by the solution plus the kernel.
		// This subspace contains the error space.
		Tinv_S := Tinv_S cat [sub<Fqm | X> + Kernel(T_b)];
	end for;
	
	// Intersect all the Tinv_S.
	// Most likely this will be exactly the error space.
	A := Tinv_S[1];
	for i in [2..#T_B] do
		A := A meet Tinv_S[i];
	end for;
	return A;
end function;
// End utilities.


/*================================================
		     Experiments
=================================================*/ 

// Play with random, so we can reproduce interesting cases.


er1 := Real(q^(r*d - (n- k))); 
er2 := Real(q^((d-1)*(r*(d+1)-m )));
er3 := Real(q^(d+(d-1)*(r*(d+1)-m )));

Expected_ProductRecovery_Error    := er1;
Expected_ErrorRecovery_Error1     := er1 + er2;
Expected_ErrorRecovery_Error2     := er1 + er3;
Expected_ErrorSuppRecovery_Error  := er1 + er2;


Test := function(trials)
	// Create the code.
	// Generate a random sequence of matrices T[k] = T_{*,*,k}.
	// This will be the 3-tensor T, generating the T-product.
	T := [Matrix(Fq, m ,m, [Random(Fq): i in [1..m*m]]): j in [1..m]];

	//Generate a subspace of dimension d.
	// Create a subspace  SuppH of (Fq)^m of dimension d.
	SuppH := sub<Fqm | [Random(Fqm): i in [1..d]]>;
	// Check column support is really of dimension d.
	while(Dimension(SuppH) lt d ) do
		SuppH := sub<Fqm | [Random(Fqm): i in [1..d]]>;
	end while;

	// Generate H_1, ..., H_{n-k} with column space givn by Supp_H.
	genH := [Transpose(Matrix(Fq, n, m, [Random(SuppH): i in [1..n]])): j in [1..n-k]];

	// Fix a base for SuppH. We will use that base to reconstruct the error.
	B := Transpose(Matrix(Basis(SuppH)));
	// Find coordinates matrices of H_j with respect to the base B.
	// That is find Y_j such that BY_j = H_j.
	Y := [Transpose(Matrix(Fq, n, d, [Solution(Transpose(B), Transpose(genH[j])[i]): i in [1..n]])): j in [1..n-k]];

	// Counters for different steps
	ProductCount       := 0;
	ErrorSupportCount  := 0;
	ErrorRecoveryCount := 0;
	Fails              :=  [0, 0, 0, 0, 0];
	for trial in [1..trials] do		
		SetSeed(trial);
		// Create a subspace  SuppE of (Fq)^m of dimension r.
		SuppE := sub<Fqm | [Random(Fqm): i in [1..r]]>;
		// Generate the T-poduct space between SuppE and SuppH.
		SuppEH := ProductSpace(SuppE, SuppH, T);
		// Check column support is really of dimension r.
		while( Dimension(SuppE) lt r) or (Dimension(SuppEH) lt r*d ) do
			SuppE := sub<Fqm | [Random(Fqm): i in [1..r]]>;
			SuppEH := ProductSpace(SuppE, SuppH, T);
		end while;
		
		// Generate an error matrix Err in (F_q)^(m x n) such that Colsp(Err) = E.
		Err := Transpose(Matrix([ElementToSequence(Random(SuppE)): i in [1..n]]));
		
		while( Rank(Err) lt r) do
			Err := Transpose(Matrix([ElementToSequence(Random(SuppE)): i in [1..n]]));
			
		end while;
		
		
		
		// Generate the syndrome matrix S.
		S := ZeroMatrix(Fq, m, n-k);
		for i in [1..m] do
			for j in [1..n-k] do
				S[i,j] := TraceProduct(Err, T[i] * genH[j]);
			end for;
		end for;
		
		// Recover Big_Supp_E_rec as Colsp(S).
		SuppEH_rec := sub<Fqm | [Fqm ! Transpose(S)[j] : j in [1..n-k]]>;
		
		// Check that SuppEH recovered is the same as SuppEH
		if SuppEH_rec eq SuppEH then
			ProductCount +:= 1;
			//print "Recovered the correct product space.";
		else
			//print "Failed to recover the product space.";
			Fails[1] +:= 1;
			continue trial;	
		end if;
		
		// Do the intersection of SuppEH_rec, verify if it is equal to SuppE.
		SuppE_rec := InvertProductSpace(SuppEH_rec, SuppH, T);
		// Check we found the support of the error.
		if SuppE_rec eq SuppE then
			//print "Recovered the correct error support.";
			ErrorSupportCount +:= 1;
		else
			//print "Failed to recover the error support.";
			Fails[2] +:= 1;
			continue trial;
		end if;	
		//In order to recover the error coordinate with our tecnique we need that dim(SuppEH) == rd.
		if Dimension(SuppEH) lt r*d then
			Fails[3] +:= 1;
			continue trial;
		end if;
		// Expand the linear system and find X such that EX = Err.
		// Fix a base for SuppE_rec. We will use that base to reconstruct the error.
		A := Transpose(Matrix(Basis(SuppE_rec)));
		// Find X such that FX = Err.
		X := Transpose(Matrix(Fq, n, r, [Solution(Transpose(A), Transpose(Err)[i]): i in [1..n]]));
		// Find coordinates Z of S w.r.t its base C.
		// C is obtained by the products between A,B.
		C := Transpose(Matrix(Fq, r*d, m, [Prod_T(T, Transpose(A)[i], Transpose(B)[j]): i in [1..r], j in [1..d]]));
		Z := Transpose(Matrix(Fq, n-k, r*d, [Solution(Transpose(C), Transpose(S)[i]): i in [1..n-k]]));	
		// Split Z in r matrices (n-k) x d by taking one column each r.
		SplitZ := [Matrix(Fq, d, n-k, [Z[j + i*r]: i in [0..d-1]]): j in [1..r]];		
		// Column concatenation of the r matrices of SplitZ.
		ZZ := [HorizontalJoin([Transpose(SplitZ[j])[i]: i in [1..n-k]]): j in [1..r]];		
		// Merge matrices in Y.
		YY := HorizontalJoin([Transpose(Y[i]) : i in [1..n-k]]);		
		// Solve the linear system.
		X_rec := ZeroMatrix(Fq, 1, n);
		for i in [1..r] do
			hasSolution, solution := IsConsistent(YY, ZZ[i]);
			if hasSolution then
				X_rec := VerticalJoin(X_rec, solution);
				//print(Rank(C));
			else 
				print "NO SOLUTION";
				print(Rank(C));
				Fails[4] +:= 1;
				continue trial;
			end if;
		end for;	
		X_rec := Submatrix(X_rec, [2..r+1], [1..n]);		
		// Check we found the correct solution.
		if Err eq A*X_rec then
			//print "Recovered the correct Solution.";
			ErrorRecoveryCount +:= 1;
		else
			//print "Wrong solution.";
			Fails[5] +:= 1;
			continue trial;
		end if;		
	end for;

	

	// Look at the counters:
	printf "[ m, n, k, r, d, eta ] = %o\n", [m,n,k,r,d,eta];
	// printf "Expected Product Space recovery: %o \nMeasured Product Space recovery: %o \n \n", Expected_ProductRecovery_Error, ProductCount;
	// printf "Expected Error Support recovery: %o \nMeasured Error Support recovery: %o \n \n", Expected_ErrorSuppRecovery_Error, 1.0*(trials - ErrorSupportCount)/trials;
	// printf "Expected Error recovery: %o \nMeasured Error recovery: %o \n \n", Expected_ErrorRecovery_Error, ErrorRecoveryCount;
	printf "Fails = %o\n", Fails;
	return  1.0*(trials - ErrorSupportCount)/trials;

end function;

trials := 2^12;
Append(~failure_rate_theo, Expected_ErrorRecovery_Error2);
Append(~failure_rate_sim, Test(trials));

// printf "DFR=%o\n", Expected_ErrorRecovery_Error2;

end for;

printf "%o\n", failure_rate_theo;
printf "%o\n", failure_rate_sim;


printf "================================\n";
printf "End of Program!\n";
System("date");
printf "================================\n";