/* =====================================================================
              Key Generation for SimpleMatrix 
The program generates a random key pair for the SimpleMatrix Encryption Scheme
public key is stored in public_key.txt and private key is stored in private_key.txt
==========================================================================*/
clear ;

q:=256;
s:=2;
m:=2*s^2;
n:=s^2;
F<w>:=GF(q);

Pol<[x]>:=PolynomialRing(F,n);
Poly<[y]>:=PolynomialRing(Pol,n);
Poll<[z]>:=PolynomialRing(Poly,m);

Vn:=VectorSpace(F,n);
Vm:=VectorSpace(F,m);

//Linear map T --------------------------------------------


T:=[];

repeat
	MT:=RandomMatrix(F,n,n);
until IsInvertible(MT) ;
cT:=Random(Vn);

for loop:=1 to n do
	T[loop]:=Pol!0;
	for i:=1 to n do
		T[loop]:=T[loop]+MT[loop][i]*x[i];
	end for;
	T[loop]:=T[loop]+cT[loop];
end for;

//Linear map S -------------------------------------------------
S:=[];
repeat 
	MS:=RandomMatrix(F,m,m);
until IsInvertible(MS) eq true;
cS:=Random(Vm);

for loop:=1 to m do
	S[loop]:=Poll!0;
	for i:=1 to m do
		S[loop]:=S[loop]+MS[loop][i]*z[i];
	end for;
	S[loop]:=S[loop]+cS[loop];
end for;

// Central map Q ----------------------------------------------------
A:=ZeroMatrix(Poly,s,s);
B:=ZeroMatrix(Poly,s,s);
C:=ZeroMatrix(Poly,s,s);
BP:=ZeroMatrix(Pol,s,s);
CP:=ZeroMatrix(Pol,s,s);


for i:=1 to s do
	for j:=1 to s do
		for k:=1 to n do
			A[i][j]:=A[i][j]+Random(F)*y[k];
		end for;
	end for;
end for;

Seq:=[];
for i:=1 to n do
	Seq:=Seq cat [i];
end for;

for i:=1 to s do
	for j:=1 to s do
		for k:=1 to n do
			a:=Random(F);
			b:=Random(F);
			B[i][j]:=B[i][j]+ a*y[k];
			C[i][j]:=C[i][j]+b*y[k];
			BP[i][j]:=BP[i][j]+a*x[k];
			CP[i][j]:=CP[i][j]+b*x[k];
		end for;
	end for;
end for;

E1:=A*B;
E2:=A*C;

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


// Public key computation
for loop:=1 to m do
	for i:=1 to n do
		Q[loop]:=Evaluate(Q[loop],y[i],T[i]);
	end for;
end for;

P:=S;
Pk:=[];
for loop:=1 to m do
	for i:=1 to m do
		P[loop]:=Evaluate(P[loop],z[i],Q[i]);
	end for;
	Pk[loop]:=MonomialCoefficient(MonomialCoefficient(P[loop],1),1);
end for;


// Output ----------------------------------------------------------------------------

printf "******************************************************* \n";
printf "*** SimpleMatrix Encryption Scheme - Key Generation *** \n";
printf "******************************************************* \n \n";

printf "q:= %o  \n \n", q;
printf "s:= %o  \n \n", s;
printf "MS:= \n%o  \n \n", MS;
printf "cS:= %o  \n \n", cS;
printf "MT:= \n%o  \n \n", MT;
printf "cT:= %o  \n \n", cT;
printf "B:= %o  \n \n", BP;
printf "C:= %o  \n \n", CP;
printf "Pk:= %o  \n \n",Pk;

printf "Write public_key.txt \n \n";
SetOutputFile("public_key.txt":Overwrite:=true);
printf "q:= %o ; \n \n", q;
printf "s:= %o ; \n \n", s;
printf "n:= s^2; \n \n";
printf "m:=2*n; \n \n";
printf "F<w>:=GaloisField(q); \n \n";
printf "Pol<[x]>:=PolynomialRing(F,m); \n \n";
printf "Pk:= %o ; \n \n",Pk;
UnsetOutputFile();
 
printf "Write private_key.txt \n \n";
SetOutputFile("private_key.txt":Overwrite:=true);
printf "q:= %o ; \n \n", q;
printf "s:= %o ; \n \n", s;
printf "n:=s^2; \n \n";
printf "m:=2*n; \n \n";
printf "F<w>:=GaloisField(q); \n \n";
printf "Pol<[x]>:=PolynomialRing(F,m); \n \n";
printf "S:= %o ; \n \n", Eltseq(MS);
printf "cS:= %o ; \n \n", Eltseq(cS);
printf "cT:= %o ; \n \n", Eltseq(cT);
printf "T:= %o ; \n \n", Eltseq(MT);
printf "B:= %o ; \n \n", Eltseq(BP);
printf "C:= %o ; \n \n", Eltseq(CP);
UnsetOutputFile();