
#include "mat.h"
#include "limits.h"
#define M 4
#define N 4

int weight;


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

Mat* bitReprMat(Mat *mat, int l);
Mat* bitReprVec(Mat *mat, int l);

Mat* keySwitch(Mat *S, Mat *c); // returns M (key switch matrix)

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

Mat *serverAddition(Mat *a, Mat *b);
Mat *serverLinearTransform(Mat *mat, Mat *c);
Mat *serverInnerProduct(Mat *c1, Mat *c2, Mat *mat);


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 *serverInnerProduct(Mat *c1, Mat *c2, Mat *mat) {
    return multiply(mat, multiply(c1, scalermultiply(transpose(c2), 0.5)));
}

Mat *serverAddition(Mat *a, Mat *b){
    return sum(a, b);
}

Mat *serverLinearTransform(Mat *mat, Mat *c){
    return multiply(mat, c);
}

Mat* bitReprMat(Mat *mat, int l) {
    Mat *b = newmat(mat->row, mat->col*l, 0);
    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++) {
            int n = (int) get(mat, i+1, j+1);
            for (int k = l - 1; k >= 0; k--) {
                set(b, i + 1, num*l + bid, pow(2, k)*n);
                bid++;
            }
            num++;
            bid = 1;
        }
        num = 0;
        bid = 1;
    }
    return b;
}

Mat* bitReprVec(Mat *mat, int l) {
    Mat *b = newmat(mat->row * l, 1, 0);
    showmat(b);
    bool ispos = true;
    unsigned id;
    int bid = 1;

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

        if (n >= 0) ispos = true;
        else {
            ispos = false;
            n = abs(n);
        }
        for (id = 1 << (l-1); id > 0; id = id / 2) {
            if (n & id) {
                if (ispos) {
                    set(b, bid, 1, 1);
                } else {
                    set(b, bid, 1, -1);
                }
            } else {
                set(b, bid, 1, 0);
            }
            bid++;
        }
    }
//    showmat(b);
    return b;
}

Mat* keySwitch(Mat *S, Mat *c) {
    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);
    Mat *sstar = bitReprMat(S, l);

    int nprime = N * l;
    Mat *T = randm(M, nprime, 0, 1);
    Mat *I = eye(M);

    Mat *sprime = hconcat(I, T);
    Mat *A = randm(nprime-M, nprime, 0, 1000);
    Mat *E = randm(M, nprime, -10, 10);
    Mat *KSM = vconcat(sum(minus(sstar, multiply(T, A)), E), A);
    Mat *cprime = serverLinearTransform(KSM, c);

    return KSM;
}

Mat* encryption(Mat *S, Mat *x) {
    Mat *e = randm(M, 1, -1, 1); // error term
    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;
}

int main() {
    srand(time(0));
    weight = rand() % 10;

    Mat *I = eye(M);
    Mat *S = scalermultiply(I, weight);
    Mat *x_1 = randm(M, 1, INT_MIN>>26, INT_MAX>>26); // initialized variable range
    Mat *x_2 = randm(M, 1, INT_MIN>>26, INT_MAX>>26); // initialized variable range
    printf("********** Client: Plaintext 1 *************\n");
    showmat(x_1);
    printf("********** Client: Plaintext 2 *************\n");
    showmat(x_2);
    Mat *c_1 = encryption(S, x_1); // ciphertext
    Mat *c_2 = encryption(S, x_2); // ciphertext
    printf("********** Client -> Server: Encrypted Text 1 *************\n");
    showmat(c_1);
    printf("********** Client -> Server: Encrypted Text 2 *************\n");
    showmat(c_2);
    printf("********** Server: Encrypted Addition *************\n");
    Mat *s_1 = serverAddition(c_1, c_2);
    showmat(s_1);
    printf("********** Server->Client: Decrypted Addition *************\n");
    Mat *addi = serverAddition(x_1, x_2);
    showmat(addi);
    printf("************** Linear Transformation ****************\n");
    printf("********** Clinet: Linear Transformation Key Switching Matrix *************\n");
    Mat *G = randm(M, N, INT_MIN>>24, INT_MAX>>24);
    Mat *KSM = clientLinearTransform(G, S, c_1);
    showmat(KSM);
    printf("********** Server: Encrypted Linear Transform, (M_l, c_1) *************\n");
    Mat *s_2 = serverLinearTransform(KSM, c_1);
    showmat(s_2);
    printf("********** Server->Client: Decrypted Linear Transform *************\n");
    Mat *M_l = randm(M+5, 1, INT_MIN>>24, INT_MAX>>24); // initialized variable range
    Mat *delin = serverLinearTransform(M_l, x_1);
    showmat(delin);
    printf("********** Server: Encrypted Weighted Inner Product *************\n");
    Mat *M_i = randm(M, M, INT_MIN>>24, INT_MAX>>24); // initialized variable range
    Mat *s_3 = serverInnerProduct(c_1, transpose(c_2), M_i);
    showmat(s_3);
    printf("********** Server->Client: Decrypted Weighted Inner Product *************\n");
    Mat *dewin = serverInnerProduct(x_1, transpose(x_2), M_i);
    showmat(dewin);
    printf("********** Client: Decrypted Plaintext 1 *************\n");
    Mat *xx_1 = decryption(S, c_1);
    showmat(xx_1);
    printf("********** Client: Decrypted Plaintext 2 *************\n");
    Mat *xx_2 = decryption(S, c_2);
    showmat(xx_2);

}
