/* ==============================================
 Key Generation of the Matsumoto Imai Cryptosystem
 The program generates a key pair of the MI cryptosystem
 and stores private and public key in the files public_key.txt and private_key.txt
===========================================================*/

time_start:=Realtime();

function squareandmultiply(m,e,p) // for fast exponentiation of polynomials
	e:=Intseq(e,2);
	s:=# e; // bitlength of e
	erg:=1;
	for i:=1 to s do
		if e[i] eq 1 then
			erg:=erg *m mod p;
		end if;
		m:=m^2 mod p;
	end for;
	return erg;
end function;

// ============ parameters =====================

SetSeed(1994655094);
q:=4; // field size
n:=3; // # variables
theta:=2; // MI parameter


timet:=Cputime();
GF<w>:=GaloisField(q);
Vectn:=VectorSpace(GF,n);

Pol<[x]>:=PolynomialRing(GF,n);
Poly<X>:=PolynomialRing(Pol);
irrpol:=Poly!(IrreduciblePolynomial(GF,n));

FIELDEQ:=[]; // field equations over ground field
for i:=1 to n do
	FIELDEQ:=FIELDEQ cat [x[i]^q-x[i]];
end for;
FIELDIdeal:=Ideal(FIELDEQ);
// =============================================

g,thetainv,_:=XGCD(q^theta+1,q^n-1); // parameter for decryption
thetainv:=thetainv mod (q^n-1);

if g gt 1 then
	printf"ERROR!!! Wrong theta";
end if;


// affine map T ----------------------------------------------------------------
repeat
	MT:=RandomMatrix(GF,n,n);
until IsInvertible(MT) eq true;
cT:=Random(Vectn);

T:=[];
for i:=1 to n do
	T[i]:=cT[i]+MT[i][1]*x[1];
	for j:= 2 to n do
		T[i]:=T[i]+MT[i][j]*x[j];
	end for;
end for;
// ----------------------------------------------------------------------------

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

repeat
	MSF:=RandomMatrix(GF,n,n);
until IsInvertible(MSF) eq true;
MS:=Matrix(Pol,n,n,ChangeUniverse(Eltseq(MSF),Pol));
cSF:=RandomMatrix(GF,n,1);
cS:=Matrix(Pol,n,1,ChangeUniverse(Eltseq(cSF),Pol));
// -----------------------------------------------------------------------

MIA:=0; // lift into the extension field
for i:=1 to n do
	MIA:=MIA+T[i]*X^(i-1);
end for;

MIB:=squareandmultiply(MIA,1+q^theta,irrpol); // MIB=MIA^(1+q^theta)

p:=[];  // move down to the vector space
for i:=1 to n do
	p[i]:=MonomialCoefficient(MIB,X^(i-1));
end for;

for i:=1 to n do // reduce modulo field equations
	p[i]:=NormalForm(p[i],FIELDIdeal);
end for;

MIB:=Pol!0;
for i:=1 to n do
	MIB:=MIB+p[i]*X^(i-1);
end for;

P:=ZeroMatrix(Pol,n,1);
for i:=1 to n do
	P[i][1]:=p[i];
end for;

Pk:=MS*P+cS; // Compute public key

for i:=1 to n do
	Pk[i][1]:=NormalForm(Pk[i][1],FIELDIdeal); // reduce modulo field equations
end for;


timet:=Cputime(timet);
print "MI keygen time:",timet;
mem_usage:=GetMemoryUsage();

// Output ==========================================================================
printf "****************************************************\n";
printf "*** Matsumoto-Imai Cryptosystem - Key Generation *** \n";
printf "**************************************************** \n \n";

printf "q:= %o  \n \n", q;
printf "n:= %o  \n \n", n;
printf "irrpol:= %o  \n \n", irrpol;
printf "theta:= %o \n \n", theta;
printf "thetainv:= %o \n \n", thetainv;
printf "S:= \n%o \n \n", MS;
printf "cS:= %o \n \n", Vectn!Eltseq(cS);

printf"Write public_key.txt \n \n";
SetOutputFile("public_key.txt":Overwrite:=true);
printf "q:= %o ; \n \n", q;
printf "n:= %o ; \n \n", n;
printf "GF<w>:=GaloisField(q); \n \n";
printf "Pol<[x]>:=PolynomialRing(GF,n); \n \n";
printf "Pk:= %o ; \n \n", Eltseq(Pk);
UnsetOutputFile();
 
printf"Write private_key.txt \n \n";
SetOutputFile("private_key.txt":Overwrite:=true);
printf "q:= %o ; \n \n", q;
printf "n:= %o ; \n \n", n;
printf "GF<w>:=GaloisField(q); \n \n";
printf "Pol<[x]>:=PolynomialRing(GF,n); \n \n";
printf "Poly<X>:=PolynomialRing(GF); \n \n";
printf "thetainv:= %o ;\n \n", thetainv;
printf "irrpol:= %o ; \n \n", irrpol;
printf "MT:= %o ; \n \n", Eltseq(MT);
printf "cT:= %o ; \n \n", Eltseq(cT);
printf "MS:= \n%o ; \n \n", Eltseq(MS);
printf "cS:= \n%o ; \n \n", Eltseq(cS);
UnsetOutputFile();
