#include "poly.h"


/* Generate a random polynomial
Require: f with a determined number of coefficients equal to 0, -1, and 1(w.r.t N-2*d-1,d,d+1) */
void trinary_poly_gen(int32_t  *f) {

	unsigned char short_seed[32];
	unsigned char long_seed[1176]; // 7*168
	int seed_len = 1176;
	
	int i;
	int pos;
	int  t = 0;
	int32_t keep;
	int d = 256;
	int32_t upper = (1 << 16) - 1;

	keccak_state state;
	rand_byts(32, short_seed);

	uint16_t nonce = (short_seed[0] << 8) | short_seed[1];
	shake128_stream_init(&state, short_seed, nonce);
	shake128_squeezeblocks(long_seed, 7, &state);
	memset(f, 0, sizeof(int32_t) * 1024);

	for (i = 0; i < d + 1; )
	{
		/*//发生概率极低
		if (t >= seed_len) {
			t = 0;
			shake128_squeezeblocks(long_seed, 1, &state);
			seed_len = SHAKE128_RATE; // 168
		}
		*/

		keep = (long_seed[t] | long_seed[t + 1] << 8);
		t = t + 2;

		pos = (upper / (i + 511)) * (i + 511);
		if (keep < pos)
		{
			pos = keep % (i + 511);
			f[i + 511] = f[pos];
			f[pos] = 1;
			i++;
		}
	}

	for (i = 0; i < d;)
	{
		if (t >= seed_len) {
			t = 0;
			shake128_squeezeblocks(long_seed, 1, &state);
			seed_len = SHAKE128_RATE; // 168
		}

		keep = (long_seed[t] | long_seed[t + 1] << 8);
		t = t + 2;

		pos = (upper / (768 + i)) * (768 + i);
		if (keep < pos)
		{
			pos = keep % (768 + i);
			f[768 + i] = f[pos];
			f[pos] = -1;
			i++;
		}
	}

}



/* alpha = (464897 - 1)/8 = 58112 */
/* every coef of a is uniform from {0,1,...,alpha-1} */
/* Log2[58112]=15.8265 */
int sample_rej(int32_t *a, unsigned int len, const unsigned char *buf, unsigned int buflen)
{
	unsigned int ctr, pos;
	int32_t t = 0;
	ctr = pos = 0;
	while (ctr < len && pos + 2 <= buflen) {
		t = buf[pos++];
		t |= buf[pos++] << 8;
		if (t < 58112)
			a[ctr++] = t-(ALPHA>>1);
	}
	return ctr;
}

void sample_r(int32_t *a) {
	int ctr;
	unsigned char short_seed[32];
	unsigned char buf[2184],buf1[168];

	keccak_state state;
	rand_byts(32, short_seed);
	uint16_t nonce = (short_seed[0] << 8) | short_seed[1];
	shake128_stream_init(&state, short_seed, nonce);
	shake128_squeezeblocks(buf, 13, &state);

	ctr = sample_rej(a, N, buf, 2184);
	while (ctr < N) {
		shake128_squeezeblocks(buf1, 1, &state);
		ctr += sample_rej(a + ctr, N - ctr, buf1, 168);
	}
}



/* 数组赋值 */
void eval_poly(int32_t *f, int32_t *g)
{
	int i;
	for (i = 0; i < N; i++) 
		f[i] = g[i];
}

/* Compute the inverse of a positive number mod Q */
int32_t inv_num(int32_t n)
{
	int32_t  u, v, v1, v2, q, r, v0;

	v1 = 0;
	v2 = 1;
	u = Q; v = n;

	while (v != 0) {
		q = u / v;
		r = u % v;
		u = v;
		v = r;
		v0 = v2;
		v2 = v1 - q * v2;
		v1 = v0;
	}
	return v1;
}


void balance(int32_t *f)
{
	int i;
	for (i = 0; i < N; i++)
	{
		f[i] = f[i] - (f[i] / QP)*Q;
	}
}

//判断向量的无穷范数是否大于a
int judge(int32_t *g, int a) {
	int i, ctr = 0;

	for (i = 0; i < N; i++) {
		ctr = ctr ||(g[i]/ (a + 1));
	}
	if (ctr == 1) return 1;
	return 0;
}







/* 求多项式的逆 */
//判断是否可逆  ***********************
int inv_poly(int32_t *f, int32_t *fq)
{
	int ctr;
	int32_t g[N], i,j;


	for (i = 0; i < N; i++) {
		g[i] = f[i];
	}

	poly_ntt_o(g);
	ctr = 0;
	for (i = 0; i < N; i++) {
		ctr = ctr | ((g[i] - 1) >> 20);
	}
	if (ctr != 0) return 1;
	for (i = 0; i < N; i++)
		fq[i] = inv_num(g[i]);
	return 0;
}


void poly_add(int32_t *a, int32_t *b, int32_t *c) {
	for (int i = 0; i < 1024; i++) {
		c[i] = a[i] + b[i];
	}
}

/* f = x^k × g */
void right_shift(int32_t *f, const int32_t *g, int k) {
	int i;
	for (i = 0; i < 1024 - k; i++) {
		f[k + i] = g[i];
	}
	for (i = 0; i < k; i++) {
		f[k - 1 - i] = -g[1024 - 1 - i];
	}
}

/* arr: array of locations of 1 of polynomial a, f is a trinary poly */
void sparse_times(uint16_t *arr, int32_t *f, int32_t *res, int32_t sz) {
	int i;
	memset(res, 0, 1024 * sizeof(int32_t));
	int y[1024];
	for (i = 0; i < sz; i++)
	{
		right_shift(y, f, arr[i]);
		poly_add(y, res, res);
	}
}