#include "sig.h"

const int32_t A[8] = { 0 + ALPHA / 2,ALPHA + ALPHA / 2,2 * ALPHA + ALPHA / 2,3 * ALPHA + ALPHA / 2,4 * ALPHA + ALPHA / 2,5 * ALPHA + ALPHA / 2,6 * ALPHA + ALPHA / 2,7 * ALPHA + ALPHA / 2 };

unsigned long long sig_get_pk_byts()
{
	return (N_BYTES);
}

unsigned long long sig_get_sk_byts()
{
	return (N_BYTES + (N2));
}

unsigned long long sig_get_sn_byts()
{
	return (ULEN + RLEN + L_BYTES);
}


/* B_l: 1024-dim binary string, wighted 44 */
/* input 32 bytes = 256 bits short seed */
/* B_l: 1024-dim binary string, wighted 44 */
/* input 32 bytes = 256 bits short seed */
void hash_to_fixed_weight(int l, unsigned char* short_seed, uint16_t *arr) {
	int i,j, t = 0;
	int16_t f[1024];
	memset(f, 0, sizeof(int16_t) * 1024);
	unsigned char long_seed[168];
	keccak_state state;
	uint16_t pos;
	uint64_t five_bytes_to_four_pos[33];
	uint16_t nonce = (short_seed[0] << 8) | short_seed[1];
	shake128_stream_init(&state, short_seed, nonce);
	int rec = 0;

	while (rec < 44) {
		shake128_squeezeblocks(long_seed, 1, &state);
		for (i = 0; i < 11; i++) 
		{
			five_bytes_to_four_pos[i] = (uint64_t)long_seed[5 * i] | ((uint64_t)long_seed[5 * i + 1] << (8)) | ((uint64_t)long_seed[5 * i + 2] << (16)) | ((uint64_t)long_seed[5 * i + 3] << (24)) | ((uint64_t)long_seed[5 * i + 4] << (32));
			for (j = 0; j < 4; j++)
			{
				pos = five_bytes_to_four_pos[i] & 0x3FFU;
				five_bytes_to_four_pos[i] >>= 10;
				if (f[pos] == 0) 
				{
					f[pos] = 1;
					arr[rec++] = pos;
				}
			}
		}
		if (rec == 44)
			return;
		for (i = 11; i < 33; i++)
		{
			five_bytes_to_four_pos[i] = (uint64_t)long_seed[5 * i] | ((uint64_t)long_seed[5 * i + 1] << (8)) | ((uint64_t)long_seed[5 * i + 2] << (16)) | ((uint64_t)long_seed[5 * i + 3] << (24)) | ((uint64_t)long_seed[5 * i + 4] << (32));
			for (j = 0; j < 4; j++)
			{
				pos = five_bytes_to_four_pos[i] & 0x3FFU;
				five_bytes_to_four_pos[i] >>= 10;
				if (f[pos] == 0)
				{
					f[pos] = 1;
					arr[rec++] = pos;
					if (rec == 44)
						return;
				}
			}
		}

		pos = long_seed[165] | ((long_seed[166] & 3) << 8); // 8+2
		if (f[pos] == 0)
		{
			f[pos] = 1;
			arr[rec++] = pos;
			if (rec == 44)
				return;
		}
		pos = (long_seed[166] >> 2) | ((long_seed[167] & 15) << 6); //6+4
		if (f[pos] == 0)
		{
			f[pos] = 1;
			arr[rec++] = pos;
			if (rec == 44)
				return;
		}

	}

}


//输入 chtohash 生成 固定重量的hash 存储其1的位置
void hash_pack(uint16_t *ad, unsigned char *chtohash, int inputlen) 
{
	unsigned char a_bytes[HASHBIT];
	shake256(a_bytes, HASHBIT, chtohash, inputlen);
	hash_to_fixed_weight(44, a_bytes, ad);
}

int sig_keygen(unsigned char *pk, unsigned long long *pk_byts, unsigned char *sk, unsigned long long *sk_byts)
{
	int i;
	int32_t f[N], g[N], fq[N], h[N];

	//generate an invertible f
	trinary_poly_gen(f);
	while (inv_poly(f, fq) == 1)
		trinary_poly_gen(f);

	//generate g and compute h = f(g+alpha)

	trinary_poly_gen(g);
	trinary2byteArray(sk, g);
	*g = *g + ALPHA;
	poly_ntt_o(g);
	poly_pointwise_o(h, g, fq);
	poly_tobytes(pk, h);
	trinary2byteArray(sk+N4, f);
	for (i = 0; i < N_BYTES; i++)
		*(sk + N2 + i) = *(pk + i);

	*pk_byts = sig_get_pk_byts();
	*sk_byts = sig_get_sk_byts();
	return 0;
}

