// ace_kem.h

template<class group> ACEKEM<group> :: 
ACEKEM(int cofactormode, KDF& K, HASH& H, int Keylen, group& Grp) 
   : KEM(Keylen), CofactorMode(cofactormode), k(K), hash(H), grp(Grp)
{
	w = RandomBnd(grp.mu);
	x = RandomBnd(grp.mu);
	y = RandomBnd(grp.mu);
	z = RandomBnd(grp.mu);

	gp = grp.mul(w, grp.g);
	c = grp.mul(x, grp.g);
	d = grp.mul(y, grp.g);
	h = grp.mul(z, grp.g);
}

template<class group> ACEKEM<group> :: 
ACEKEM(int cofactormode, KDF& K, HASH& H, int Keylen, 
                group& Grp, std::istream& ifile) : 
      KEM(Keylen), CofactorMode(cofactormode), k(K), hash(H), grp(Grp)
{
	ifile >> w >> x >> y >> z;
	grp.readfromfile(ifile, gp);
	grp.readfromfile(ifile, c);
	grp.readfromfile(ifile, d);
	grp.readfromfile(ifile, h);
}


template<class group> void ACEKEM<group> ::
PrettyPrint(std::ostream& tofile)
{
	tofile << "\n-------------";
	tofile << "\nACE-KEM";
	tofile << "\n-------------";
	tofile << "\n\nKdf=";
	k.PrettyPrint(tofile);
	tofile << "\nHash=";
	hash.PrettyPrint(tofile);
	tofile << "\nKeylen=" << keylen;
	tofile << "\nCofactorMode=" << CofactorMode;
	tofile << "\n\n-----\n";
	tofile << "\nGroup=";
	grp.PrettyPrint(tofile);
	tofile << "\n\n-----\n";
	tofile << "\nPublic Key";
	prettyprint("\ng'", gp, tofile);
	prettyprint("\nc", c, tofile);
	prettyprint("\nd", d, tofile);
	prettyprint("\nh", h, tofile);
	tofile << "\n\n-----\n";
	tofile << "\nPrivate Key";
	tofile << "\n\nw = "; prettyprint(tofile, w);
	tofile << "\n\nx = "; prettyprint(tofile, x);
	tofile << "\n\ny = "; prettyprint(tofile, y);
	tofile << "\n\nz = "; prettyprint(tofile, z);
	tofile << "\n\n";
}

 
template<class group> void ACEKEM<group> :: 
writetofile(std::ostream& ofile)
{
	ofile << w << "\n" << x << "\n" << y << "\n" << z << "\n"; 
        grp.writetofile(ofile, gp);  ofile << "\n";
        grp.writetofile(ofile, c);  ofile << "\n";
        grp.writetofile(ofile, d);  ofile << "\n";
        grp.writetofile(ofile, h);  ofile << "\n";
}

template<class group> ostring ACEKEM<group> :: 
encrypt(ostring& C0)
{
	typedef typename group::GrpElt GrpElt;
	const GrpElt& g = grp.g;
	const NTL::ZZ& mu = grp.mu;

	NTL::ZZ r, R, alpha;
	GrpElt U, V, Up, hp;
	ostring K, EU, EUp, EV, PEH;

	r = RandomBnd(mu);

	U = grp.mul(r, g);
	Up = grp.mul(r, gp);
	hp = grp.mul(r, h);

	if (logfile)
	{
		*logfile << "\n-------------------------------";
		*logfile << "\nTrace for ACE-KEM encrypt";
		*logfile << "\n-------------------------------";
		grp.print_fmt(*logfile);
		*logfile << "\n\nr = "; prettyprint(*logfile, r);
		prettyprint("\nu", U, *logfile);
		prettyprint("\nu'", Up, *logfile);
		prettyprint("\nh~", hp, *logfile);
	}

	EU = grp.encode(U); 
	EUp = grp.encode(Up);

	alpha = os2ip(hash.eval(EU + EUp));

	R = (alpha * r) % mu;
	V = grp.add(grp.mul(r, c), grp.mul(R, d));

	if (logfile)
	{
		*logfile << "\n\nEU = "; prettyprint(*logfile, EU);
		*logfile << "\n\nEU' = "; prettyprint(*logfile, EUp);
		*logfile << "\n\nalpha = "; prettyprint(*logfile, alpha);
		*logfile << "\n\nr' = "; prettyprint(*logfile, R);
		prettyprint("\nv", V, *logfile);
	}

	EV = grp.encode(V);
	PEH = grp.pencode(hp);

	C0 = EU + EUp + EV;
	K = k.eval(EU + PEH, keylen);

	if (logfile)
	{
		*logfile << "\n\nEV = "; prettyprint(*logfile, EV);
		*logfile << "\n\nPEH = "; prettyprint(*logfile, PEH);
		*logfile << "\n\nC0 = "; prettyprint(*logfile, C0);
		*logfile << "\n\nK = "; prettyprint(*logfile, K);
		*logfile << "\n\n";
	}

	return K;
}


template<class group> ostring ACEKEM<group> :: 
decrypt(const ostring& C, int& nbytes)
{
	typedef typename group::GrpElt GrpElt;
	const GrpElt& g = grp.g;
	const NTL::ZZ& mu = grp.mu;
	const NTL::ZZ& nu = grp.nu;

	NTL::ZZ W, X, Y, Z, alpha, t, nu_inv, u;
	GrpElt U, V, Up, Ue, H;

	ostring PEH, K, EU, EUp;

	unsigned int fmt1, fmt2, fmt3;

	int nbtotal = 0;
	int nb;

	nb = grp.decode1(U, C, nbtotal, fmt1);
	if (!nb) throw Decrypt_error();
	EU = C.substr(nbtotal, nb);
	nbtotal += nb;

	nb = grp.decode1(Up, C, nbtotal, fmt2);
	if (!nb) throw Decrypt_error();
	EUp = C.substr(nbtotal, nb);
	nbtotal += nb;
	
	nb = grp.decode1(V, C, nbtotal, fmt3);
	if (!nb) throw Decrypt_error();
	nbtotal += nb;

	if (!consistent(fmt1, fmt2, fmt3))
		throw Decrypt_error();

	if(CofactorMode == 1)
	{
		Ue = grp.mul(nu, U);

		// in a more efficient implementation, nu_inv, W, X, Y, Z
		// could be precomputed

		rem(u, nu, mu);    
                nu_inv = InvMod(u, mu);

		W = MulMod(nu_inv, w, mu);
		X = MulMod(nu_inv, x, mu);
		Y = MulMod(nu_inv, y, mu);
		Z = MulMod(nu_inv, z, mu);
	}
	else
	{
		Ue = U;
		W = w;
		X = x;
		Y = y;
		Z = z;
	}

	if(CofactorMode != 1)
	{
		if(!grp.subgroup(U))
			throw Decrypt_error();
	}

	alpha = os2ip(hash.eval(EU + EUp));

	t = (X + Y*alpha) % mu;

	if( !grp.equal(Up, grp.mul(W, Ue)) || !grp.equal(V, grp.mul(t, Ue)) )
		throw Decrypt_error();

	H = grp.mul(Z, Ue);

	PEH = grp.pencode(H);

	K = k.eval(EU + PEH, keylen);

	nbytes = nbtotal;
	return K;
}

