/* Copyright 2020 The FedLearn Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package com.jdt.fedlearn.core.encryption.distributedPaillier;

import com.jdt.fedlearn.core.encryption.distributedPaillier.DistributedPaillier.DistPaillierPrivkey;
import com.jdt.fedlearn.core.encryption.distributedPaillier.DistributedPaillier.DistPaillierPubkey;
import com.jdt.fedlearn.core.encryption.distributedPaillier.DistributedPaillierNative.signedByteArray;

import java.nio.charset.StandardCharsets;
import java.util.stream.IntStream;


/**
 * This class provides an interface combining FAKE encryption and Distributed Paillier encryption.
 * When usingFake==true, encryption is simply converting a float/long/double from plaintext to byte.
 * When usingFake==false, real Distributed Paillier encryption is done.
 */
public class HomoEncryptionUtil {

    final static double HOMOEPSILON = 1E-1;

    final boolean usingFake;

    final int n, t, length;
    signedByteArray[] zeros; // 缓存生成的0向量
    DistPaillierPrivkey[] skAll; // 单机版时存储所有party的sk，用于debug和实验。
    DistPaillierPrivkey sk; // private key
    DistPaillierPubkey pk; // public key

    public HomoEncryptionUtil(int n, int length, boolean usingFake) {

        this.usingFake = usingFake; // using_fake==1 --> do not use any encryption
        this.n = n; // total number of parties

        /*
         * About choosing value for t:
         * t is the threshold for Generating Prime Shares during Key Generation.
         * In Paillier Decryption stage, since that shares are multiplied once, we need at least 2*t+1 parties
         * to reconstruct secret and decrypt. So, 2*t+1 is the SMALLEST number of party joined during decryption.
         * In federated learning, all parties want to join decryption, so we set t = (n-1)/2.
         */
        this.t = (n-1)/2;

        /*
         * About choosing value for length:
         * length is bit length for the Prime number generated by each client during Key Generation stage.
         * Since that shares are multiplied once, the generated Paillier Key length will be 2*length,
         * i.e. if we want a key length of 1024 for Paillier, we need to set length = 512 here.
         */
        this.length = length/2;

        this.skAll = new DistPaillierPrivkey[n];
        for(int i = 0; i < skAll.length; i++) {
            skAll[i] = new DistPaillierPrivkey();
        }
        this.sk = new DistPaillierPrivkey();
        this.pk = new DistPaillierPubkey();
    }

    // init from debugUtil, used for debug
    public void initKeyFromDebugUtil(HomoEncryptionDebugUtil dec_helper){
        this.setPk(dec_helper.pk);
        this.setSkAll(dec_helper.sk_lst);
    }

    /*
     * ========================================
     *  Encryption APIs
     * ========================================
     */
    public void generateKeyStandalone() {
        if (!usingFake) {
            DistributedPaillier.genPrivpubKeysStandalone(skAll, pk, length, t, n);
        } else {
            for(int i = 0; i < skAll.length; i++) {
                skAll[i] = new DistPaillierPrivkey();
            }
        }
    }

    public void setPk(String pk_str) {
        this.pk.parseJson(pk_str);
    }

    public void setSk(String sk_str) {
        this.sk.parseJson(sk_str);
    }


    public void setZeros(signedByteArray[] zeros) {
        this.zeros = zeros;
    }

    public void setSkAll(DistPaillierPrivkey[] skAll) {
        this.skAll = skAll;
    }

    public void setSk(DistPaillierPrivkey sk) {
        this.sk = sk;
    }

    public void setPk(DistPaillierPubkey pk) {
        this.pk = pk;
    }

    public signedByteArray encryption(double data, DistPaillierPubkey pubKey) {
        if (!usingFake) {
            return DistributedPaillier.encDouble(data, pubKey);
        } else {
            return toSignedByteArr(data);
        }
    }
    public signedByteArray encryption(long data, DistPaillierPubkey pubKey) {
        if (!usingFake) {
            return DistributedPaillier.enc(data, pubKey);
        } else {
            return toSignedByteArr(data);
        }
    }

    public signedByteArray[] encryption(double[] data, DistPaillierPubkey pubKey) {
        if (!usingFake) {
            return DistributedPaillier.encDoubleList(data, pubKey);
        } else {
            return toSignedByteArrLst(data);
        }
    }