int sig_sign(unsigned char *sk, unsigned long long sk_byts, const unsigned char *m, unsigned long long mlen, unsigned char *sm, unsigned long long *smlen)
{
	int32_t f[N], g[N], h[N], r[N], w[N], rw[N], a[N], c1[N], c2[N], rcopy[N];
	unsigned char seed[32], k[RLEN], qw_bytes[N2], chtohash[(RLEN)+N2 + MLEN];  
	uint16_t ad[L], qw[N];
	int i,j;

	*smlen = sig_get_sn_byts();

	//recover the polynomial sk and pk
	byteArray2trinary(g, sk);
	byteArray2trinary(f, sk+N4);
	poly_frombytes(h, sk+N2);

	//generate k	
	rand_byts(32, seed);
	shake256(k, RLEN, seed, 32);
	
	//compute the input of the hash (k||m|| .) except the quotient
	
	for (i = 0; i < RLEN; i++)
		chtohash[i] = k[i];
	for (i = 0; i < MLEN; i++)
		chtohash[i + RLEN] = m[i];


pos_1:
    //generate r
	sample_r(r);
	eval_poly(rcopy, r);
	poly_ntt_o(rcopy);
	poly_pointwise_o(w, h, rcopy);
	poly_invntt_o(w);
	for (i = 0; i < N; i++)
	{
		qw[i] = (w[i] - 1) / (ALPHA);
	}

	//generate a
	poly_tobytes_qw(qw_bytes, qw);
	for (i = 0; i < N2; i++)
		chtohash[i + RLEN + MLEN] = qw_bytes[i];
	hash_pack(ad, chtohash, RLEN + N2 + MLEN);



	sparse_times(ad, f, c2, L);

	if (judge(c2, GAMMA)) goto pos_1;
	
	//check af+r
	for (i = 0; i < N; i++)
	{
		c2[i] = c2[i] + r[i];
	}
	if (judge(c2, ALPHA / 2 - GAMMA-1)) goto pos_1;
	
	//check ag
	sparse_times(ad, g, c1, L);
	if (judge(c1, GAMMA)) goto pos_1;

	//check ag+R(w)
	for (i = 0; i < N; i++)
	{
		rw[i] = w[i] - A[qw[i]];
		c1[i] = c1[i] + rw[i];
	}
	if (judge(c1, ALPHA / 2 - GAMMA)) goto pos_1;
	
	//pack signature  (k || a||u)
	for (i = 0; i < RLEN; i++) sm[i] = k[i];
	poly_tobytes_ten(sm+RLEN, ad, L);
	poly_tobytes_u(sm + L_BYTES + RLEN, c2);
	return 0;
}

int sig_verf(unsigned char *pk, unsigned long long pk_byts, unsigned char *sm, unsigned long long smlen, unsigned char *m, unsigned long long mlen)
{
	int32_t  h[N], u[N], w[N];
	uint16_t ad[L], adp[L], qw[N];
	unsigned char  qw_bytes[N_BYTES], hash[(RLEN)+N2 + MLEN];
	int  i;
	
	//load the pk and the sig
	poly_frombytes(h, pk);
	poly_frombytes_ten(ad, sm + RLEN, L);
	poly_frombytes_u(u, sm + (L_BYTES + RLEN));

	//check the norm of u
	if (judge(u, ALPHA / 2 - GAMMA - 1))
	{
		return 0;
	}

	//generate the input of  the hash
	poly_ntt_o(u);
	poly_pointwise_o(w, h, u);
	poly_invntt_o(w);
	for (i = 0; i < L; i++)
	{
		w[ad[i]] = (w[ad[i]] - ALPHA + Q) % Q;

	}
	for (i = 0; i < N; i++)
	{
		qw[i] = (w[i] - 1) / (ALPHA);
	}
	poly_tobytes_qw(qw_bytes, qw);

	for (i = 0; i < RLEN; i++)
		hash[i] = sm[i];
	for (i = 0; i < MLEN; i++)
		hash[i + RLEN] = m[i];
	for (i = 0; i < N2; i++)
		hash[i + RLEN + MLEN] = qw_bytes[i];

	//generate the hash and check
	hash_pack(adp, hash, RLEN + N2 + MLEN);
	
	for (i = 0; i < L; i++)
	{
		if (adp[i] != ad[i]) return 0;
	}
	return 1;
}
