#include "ffi.h"
#include "ffi_field.h"
#include "ffi_elt.h"
#include "ffi_vec.h"
#include "ffi_mat.h"
#include "rng.h"
#include "parameters.h"
#include "gabidulin.h"
#include "parsing.h"
#include "piglet1.h"

//全局参数
unsigned char Public_Key_Seed[PK_SEED_LEN];//公钥种子
unsigned char Public_Key[PUBLIC_KEY_BYTES];//公钥
unsigned char Gabidulin_Seed[GD_SEED_LEN];//Gabidulin code 种子

unsigned char Secret_Key_Seed[SK_SEED_LEN];//私钥种子
unsigned char Secret_Key[SECRET_KEY_BYTES];//私钥

/** 
* Keygen of the piglet1 IND-CPA scheme.
* The public key is composed of the matrix Q as well as two seeds for vector g and matrix H respectively.
* The secret key is composed of two seeds for matrix X and vector g respectively.
*
* \param[out] pk String containing the public key
* \param[out] sk String containing the secret key
*/
void piglet1_pke_keygen(unsigned char *pk, unsigned char *sk)
{
	unsigned char pk_seed[PK_SEED_LEN];
    randombytes(pk_seed, PK_SEED_LEN);
	AES_XOF_struct* pk_seedexpander = (AES_XOF_struct*) malloc(sizeof(AES_XOF_struct));
	seedexpander_init(pk_seedexpander, pk_seed, pk_seed + 32, SEEDEXPANDER_MAX_LENGTH);
	
	unsigned char sk_seed[SK_SEED_LEN];
	randombytes(sk_seed, SK_SEED_LEN);
	AES_XOF_struct* sk_seedexpander = (AES_XOF_struct*) malloc(sizeof(AES_XOF_struct));
    seedexpander_init(sk_seedexpander, sk_seed, sk_seed + 32, SEEDEXPANDER_MAX_LENGTH);
	// Compute secret key X, Y
	ffi_vec F, X[PARAM_K], Y[PARAM_K];
	ffi_vec_set_random_full_rank_with_one_using_seedexpander(F, PARAM_W, sk_seedexpander);
	for(int i = 0; i < PARAM_K; i++)
    {
		ffi_vec_set_random_from_support_using_seedexpander(X[i], PARAM_R, F, PARAM_W, sk_seedexpander);
		ffi_vec_set_random_from_support_using_seedexpander(Y[i], PARAM_R, F, PARAM_W, sk_seedexpander);
    }
	// Compute public key
	ffi_vec H[PARAM_K*PARAM_K], s[PARAM_K];
	// set dimension (length) of s and H
	for(int i = 0; i < PARAM_K; i++)
	{
		s[i].SetLength(PARAM_R);
	}
	for(int i = 0; i < PARAM_K; i++)
		for(int j = 0; j < PARAM_K; j++)
			H[i*PARAM_K + j].SetLength(PARAM_R);

	ffi_vec sum, tmp;
	// generate random H
	for(int i = 0; i < PARAM_K; i++)
	{
		for(int j = 0; j < PARAM_K; j++)
		{
			ffi_vec_set_random_using_seedexpander(H[i*PARAM_K + j], PARAM_R, pk_seedexpander); 
		}
	}
	// s = xH + y
	for(int i = 0; i < PARAM_K; i++)
	{
	  ffi_vec_set_zero(tmp, PARAM_R);
	  ffi_vec_set_zero(sum, PARAM_R);
	  for(int j = 0; j < PARAM_K; j++)
	  {
		  ffi_vec_mul(tmp, X[j], H[j*PARAM_K + i], PARAM_R);
		  ffi_vec_add(sum, tmp, sum, PARAM_R);
	  }
	  ffi_vec_add(s[i], sum, Y[i], PARAM_R);
	}

	// Set public g_seed for generate matrix
    unsigned char g_seed[GD_SEED_LEN];
    randombytes(g_seed, GD_SEED_LEN);

	// Parse keys to string
    piglet1_public_key_to_string(pk, pk_seed, s, g_seed);
    piglet1_secret_key_to_string(sk, sk_seed, g_seed, pk);

	free(pk_seedexpander);
	free(sk_seedexpander);
}

