load "private_key.txt";
load "ciphertext.txt";

function Relinearization(pol1,m,n,solp);
	Mat:=ZeroMatrix(F,m,Integers()!(n*(n+1)/2));
	counter:=1;
	for j:=1 to n do
		for k:=j to n do
			for i:=1 to m do
				Mat[i][counter]:=MonomialCoefficient(pol1[i],Pol.j*Pol.k);
			end for;
			counter:=counter+1;
		end for;
	end for;
	solp:=VectorSpace(F,m)!solp;
	_, vec,ker:=IsConsistent(Transpose(Mat),solp);
	recsol:=Zero(VectorSpace(F,n));
	for i:=1 to n do
		pl:=Integers()!((n*(n+1)-(n-i+1)*(n-i+2))/2+1);
		recsol[i]:=vec[pl]^Integers()!(q div 2);
	end for;
	return recsol;
end function;

function matrixprod(A,B,m,n,o);
	Erg:=ZeroMatrix(Pol,m,o);
	for i:=1 to m do
		for j:=1 to o do
			for k:=1 to n do
				Erg[i][j]:=Erg[i][j]+A[i][k]*B[k][j];
			end for;
		end for;
	end for;
	return Erg;
end function;

S:=Matrix(F,m,m,S);
T:=Matrix(F,n,n,T);
A:=Matrix(Poly,s,r,A);
B:=Matrix(Pol,r,u,B);
C:=Matrix(Pol,r,v,C);
Vn:=VectorSpace(F,n);
Vd:=VectorSpace(F,r*(u+v));

// private key computation ---------------------------------------------------------------------
B2:=ZeroMatrix(Poly,r,u);
C2:=ZeroMatrix(Poly,r,v);
for i:=1 to r do
	for j:=1 to u do
		for k:=1 to n do
			B2[i][j]:=B2[i][j]+MonomialCoefficient(B[i][j],x[k])*y[k];
			C2[i][j]:=C2[i][j]+MonomialCoefficient(C[i][j],x[k])*y[k];
		end for;
	end for;
end for;

E1:=A*B2;
E2:=A*C2;

Q:=[];
for i:=1 to s do
	for j:=1 to u do
		Q:=Q cat [E1[i][j]];
	end for;
end for;
for i:=1 to s do
	for j:=1 to v do
		Q:=Q cat [E2[i][j]];
	end for;
end for;

// first step ------------------------------------------------------------------------------------------------------------------------------
ix:=ciphertext*Transpose(S^(-1));

E1:=ZeroMatrix(F,s,u); // build matrices E1 and E2
E2:=ZeroMatrix(F,s,v);
for i:=1 to s do
	for j:=1 to u do
		E1[i][j]:=ix[(i-1)*u+j];
	end for;
for j:=1 to v do
		E2[i][j]:=ix[s*u+(i-1)*v+j];
	end for;
end for;


//Second Step-------------------------------------------------------------------------------------------------------------
		
W:=ZeroMatrix(Pol,r,s);
for i:=1 to r do
	for j:=1 to s do
		W[i][j]:=x[n+(i-1)*s+j];
	end for;
end for;

System:=ZeroMatrix(F,r*(u+v),n+r*s); //linear system

Sys1:=matrixprod(W,E1,r,s,u)-B;
	 
	for i:=1 to r do
		for j:=1 to u do
			for k:=1 to n+r*s do
				System[(i-1)*u+j][k]:=MonomialCoefficient(Sys1[i][j],x[k]);
			end for;
		end for;
	end for;
	
Sys2:=matrixprod(W,E2,r,s,v)-C;
	for i:=1 to r do
		for j:=1 to v do
			for k:=1 to n+r*s do
				System[u*r+(i-1)*v+j][k]:=MonomialCoefficient(Sys2[i][j],x[k]);
			end for;
		end for;
	end for;

//printf"System = \n%o \n \n", System;
ReverseColumns(~System);
//printf"System= \n%o \n \n", System;

System:=EchelonForm(HorizontalJoin(System,ZeroMatrix(F,r*(u+v),1)));
//printf"System= \n%o \n \n", System;
for i:=1 to r*s do
	if Submatrix(System,1,1,1,r*s) ne VectorSpace(F,r*s)!0 then
		RemoveRow(~System,1);
	end if;
end for;
for i:=1 to r*s do
	RemoveColumn(~System,1);
end for;
RemoveZeroRows(~System);
printf"Found %o equations in the %o plaintext variables. \n \n", Nrows(System), n;
pol:=[];
for i:=1 to Nrows(System) do
	pol[n-i+1]:=Pol!0;
	for j:=1 to n do
		pol[n-i+1]:=pol[n-i+1]+System[i][j]*x[n+1-j];
	end for;
end for; 


Q1:=[];
for i:=1 to m do
	for j:=n to n-Nrows(System)+1 by -1 do
		Q[i]:=Evaluate(Q[i],y[j],pol[j]-x[j]);
	end for;
	for j:=1 to n-Nrows(System) do
		Q[i]:=Evaluate(Q[i],y[j],x[j]);
	end for;
	Q1[i]:=MonomialCoefficient(Q[i],1);
	Q1[i]:=Pol!(Q1[i]);
end for;

erg:=Relinearization(Q1,m,n-Nrows(System),ix);
yps:=Eltseq(erg);

for i:=n-Nrows(System)+1 to n do
	a:=pol[i]-x[i];
	for i:=1 to n-Nrows(System) do
		a:=Evaluate(a,x[i],yps[i]);
	end for;
	yps:= yps cat [a];
end for;

// third step ------------------------------------------------------------------------------------
 recmess:=Vn!yps*Transpose(T)^(-1);
printf "recovered message: %o \n \n", recmess;