package main

import (
 "fmt"
 "github.com/tuneinsight/lattigo/v3/ckks"
 "github.com/tuneinsight/lattigo/v3/dckks"
)

func main() {
 // Parameters
 params := ckks.DefaultParams[ckks.PN12QP109]
 // Number of parties
 nbParties := 2
 // Number of values to encrypt
 nbSlots := 16

 // Create the CKKS scheme
 scheme := ckks.NewScheme()
 scheme.SetParams(params)

 // Create the distributed CKKS scheme
 dckksScheme := dckks.NewDCKKS(scheme)

 // Create the collective public key generator
 collectivePublicKeyGen := dckksScheme.NewCKKSCollectivePublicKeyGenerator()

 // Generate the collective public key
 collectivePublicKey := collectivePublicKeyGen.GenCollectiveKey()

 // Create the encryptors for each party
 encryptors := make([]*dckks.CKKSLocalEncryptor, nbParties)
 for i := range encryptors {
  encryptors[i] = dckksScheme.NewCKKSLocalEncryptor(collectivePublicKeyGen.LocalSecretKeys()[i])
 }

 // Create the evaluator
 evaluator := dckksScheme.NewCKKSEvaluator()

 // Create the plaintexts for each party
 values1 := make([]complex128, nbSlots)
 values2 := make([]complex128, nbSlots)
 for i := range values1 {
  values1[i] = complex(float64(i), 0)
  values2[i] = complex(float64(i+1), 0)
 }
 fmt.Println("Values 1:", values1)
 fmt.Println("Values 2:", values2)

 pt1 := scheme.NewPlaintext(params.MaxLevel(), params.Scale())
 pt2 := scheme.NewPlaintext(params.MaxLevel(), params.Scale())
 scheme.EncodeNTT(pt1, values1, nbSlots)
 scheme.EncodeNTT(pt2, values2, nbSlots)

 // Encrypt the plaintexts for each party
 ct1 := encryptors[0].EncryptNew(pt1)
 ct2 := encryptors[1].EncryptNew(pt2)

 // Add the ciphertexts
 ctSum := evaluator.AddNew(ct1, ct2)

 // Create the decryptor for each party
 decryptors := make([]*dckks.CKKSLocalDecryptor, nbParties)
 for i := range decryptors {
  decryptors[i] = dckksScheme.NewCKKSLocalDecryptor(collectivePublicKeyGen.LocalSecretKeys()[i])
 }

 // Decrypt the ciphertexts for each party
 ptShares := make([]*ckks.Plaintext, nbParties)
 for i := range ptShares {
  ptShares[i] = decryptors[i].DecryptNew(ctSum)
 }

 // Recombine the plaintext shares
 ptResult := scheme.NewPlaintext(params.MaxLevel(), params.Scale())
 evaluator.Recombine(ptShares, ptResult)

 // Decode the result
 valuesResult := make([]complex128, nbSlots)
 scheme.DecodeNTT(ptResult, valuesResult, nbSlots)
 fmt.Println("Result:", valuesResult)
}