/** 
 * \fn void piglet1_pke_encryption(const ffi_vec& m, const unsigned char *pk, unsigned char* theta, ffi_mat& C1, ffi_mat& C2)
 * \brief Encryption of the piglet1_PKE IND-CPA scheme
 *
 * The ciphertext is composed of the matrices C1 and C2.
 *
 * \param[out] C1 Matrix C1 (first part of the ciphertext)
 * \param[out] C2 Matrix C2 (second part of the ciphertext)
 * \param[in] m Vector representing the message to encrypt
 * \param[in] theta Seed used to derive randomness required for encryption
 * \param[in] pk String containing the public key
 */
void piglet1_pke_encryption(const ffi_vec& m, 
							const unsigned char *pk, 
							unsigned char* theta, 
							ffi_vec* u, 
							ffi_vec& v) 
{
	// Create seed_expander from theta
    AES_XOF_struct* seedexpander = (AES_XOF_struct*) malloc(sizeof(AES_XOF_struct));
    seedexpander_init(seedexpander, theta, theta + 32, SEEDEXPANDER_MAX_LENGTH);

	// Retrieve h and s from public key
	ffi_vec H[PARAM_K*PARAM_K], s[PARAM_K];
	for(int i = 0; i < PARAM_K; i++)
	{
		s[i].SetLength(PARAM_R);
	}
	for (int i = 0; i < PARAM_K; i++)
	{
		for (int j = 0; j < PARAM_K; j++)
		{
			H[i*PARAM_K + j].SetLength(PARAM_R);
		}
	}

	unsigned char g_seed[GD_SEED_LEN];
	piglet1_public_key_from_string(H, s, g_seed, pk);

    // Generate e, e1 and r
    ffi_vec support_we, E[PARAM_K], R[PARAM_K], E1;
    //ffi_vec_set_random_full_rank_with_one_using_seedexpander(support_we, PARAM_WE, seedexpander);
	ffi_vec_set_random_full_rank_using_shake256(support_we, PARAM_WE, theta, GD_SEED_LEN);
	for(int i = 0; i < PARAM_K; i++)
    {
		ffi_vec_set_random_from_support_using_seedexpander(E[i], PARAM_R, support_we, PARAM_WE, seedexpander);
		ffi_vec_set_random_from_support_using_seedexpander(R[i], PARAM_R, support_we, PARAM_WE, seedexpander);
    }
	ffi_vec_set_random_from_support_using_seedexpander(E1, PARAM_R, support_we, PARAM_WE, seedexpander);

    // Compute u = H*r^t + e^t
	ffi_vec sum, tmp;
	for(int i = 0; i < PARAM_K; i++)
	{
		ffi_vec_set_zero(tmp, PARAM_R);
		ffi_vec_set_zero(sum, PARAM_R);
		for(int j = 0; j < PARAM_K; j++)
		{
			ffi_vec_mul(tmp, H[i*PARAM_K + j], R[j], PARAM_R);
			ffi_vec_add(sum, tmp, sum, PARAM_R);
		}
		ffi_vec_add(u[i], sum, E[i], PARAM_R);
	}

	// Compute v = s*r^t + E1 + mG
  	ffi_vec_set_zero(tmp, PARAM_R);
	ffi_vec_set_zero(sum, PARAM_R);
	for(int i = 0; i < PARAM_K; i++)
	{
		ffi_vec_mul(tmp, s[i], R[i], PARAM_R);
		ffi_vec_add(sum, tmp, sum, PARAM_R);	 
	}
	ffi_vec_add(v, sum, E1, PARAM_R);

//	unsigned char gc_seed[SK_SEED_LEN];
//	randombytes(gc_seed, SK_SEED_LEN);
//	AES_XOF_struct* gc_seedexpander = (AES_XOF_struct*)malloc(sizeof(AES_XOF_struct));
//	seedexpander_init(gc_seedexpander, gc_seed, gc_seed + 32, SEEDEXPANDER_MAX_LENGTH);

	// Use g_seed to create g (a Gabidulin code)
	ffi_vec g;
	g.SetLength(PARAM_R);
	ffi_vec_set_random_full_rank_using_shake256(g, PARAM_R, g_seed, GD_SEED_LEN);
//	ffi_vec_set_random_full_rank_with_one_using_seedexpander(g, PARAM_R, gc_seedexpander);

	// Compute tmp_v = mG by encoding the message
    ffi_vec mG;
	mG.SetLength(PARAM_R);

    gabidulin_code *code = gabidulin_code_init(g, PARAM_L, PARAM_R);

	//printf("\ngc in ENCODE------------------------------\n");
	//for (long i = 0; i < PARAM_R; i++) {
	//	printf("\ng[%02d] = ", i);
	//	GF2X::HexOutput = 1;
	//	ffi_elt_print(code->g.at(i));
	//	//printf("\ng[%02d] = ", i);
	//	//GF2X::HexOutput = 0;
	//	//ffi_elt_print(code->g.at(i));
	//	//printf("\ng saved 0x: ");
	//	//for (int j = 0; j < code->g.at(i)._GF2E__rep.xrep.length(); j++) {
	//	//	printf("%08X ", *(code->g.at(i)._GF2E__rep.xrep.elts() + j));
	//	//}
	//}

    gabidulin_code_encode(mG, code, m);
	ffi_vec_add(v, mG, v, PARAM_R);

	gabidulin_code_clear(code);
}

