#include <stdio.h>
#include "switching_matrix.h"
#include "encrypted_domain.h"
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>
#include <math.h>
#include <stdio.h>

#include "./clinet_server.c"
#include "./key_switching.c"
#include "./key_switching.h"
#include "./server_functions.c"
#include "./mat.h"

/* reset the accelerator */
void key_switching_reset(const key_switching_status_t *c)
{
  key_switching_arg_t ksa;
  ksa.status_info = *c;
  if (ioctl(key_switching_fd, KEY_SWITCHING_RESET_WRITE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_RESET_WRITE) failed");
      return;
  }
}

void key_switching_all_loaded(const key_switching_status_t *c)
{
  key_switching_arg_t ksa;
  ksa.status_info = *c;
  if (ioctl(key_switching_fd, KEY_SWITCHING_RESET_WRITE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_LOADED_WRITE) failed");
      return;
  }

static void key_switching_write_op_type(key_switching_loading_t load_info)
{
  key_switching_arg_t ksa;
  ksa.load_info = *c;
  if (ioctl(key_switching_fd, KEY_SWITCHING_OP_TYPE_WRITE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_OP_TYPE_WRITE) failed");
      return;
  }
}

static void key_switching_write_width(key_switching_loading_t load_info)
{
  key_switching_arg_t ksa;
  ksa.load_info = *c;
  if (ioctl(key_switching_fd, KEY_SWITCHING_WIDTH_WRITE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_WIDTH_WRITE) failed");
      return;
  }
}

static void key_switching_write_length(key_switching_loading_t load_info)
{
  key_switching_arg_t ksa;
  ksa.load_info = *c;
  if (ioctl(key_switching_fd, KEY_SWITCHING_LENGTH_WRITE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_LENGTH_WRITE) failed");
      return;
  }
}

static void key_switching_write_ell(key_switching_loading_t load_info)
{
  key_switching_arg_t ksa;
  ksa.load_info = *c;
  if (ioctl(key_switching_fd, KEY_SWITCHING_ELL_WRITE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_ELL_WRITE) failed");
      return;
  }
}

/* write to the accelerator */
void key_switching_write_data(const key_switching_loading_t *c)
{
  key_switching_arg_t ksa;
  ksa.load_info = *c;
  if (ioctl(key_switching_fd, KEY_SWITCHING_WRITE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_WRITE) failed");
      return;
  }
}

/* read from the accelerator */
int key_switching_read_data(const key_switching_loading_t *c)
{
  key_switching_arg_t ksa;

  if (ioctl(key_switching_fd, KEY_SWITCHING_READ, &ksa)) {
      perror("ioctl(KEY_SWITCHING_READ) failed");
      return;
  } else ksa.output_info.output;
}


static void key_switching_read_length(key_switching_output_t output_info)
{	
  key_switching_arg_t ksa;

  if (ioctl(key_switching_fd, KEY_SWITCHING_READ_LENGTH, &ksa)) {
      perror("ioctl(KEY_SWITCHING_READ_LENGTH) failed");
      return;
  } else ksa.output_info.output;
}

static void key_switching_read_width(key_switching_output_t output_info)
{	
  key_switching_arg_t ksa;

  if (ioctl(key_switching_fd, KEY_SWITCHING_READ_WDITH, &ksa)) {
      perror("ioctl(KEY_SWITCHING_READ_WDITH) failed");
      return;
  } else ksa.output_info.output;
}

static void key_switching_read_done(key_switching_output_t output_info)
{	
  key_switching_arg_t ksa;

  if (ioctl(key_switching_fd, KEY_SWITCHING_READ_DONE, &ksa)) {
      perror("ioctl(KEY_SWITCHING_READ_DONE) failed");
      return;
  } else ksa.output_info.output;
}

Mat* encryption(Mat *S, Mat *x);
Mat* decryption(Mat *S, Mat *c);

Mat *clientLinearTransform(Mat *G, Mat *S, Mat *c);
Mat *clientInnerProduct(Mat *S1, Mat *S2, Mat *H, Mat *c1, Mat *c2);

/*---------------- key-switching accelerator ----------------*/
Mat* keySwitch(Mat *S, Mat *c);
// accelerator part
Mt* bitReprMat(Mat *mat, int l, key_switching_arg_t *ksa);
Mat* bitReprVec(Mat *mat, int l, key_switching_arg_t *ksa);

Mat* getRandomMat(int m, int n, key_switching_arg_t *ksa);
Mat* getNoiseMat(int m, int n, key_switching_arg_t *ksa);
/*---------------- key-switching accelerator ----------------*/

Mat *clientLinearTransform(Mat *G, Mat *S, Mat *c) {
    return keySwitch(multiply(G, S), c);
}

Mat *clientInnerProduct(Mat *S1, Mat *S2, Mat *H, Mat *c1, Mat *c2) {
    Mat *S = transpose(multiply(multiply(transpose(S1), H), S2));
    Mat *c = multiply(c1, scalermultiply(transpose(c2), 1.0/weight));
    return keySwitch(S, c);
}

Mat* bitReprMat(Mat *mat, int l, key_switching_arg_t *ksa) {
    Mat *b = newmat(mat->row, mat->col*l, 0);

    ksa.status_info.reset = 1;
    key_switching_reset(&ksa.status_info);
    ksa.status_info.reset = 0;
    key_switching_reset(&ksa.status_info);
    
    // load operation type, width, length, ell
    // tell hardware to evoke bit_repr_matrix module
    ksa.load_info.op_type = 1;
    key_switching_write_op_type(&ksa.load_info);
    ksa.load_info.width = mat->col;
    key_switching_write_width(&ksa.load_info);
    ksa.load_info.length = mat->row;
    key_switching_write_length(&ksa.load_info);
    ksa.load_info.ell = l;
    key_switching_write_ell(&ksa.load_info);
    // ready
    ksa.status_info.all_loaded = 1;
    key_switching_all_loaded(&ksa.status_info);

    int bid = 1; // shifting index for bit repr
    int num = 0; // offset for each number in mat
    for (int i = 0; i < mat->row; i++) {
        for (int j = 0; j < mat->col; j++) {
            ksa.load_info.input = get(mat, i, j);
            key_switching_write_data(&ksa.load_info);
        }
    }

    for (int i = 0; i < mat->row; i++) {
        for (int j = 0; j < mat->col; j++) {
            for (int k = l - 1; k >= 0; k--) {
                set(b, i + 1, num*l + bid, key_switching_read_data(&ksa.output_info););
                bid++;
            }
            num++;
            bid = 1;
        }
        num = 0;
        bid = 1;
    }

    // stop hardware
    ksa.status_info.all_loaded = 0;
    key_switching_all_loaded(&ksa.status_info);
    return b;
}

Mat* bitReprVec(Mat *mat, int l, key_switching_arg_t *ksa) {
    Mat *b = newmat(mat->row * l, 1, 0);
    
    unsigned id;
    int bid = 1;

    ksa.status_info.reset = 1;
    key_switching_reset(&ksa.status_info);
    ksa.status_info.reset = 0;
    key_switching_reset(&ksa.status_info);
    
    // load operation type, width, ell
    // tell hardware to evoke bit_repr_vector module
    ksa.load_info.op_type = 1;
    key_switching_write_op_type(&ksa.load_info);
    ksa.load_info.width = mat->col;
    key_switching_write_width(&ksa.load_info);
    ksa.load_info.length = mat->row;
    key_switching_write_ell(&ksa.load_info);
    ksa.load_info.ell = l;
    key_switching_write_ell(&ksa.load_info);

    // ready
    ksa.status_info.all_loaded = 1;
    key_switching_all_loaded(&ksa.status_info);

    for (int i = 0; i < mat->row; i++) {
        ksa.load_info.input = get(mat, 1, i);
        key_switching_write_data(&ksa.load_info);
    }

    for (int i = 0; i < mat->row; i++) {
        int n = (int) get(mat, i+1, 1);

        for (id = 1 << (l-1); id > 0; id = id / 2) {
            set(b, 1, bid, key_switching_read_data(&ksa.output_info));
            bid++;
        }
    }

    // stop hardware
    ksa.status_info.all_loaded = 0;
    key_switching_all_loaded(&ksa.status_info);
    return b;
}

Mat* getRandomMat(int r,int c, key_switching_arg_t *ksa){
    Mat* R=newmat(r,c,1);
    int k=0;


    ksa.status_info.reset = 1;
    key_switching_reset(&ksa.status_info);
    ksa.status_info.reset = 0;
    key_switching_reset(&ksa.status_info);
    
    // load operation type, width, length
    // tell hardware to evoke get_random_matrix module
    ksa.load_info.op_type = 2;
    key_switching_write_op_type(&ksa.load_info);
    ksa.load_info.width = c;
    key_switching_write_width(&ksa.load_info);
    ksa.load_info.length = r;
    key_switching_write_length(&ksa.load_info);

    // ready
    ksa.status_info.all_loaded = 1;
    key_switching_all_loaded(&ksa.status_info);

    for(int i=1;i<=r;i++){
        for(int j=1;j<=c;j++){
            int r = key_switching_read_data(&ksa.output_info);
            R->entries[k++] = r;
        }
    }
    // stop hardware
    ksa.status_info.all_loaded = 0;
    key_switching_all_loaded(&ksa.status_info);
    return R;
}

Mat* getNoiseMat(int r,int c, key_switching_arg_t *ksa){
    Mat* R=newmat(r,c,1);
    int k=0;

    ksa.status_info.reset = 1;
    key_switching_reset(&ksa.status_info);
    ksa.status_info.reset = 0;
    key_switching_reset(&ksa.status_info);
    
    // load operation type, width, length
    // tell hardware to evoke get_noise_matrix module
    ksa.load_info.op_type = 3;
    key_switching_write_op_type(&ksa.load_info);
    ksa.load_info.width = c;
    key_switching_write_width(&ksa.load_info);
    ksa.load_info.length = r;
    key_switching_write_length(&ksa.load_info);

    // ready
    ksa.status_info.all_loaded = 1;
    key_switching_all_loaded(&ksa.status_info);

    for(int i=1;i<=r;i++){
        for(int j=1;j<=c;j++){
            int r = key_switching_read_data(&ksa.output_info);
            R->entries[k++] = r;
        }
    }
    // stop hardware
    ksa.status_info.all_loaded = 0;
    key_switching_all_loaded(&ksa.status_info);
    return R;
}

Mat* keySwitch(Mat *S, Mat *c, key_switching_arg_t *ksa) {
    int l = 0;
    double norm_val = norm(c);
    printf("norm: %f\n", norm_val);
    while (pow(2, l) < norm_val) {
        l++;
    }
    printf("l:%d\n", l);
    Mat *cstar = bitReprVec(c, l, &ksa);
    Mat *sstar = bitReprMat(S, l, &ksa);

    int nprime = N * l;
    Mat *T = getRandomMat(M, nprime, &ksa);
    Mat *I = eye(M);

    Mat *sprime = hconcat(I, T);
    Mat *A = getRandomMat(nprime-M, nprime, &ksa);
    Mat *E = getNoiseMat(M, nprime, &ksa);
    Mat *KSM = vconcat(sum(minus(sstar, multiply(T, A)), E), A);
    Mat *cprime = serverLinearTransform(KSM, c);

    return KSM;
}

Mat* encryption(Mat *S, Mat *x, key_switching_arg_t *ksa) {
    Mat *e = getRandomMat(M, 1, &ksa);
    Mat *Sinv = inverse(S);
    x = scalermultiply(x, weight);
    x = sum(x, e);
    Mat *res = multiply(Sinv, x);
    return res;
}

Mat* decryption(Mat *S, Mat *c) {
    Mat *res = multiply(S, c);
    res = scalermultiply(res, 1./weight);
    return res;
}