    public signedByteArray[] encryption(long[] data, DistPaillierPubkey pubKey) {
        if (!usingFake) {
            return DistributedPaillier.enc_list(data, pubKey);
        } else {
            return toSignedByteArrLst(data);
        }
    }

    public signedByteArray add(signedByteArray a, signedByteArray b, DistPaillierPubkey pubKey) {
        if (!usingFake) {
            return DistributedPaillier.add(a, b, pubKey);
        } else {
            double a_d = toDouble(a);
            double b_d = toDouble(b);
            return toSignedByteArr(a_d+b_d);
        }
    }

    public signedByteArray[] add(signedByteArray[] a, signedByteArray[] b, DistPaillierPubkey pubKey) {
        assert (a.length == b.length);
        if (!usingFake) {
            return DistributedPaillier.add(a, b, pubKey);
        } else {
            double[] a_d = toDoubleLst(a);
            double[] b_d = toDoubleLst(b);

            double[] res = IntStream.range(0, a_d.length).mapToDouble(i -> a_d[i]+b_d[i]).toArray();

            return toSignedByteArrLst(res);
        }
    }

    public signedByteArray[] add(signedByteArray[] a, double[] b, DistPaillierPubkey pubKey) {
        assert (a.length == b.length);
        if (!usingFake) {
            signedByteArray[] b_enc = DistributedPaillier.encDoubleList(b, pubKey);
            return DistributedPaillier.add(a, b_enc, pubKey);
        } else {
            double[] a_d = toDoubleLst(a);
            double[] res = IntStream.range(0, a_d.length).mapToDouble(i -> a_d[i]+b[i]).toArray();
            return toSignedByteArrLst(res);
        }
    }


    public signedByteArray[] mul(signedByteArray[] a, double[] b, DistPaillierPubkey pubKey) {
        assert (a.length == b.length);
        if (!usingFake) {
            return DistributedPaillier.mulDoubleElementwise(a, b, pubKey);

        } else {
            double[] a_d = toDoubleLst(a);
            double[] res = IntStream.range(0, a_d.length).mapToDouble(i -> a_d[i]*b[i]).toArray();

            return toSignedByteArrLst(res);
        }
    }

    public signedByteArray[] mul(signedByteArray[] a, long[] b, DistPaillierPubkey pubKey) {
        assert (a.length == b.length);
        if (!usingFake) {
            return DistributedPaillier.mulLongElementwise(a, b, pubKey);
        } else {
            double[] a_d = toDoubleLst(a);
            double[] res = IntStream.range(0, a_d.length).mapToDouble(i -> a_d[i]*b[i]).toArray();

            return toSignedByteArrLst(res);
        }
    }

    public signedByteArray mul(signedByteArray a, double b, DistPaillierPubkey pubKey) {
        if (!usingFake) {
            return  DistributedPaillier.mulDouble(a, b, pubKey);
        } else {
            String str = new String(a.byteArr, StandardCharsets.UTF_8);
            double a_d = Double. parseDouble(str);
            return toSignedByteArr(a_d*b);
        }
    }

    public signedByteArray[] div(signedByteArray[] a, long[] b, DistPaillierPubkey pubKey) {
        assert (a.length == b.length);
        if (!usingFake) {
            return DistributedPaillier.divLongElementwise(a, b, pubKey);
        } else {
            double[] a_d = toDoubleLst(a);
            double[] res = IntStream.range(0, a_d.length).mapToDouble(i -> a_d[i]/b[i]).toArray();

            return toSignedByteArrLst(res);
        }
    }

    public signedByteArray innerProduct(signedByteArray[] a, double[] b, DistPaillierPubkey pubKey) {
        assert (a.length == b.length);
        if (!usingFake) {
            return DistributedPaillier.innerProductDouble(a, b, pubKey);
        } else {
            double[] a_d = toDoubleLst(a);
            double res = IntStream.range(0, a_d.length).mapToDouble(i -> a_d[i]*b[i]).sum();
            return toSignedByteArr(res);
        }
    }