/*
	m 明文
*/
void piglet1_pke_decryption(const unsigned char* sk, ffi_vec& m, const ffi_vec* u, const ffi_vec& v) {
	#ifdef VERBOSE
        printf("\n\n### DECRYPTION of piglet1 ###");
    #endif

    // Retrieve x, y and their support from secret key
	ffi_vec F, X[PARAM_K], Y[PARAM_K];
	F.SetLength(PARAM_W);
	for(int i = 0; i < PARAM_K; i++)
	{
		X[i].SetLength(PARAM_R);
		Y[i].SetLength(PARAM_R);
	}

	ffi_vec g;
	g.SetLength(PARAM_R);

	piglet1_secret_key_from_string(X, Y, g, sk);

	//printf("\ng restored---------------------\n");
	//for (long i = 0; i < PARAM_R; i++) {
	//	printf("\ng[%02d] = ", i);
	//	ffi_elt_print(g.at(i));
	//}

	// Compute v - x*u
	ffi_vec tmp, sum, ec;
	ec.SetLength(PARAM_R);
	ffi_vec_set_zero(tmp, PARAM_R);
	ffi_vec_set_zero(sum, PARAM_R);
	for(int i = 0; i < PARAM_K; i++)
	{
		ffi_vec_mul(tmp, X[i], u[i], PARAM_R);
		ffi_vec_add(sum, tmp, sum, PARAM_R);	 
	}
	ffi_vec_add(ec, v, sum, PARAM_R);

	// Decode from dec_mG and store the result in dec_m
    gabidulin_code *code = gabidulin_code_init(g, PARAM_L, PARAM_R);

	//printf("\ngc in DECODE------------------------------\n");
	//for (long i = 0; i < PARAM_R; i++) {
	//	printf("\ng[%02d] = ", i);
	//	//GF2X::HexOutput = 1;
	//	ffi_elt_print(code->g.at(i));
	//}

    gabidulin_code_decode(m, code, ec);

    gabidulin_code_clear(code);
}