/* ---------------------------------------------------------
Generation of a Rainbow key pair over GF(q) 
public key is stored in public_key.txt
private key is stored in private_key.txt
------------------------------------------------------------*/


q:=4; // field size
parameters:=[3,2,2]; // Rainbow parameters (v_1, o_1, ..., o_u)
u:= #parameters -1;

v:=[];
o:=[];

v[1]:=parameters[1];
for i:=1 to u do
	o[i]:=parameters[i+1];
	v[i+1]:=v[i]+o[i];
end for;

n:=v[u+1];
m:=n-v[1];
F<w>:=GaloisField(q);
Vn:=VectorSpace(F,n);
Vm:=VectorSpace(F,m);

P<[x]>:=PolynomialRing(F,n);
Pol<[y]>:=PolynomialRing(P,n);


// key generation
// affine map T -----------------------------------------------------------------------------
repeat
	MT:=RandomMatrix(F,n,n);
until IsInvertible(MT) eq true;
cT:= Random(Vn);

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

// central map Q ---------------------------------------------------------------------------
Qc:=[]; Q:=[];

for greatloop:=1 to u do

	for loop:=v[greatloop]-v[1]+1 to v[greatloop+1]-v[1] do // greatloop-th Layer --------------------------------
		Q[loop]:=Pol!0;

		for i:=1 to v[greatloop] do
			for j:=1 to v[greatloop+1] do
				Q[loop]:=Q[loop] + Random(F)*Pol.i*Pol.j;
			end for;
		end for;

		for i:=1 to v[greatloop+1] do
			Q[loop]:= Q[loop] + Random(F)*Pol.i;
		end for;

		Q[loop]:=Q[loop]+Random(F);

	end for;
end for; 

// affine map S ---------------------------------------------------------------------------------------

repeat
	MSF:=RandomMatrix(F,m,m);
until IsInvertible(MSF) eq true;
cS:=Matrix(Pol,m,1,ChangeUniverse(Eltseq(Random(Vm)),Pol));
MS:=Matrix(Pol,m,m,ChangeUniverse(Eltseq(MSF),Pol));

// public key Pk ----------------------------------------------------------------------------
QT:=ZeroMatrix(Pol,m,1);
for i:=1 to m do
	QT[i][1]:=Evaluate(Q[i],y[1], T[1]);
	for j:=2 to n do
		QT[i][1]:=Evaluate(QT[i][1],y[j], T[j]);
	end for;
end for;

Pk:=MS*QT+cS;

D:=[];
for i:=1 to m do
	D[i]:=MonomialCoefficient(Pk[i][1],1);
end for;


// Output ----------------------------------------------------------------------------------
printf "************************************************* \n";
printf "*** Rainbow Signature Scheme - Key Generation *** \n";
printf "************************************************* \n \n";
printf "q:= %o ; \n \n", q;
printf "parameters:= %o ; \n \n", parameters;
printf "MS:= \n%o  \n \n", MS;
printf "cS:= %o  \n \n", Vm!Eltseq(cS);
printf "MT:= \n%o  \n \n", MT;
printf "cT:= %o  \n \n", cT;
printf "Q:= %o  \n \n", Q;
printf "Pk:= %o  \n \n",D ;


printf"Write public_key.txt \n \n";
SetOutputFile("public_key.txt":Overwrite:=true);
printf "q:= %o ; \n \n", q;
printf "m:= %o ; \n \n", m;
printf "n:= %o ; \n \n", n;
printf "parameters:= %o ; \n \n", parameters;
printf "F<w>:=GaloisField(q); \n \n";
printf "Pol<[x]>:=PolynomialRing(F,n); \n \n";
printf "Pk:= %o ; \n \n",D ;
UnsetOutputFile();
 
printf"Write private_key.txt \n \n";
SetOutputFile("private_key.txt":Overwrite:=true);
printf "q:= %o ; \n \n", q;
printf "v:= %o ; \n \n", v;
printf "n:= %o ; \n \n", n;
printf "m:= %o ; \n \n", m;
printf "F<w>:=GaloisField(q); \n \n";
printf "Pol<[y]>:=PolynomialRing(F,n); \n \n";
printf "Q:= %o ; \n \n", Q;
printf "MT:= %o ; \n \n", Eltseq(MT);
printf "cT:= %o ; \n \n", Eltseq(cT);
printf "MS:= %o ; \n \n", Eltseq(MS);
printf "cS:= %o ; \n \n", Eltseq(cS);
UnsetOutputFile();