    public signedByteArray innerProduct(signedByteArray[] a, long[] b, DistPaillierPubkey pubKey) {
        assert (a.length == b.length);
        if (!usingFake) {
            return DistributedPaillier.innerProductLong(a, b, pubKey);
        } else {
            double[] a_d = toDoubleLst(a);
            double res = IntStream.range(0, a_d.length).mapToDouble(i -> a_d[i]*b[i]).sum();
            return toSignedByteArr(res);
        }
    }

    /**
     * First-time Decryption
     * @param data: 密文数组
     * @param privkey:
     */
    public  signedByteArray[] decryptPartial(signedByteArray[] data,
                                             DistPaillierPrivkey privkey) {
        if (!usingFake) {
            assert( (privkey.getRank()>0) && (privkey.getRank()<= this.n) );
            return  DistributedPaillier.decLstPartial(data, privkey.getRank(), privkey);
        } else {
            return data;
        }
    }

    /**
     * Final Decryption
     * @param im_res: decrypt_partial函数输出的解密中间结果, 二维数组 shape = (cypherText.length, number of parties)
     * @param cypherText: 密文数组
     */
    public double decryptFinal(signedByteArray[] im_res,
                               signedByteArray cypherText,
                               DistPaillierPrivkey privkey) {
        if (!usingFake) {
            return  DistributedPaillier.decFinalDouble(im_res, cypherText, privkey);
        } else {
            return toDouble(cypherText);
        }
    }

    /**
     * NOTE that im_res is in shape (n, number_of_instances)
     */
    public long[] decryptFinalLong(signedByteArray[][] im_res,
                                   signedByteArray[] cypherText,
                                   DistPaillierPrivkey privkey) {
        if (!usingFake) {
            long[] ret = new long[cypherText.length];
            for(int i = 0; i < cypherText.length; i++) {
                signedByteArray[] im_res_tmp = new signedByteArray[n];
                for(int j = 0; j < n; j++ ){
                    im_res_tmp[j] = im_res[j][i];
                }
                ret[i] =  DistributedPaillier.decFinalLong(im_res_tmp, cypherText[i], privkey);
            }
            return ret;
        } else {
            return toLongLst(cypherText);
        }
    }

    public double[] decryptFinal(signedByteArray[][] im_res,
                                 signedByteArray[] cypherText,
                                 DistPaillierPrivkey privkey) {
        if (!usingFake) {
            double[] ret = new double[cypherText.length];
            for(int i = 0; i < cypherText.length; i++) {
                signedByteArray[] im_res_tmp = new signedByteArray[n];
                for(int j = 0; j < n; j++ ){
                    im_res_tmp[j] = im_res[j][i];
                }
                ret[i] =  DistributedPaillier.decFinalDouble(im_res_tmp, cypherText[i], privkey);
            }
            return ret;
        } else {
            return toDoubleLst(cypherText);
        }
    }

    /*
     * ========================================
     *  Some Utils
     * ========================================
     */

    /**
     * 从src 深拷贝到 target， target必须预先初始化。
     * Parameters are the same as System.arraycopy
     */
    public static void arrayCopy(signedByteArray[] src, int start_src, signedByteArray[] target, int start_target, int len) {
        assert((src.length >= len) && (target.length>=len));
        for(int i = 0; i<len; i++){
            target[start_target+i] = src[start_src+i].deep_copy();
        }
    }

    public  signedByteArray[] getAllZero(int size) {

        signedByteArray[] ret = new signedByteArray[size];
        if(!usingFake) {
            if (zeros != null && zeros.length >= size) {
                arrayCopy(zeros, 0, ret, 0, size); // copy cached value directly and return
            } else {
                double[] tmp = new double[size];
                ret = encryption(tmp, pk);
                zeros = new signedByteArray[size];
                arrayCopy(ret, 0, zeros, 0, size); // 增长缓存的0向量
            }
        } else {
            ret = toSignedByteArrLst(new double[size]);
        }
        return ret;
    }

    public static byte[] toByteArr(double d) {
        return (Double.toString(d)).getBytes(StandardCharsets.UTF_8);
    }
    public static byte[] toByteArr(int d) {
        byte[] ret = new byte[1];
        ret[0] = (byte)d;
        return ret;
    }


    public static double toDouble(signedByteArray data) {
        String str = new String(data.byteArr, StandardCharsets.UTF_8);
        return Double.parseDouble(str);
    }

    public static double[] toDoubleLst(signedByteArray[] data) {
        double[] ret = new double[data.length];
        for(int i = 0; i< data.length; i++) {
            String str = new String(data[i].byteArr, StandardCharsets.UTF_8);
            ret[i] = Double. parseDouble(str);
        }
        return ret;
    }

    public static long toLong(signedByteArray data) {
        String str = new String(data.byteArr, StandardCharsets.UTF_8);
        return (long)Double. parseDouble(str);
    }

    public static long[] toLongLst(signedByteArray[] data) {
        long[] ret = new long[data.length];
        for(int i = 0; i< data.length; i++) {
            ret[i] = toLong(data[i]);
        }
        return ret;
    }

    public static int[] toIntLst(signedByteArray[] data) {
        int[] ret = new int[data.length];
        for(int i = 0; i< data.length; i++) {
            ret[i] = data[i].byteArr[0];
        }
        return ret;
    }

    public static signedByteArray toSignedByteArr(double data) {
        return new signedByteArray(toByteArr(data), false);
    }

    public static signedByteArray[] toSignedByteArrLst(double[] data) {
        signedByteArray[] ret = new signedByteArray[data.length];
        for(int i = 0; i< data.length; i++) {
            ret[i] = new signedByteArray(toByteArr(data[i]), false);
        }
        return ret;
    }

    public static signedByteArray[] toSignedByteArrLst(int[] data) {
        signedByteArray[] ret = new signedByteArray[data.length];
        for(int i = 0; i< data.length; i++) {
            ret[i] = new signedByteArray(toByteArr(data[i]), false);
        }
        return ret;
    }

    public static signedByteArray[] toSignedByteArrLst(long[] data) {
        signedByteArray[] ret = new signedByteArray[data.length];
        for(int i = 0; i< data.length; i++) {
            ret[i] = new signedByteArray(toByteArr(data[i]), false);
        }
        return ret;
    }

    public static boolean anyPrecisionEqual(double v, double i, double precision) {
        return Math.abs(i - v) < precision;
    }

    public static boolean listAnyPrecisionEqual(double[] v, double[] i, double precision) {
        if(v.length!=i.length){
            return false;
        }
        for(int j = 0; j < i.length; j++) {
            if( !anyPrecisionEqual(v[j], i[j], precision)) {
                return false;
            }
        }
        return true;
    }

    public static boolean listAnyPrecisionEqual(long[] v, double[] i, double precision) {
        if(v.length!=i.length){
            return false;
        }
        for(int j = 0; j < i.length; j++) {
            if( !anyPrecisionEqual(v[j], i[j], precision)) {
                return false;
            }
        }
        return true;
    }

    public static boolean doubleEqual(double v, double i) {
        return Math.abs(i - v) < HOMOEPSILON;
    }

    public static boolean doubleLstEqual(double[] v, double[] i) {
        if(v.length!=i.length){
            return false;
        }
        for(int j = 0; j < i.length; j++) {
            if( !doubleEqual(v[j], i[j])) {
                return false;
            }
        }
        return true;
    }


    public boolean isUsingFake() {
        return usingFake;
    }

    public int getN() {
        return n;
    }

    public int getT() {
        return t;
    }

    public int getLength() {
        return length;
    }

    public DistPaillierPrivkey[] getSkAll() {
        return skAll;
    }

    public DistPaillierPubkey getPk() {
        return pk;
    }

    public DistPaillierPrivkey getSk() {
        return sk;
    }

    public static void main(String[] args) {
        HomoEncryptionUtil test = new HomoEncryptionUtil(3, 1024, false);
        test.generateKeyStandalone();
        int N = 20;
        double [] a = new double[N];
        double [] b = new double[N];
        for(int i = 0; i< N; i++) {
            a[i] = i;
            b[i] = i;
        }
        signedByteArray[] a_c = test.encryption(a, test.getPk());
        signedByteArray[] b_c = test.encryption(b, test.getPk());
        test.add(a_c, b_c, test.getPk());
        test.innerProduct(a_c, b, test.getPk());
    }
}
