# Install the Lattice library
go get github.com/luxfi/lattice/v6

# Import required packages
import (
    "github.com/luxfi/lattice/v6/core/rlwe"
    "github.com/luxfi/lattice/v6/schemes/ckks"
    "github.com/luxfi/lattice/v6/schemes/bgv"
    "github.com/luxfi/lattice/v6/multiparty"
)

package main

import (
    "fmt"
    "log"
    "github.com/luxfi/lattice/v6/core/rlwe"
    "github.com/luxfi/lattice/v6/schemes/ckks"
)

func main() {
    // Initialize parameters for 128-bit security
    params, err := ckks.NewParametersFromLiteral(
        ckks.ParametersLiteral{
            LogN:            14,
            LogQ:            []int{50, 40, 40, 40},
            LogP:            []int{60},
            LogDefaultScale: 40,
            Sigma:           3.2,
        },
    )
    if err != nil {
        log.Fatal(err)
    }

    fmt.Printf("Ring dimension: %d\n", params.N())
    fmt.Printf("Ciphertext levels: %d\n", params.MaxLevel())
    fmt.Printf("Precision: %d bits\n", params.LogDefaultScale())
}

func encryptedArithmetic() {
    // Setup
    params, _ := ckks.NewParametersFromLiteral(ckks.PN14QP438)
    kgen := rlwe.NewKeyGenerator(params)
    sk := kgen.GenSecretKeyNew()
    pk := kgen.GenPublicKeyNew(sk)
    rlk := kgen.GenRelinearizationKeyNew(sk)
    evk := rlwe.NewMemEvaluationKeySet(rlk)

    encoder := ckks.NewEncoder(params)
    encryptor := rlwe.NewEncryptor(params, pk)
    decryptor := rlwe.NewDecryptor(params, sk)
    evaluator := ckks.NewEvaluator(params, evk)

    // Prepare data
    x := []complex128{1.5, 2.5, 3.5, 4.5}
    y := []complex128{0.5, 1.5, 2.5, 3.5}

    // Encode and encrypt
    ptX := ckks.NewPlaintext(params, params.MaxLevel())
    ptY := ckks.NewPlaintext(params, params.MaxLevel())
    encoder.Encode(x, ptX)
    encoder.Encode(y, ptY)

    ctX, _ := encryptor.EncryptNew(ptX)
    ctY, _ := encryptor.EncryptNew(ptY)

    // Homomorphic operations
    // Addition: z = x + y
    ctSum, _ := evaluator.AddNew(ctX, ctY)

    // Multiplication: z = x * y
    ctProd, _ := evaluator.MulRelinNew(ctX, ctY)
    evaluator.Rescale(ctProd, ctProd)

    // Square: z = x^2
    ctSquare, _ := evaluator.MulRelinNew(ctX, ctX)
    evaluator.Rescale(ctSquare, ctSquare)

    // Decrypt and decode results
    ptSum := decryptor.DecryptNew(ctSum)
    ptProd := decryptor.DecryptNew(ctProd)
    ptSquare := decryptor.DecryptNew(ctSquare)

    resSum := encoder.Decode(ptSum, params.LogSlots())
    resProd := encoder.Decode(ptProd, params.LogSlots())
    resSquare := encoder.Decode(ptSquare, params.LogSlots())

    fmt.Println("x + y =", resSum[:4])
    fmt.Println("x * y =", resProd[:4])
    fmt.Println("x^2 =", resSquare[:4])
}

func evaluatePolynomial() {
    params, _ := ckks.NewParametersFromLiteral(ckks.PN14QP438)
    // ... (setup as before)

    // Evaluate f(x) = 3x^3 - 2x^2 + x - 5
    coefficients := []complex128{-5, 1, -2, 3}

    // Input values
    x := []complex128{1.0, 2.0, 3.0, 4.0}
    ptX := ckks.NewPlaintext(params, params.MaxLevel())
    encoder.Encode(x, ptX)
    ctX, _ := encryptor.EncryptNew(ptX)

    // Horner's method: f(x) = -5 + x(1 + x(-2 + x(3)))
    result := ckks.NewCiphertext(params, 1, params.MaxLevel())

    // Start with highest coefficient
    encoder.Encode([]complex128{3, 3, 3, 3}, result.Plaintext())
    encryptor.EncryptZero(result)
    evaluator.Add(result, coefficients[3], result)

    for i := 2; i >= 0; i-- {
        // result = result * x + coefficients[i]
        evaluator.MulRelin(result, ctX, result)
        evaluator.Rescale(result, result)

        // Add constant
        constant := make([]complex128, params.Slots())
        for j := range constant {
            constant[j] = coefficients[i]
        }
        ptConst := ckks.NewPlaintext(params, result.Level())
        encoder.EncodeAtLevelNew(constant, result.Level(), ptConst)
        evaluator.Add(result, ptConst, result)
    }

    // Decrypt result
    ptResult := decryptor.DecryptNew(result)
    values := encoder.Decode(ptResult, params.LogSlots())

    fmt.Println("f(x) =", values[:4])
    // Expected: f(1)=-3, f(2)=9, f(3)=49, f(4)=135
}

type EncryptedLinearRegression struct {
    params    ckks.Parameters
    evaluator ckks.Evaluator
    encoder   ckks.Encoder
}

func (lr *EncryptedLinearRegression) Predict(
    ctWeights *ckks.Ciphertext,
    ctFeatures []*ckks.Ciphertext,
) (*ckks.Ciphertext, error) {
    // Compute y = w0 + w1*x1 + w2*x2 + ... + wn*xn

    if len(ctFeatures) == 0 {
        return nil, errors.New("no features provided")
    }

    // Initialize with bias (first weight)
    result := ckks.NewCiphertext(lr.params, 1, ctWeights.Level())
    lr.evaluator.Mul(ctWeights, ctFeatures[0], result)

    // Add weighted features
    for i := 1; i < len(ctFeatures); i++ {
        temp := ckks.NewCiphertext(lr.params, 1, ctWeights.Level())

        // Rotate weights to align with next feature
        rotatedWeights := lr.evaluator.RotateNew(ctWeights, i)

        // Multiply and add
        lr.evaluator.Mul(rotatedWeights, ctFeatures[i], temp)
        lr.evaluator.Add(result, temp, result)
    }

    // Rescale for precision management
    lr.evaluator.Rescale(result, result)

    return result, nil
}

// Training encrypted model using gradient descent
func (lr *EncryptedLinearRegression) TrainStep(
    ctWeights *ckks.Ciphertext,
    ctFeatures []*ckks.Ciphertext,
    ctLabels *ckks.Ciphertext,
    learningRate float64,
) (*ckks.Ciphertext, error) {
    // Predict with current weights
    predictions, _ := lr.Predict(ctWeights, ctFeatures)

    // Compute error: error = predictions - labels
    error := lr.evaluator.SubNew(predictions, ctLabels)

    // Compute gradients
    gradients := make([]*ckks.Ciphertext, len(ctFeatures))
    for i, feature := range ctFeatures {
        // gradient[i] = mean(error * feature[i])
        gradients[i] = lr.evaluator.MulNew(error, feature)
        lr.evaluator.Rescale(gradients[i], gradients[i])
    }

    // Update weights: w = w - learningRate * gradient
    ptLR := ckks.NewPlaintext(lr.params, ctWeights.Level())
    lr.encoder.Encode([]complex128{complex(learningRate, 0)}, ptLR)

    for i, grad := range gradients {
        scaledGrad := lr.evaluator.MulNew(grad, ptLR)
        lr.evaluator.Rescale(scaledGrad, scaledGrad)

        // Update specific weight
        rotatedWeights := lr.evaluator.RotateNew(ctWeights, i)
        lr.evaluator.Sub(rotatedWeights, scaledGrad, rotatedWeights)

        // Rotate back
        ctWeights = lr.evaluator.RotateNew(rotatedWeights, -i)
    }

    return ctWeights, nil
}

type EncryptedNeuralNetwork struct {
    layers []EncryptedLayer
}

type EncryptedLayer interface {
    Forward(input *ckks.Ciphertext) (*ckks.Ciphertext, error)
}

// Fully connected layer
type EncryptedDenseLayer struct {
    weights *ckks.Ciphertext
    bias    *ckks.Ciphertext
    eval    ckks.Evaluator
}

func (l *EncryptedDenseLayer) Forward(input *ckks.Ciphertext) (*ckks.Ciphertext, error) {
    // Compute output = weights @ input + bias
    output := l.eval.MulRelinNew(l.weights, input)
    l.eval.Rescale(output, output)
    l.eval.Add(output, l.bias, output)
    return output, nil
}

// ReLU activation using polynomial approximation
type EncryptedReLU struct {
    eval       ckks.Evaluator
    polynomial []complex128
}

func (r *EncryptedReLU) Forward(input *ckks.Ciphertext) (*ckks.Ciphertext, error) {
    // Approximate ReLU with polynomial
    // ReLU(x) ≈ 0.5x + 0.25x^2 + 0.125x^3 for x ∈ [-2, 2]

    x2 := r.eval.MulRelinNew(input, input)
    r.eval.Rescale(x2, x2)

    x3 := r.eval.MulRelinNew(x2, input)
    r.eval.Rescale(x3, x3)

    // Combine terms
    result := r.eval.MulNew(input, 0.5)
    r.eval.Add(result, r.eval.MulNew(x2, 0.25), result)
    r.eval.Add(result, r.eval.MulNew(x3, 0.125), result)

    return result, nil
}

// Run inference
func (nn *EncryptedNeuralNetwork) Predict(input *ckks.Ciphertext) (*ckks.Ciphertext, error) {
    current := input

    for _, layer := range nn.layers {
        next, err := layer.Forward(current)
        if err != nil {
            return nil, err
        }
        current = next
    }

    return current, nil
}

func privateDatabaseQuery() {
    // Setup BGV for exact computation
    params, _ := bgv.NewParametersFromLiteral(
        bgv.ParametersLiteral{
            LogN:             13,
            LogQ:             []int{60, 60, 60},
            PlaintextModulus: 65537,
        },
    )

    // Database with 1024 entries
    database := make([]uint64, 1024)
    for i := range database {
        database[i] = uint64(i * i) // Example data
    }

    // Client wants to query index 42 privately
    queryIndex := 42

    // Create one-hot encoded query vector
    query := make([]uint64, len(database))
    query[queryIndex] = 1

    // Encrypt query
    encoder := bgv.NewEncoder(params)
    encryptor := bgv.NewEncryptor(params, pk)

    ptQuery := bgv.NewPlaintext(params)
    encoder.Encode(query, ptQuery)
    ctQuery, _ := encryptor.EncryptNew(ptQuery)

    // Server computes response
    // result = sum(database[i] * query[i]) for all i
    evaluator := bgv.NewEvaluator(params, evk)

    ctResult := bgv.NewCiphertext(params, 1, params.MaxLevel())
    for i, value := range database {
        if value != 0 {
            // Rotate query to align with database entry
            rotated := evaluator.RotateNew(ctQuery, i)

            // Multiply by database value
            evaluator.MulScalar(rotated, value, rotated)

            // Add to result
            evaluator.Add(ctResult, rotated, ctResult)
        }
    }

    // Client decrypts result
    ptResult := decryptor.DecryptNew(ctResult)
    result := encoder.DecodeUint(ptResult)

    fmt.Printf("Database[%d] = %d\n", queryIndex, result[0])
}

func thresholdDecryption() {
    params, _ := ckks.NewParametersFromLiteral(ckks.PN14QP438)

    // Setup: 3 parties, threshold = 2
    numParties := 3
    threshold := 2

    // Each party generates their secret key share
    parties := make([]*Party, numParties)
    for i := range parties {
        parties[i] = &Party{
            ID: i,
            SK: rlwe.NewKeyGenerator(params).GenSecretKeyNew(),
        }
    }

    // Threshold key generation protocol
    tkg := multiparty.NewTKGProtocol(params)

    // Each party creates their share
    shares := make([]*multiparty.TKGShare, numParties)
    for i, party := range parties {
        shares[i] = tkg.GenShare(party.SK)
    }

    // Aggregate to get collective public key
    collectivePK := tkg.AggregateShares(shares)

    // Encrypt data with collective public key
    encoder := ckks.NewEncoder(params)
    encryptor := rlwe.NewEncryptor(params, collectivePK)

    data := []complex128{1.23, 4.56, 7.89}
    pt := ckks.NewPlaintext(params, params.MaxLevel())
    encoder.Encode(data, pt)
    ct, _ := encryptor.EncryptNew(pt)

    // Threshold decryption with any 2 parties
    pcks := multiparty.NewPCKSProtocol(params)

    // Select subset of parties (meets threshold)
    activeParties := parties[:threshold]

    // Generate decryption shares
    decShares := make([]*multiparty.PCKSShare, threshold)
    for i, party := range activeParties {
        decShares[i] = pcks.GenShare(party.SK, ct)
    }

    // Combine shares to decrypt
    ptResult := pcks.AggregateShares(decShares, ct)
    result := encoder.Decode(ptResult, params.LogSlots())

    fmt.Println("Decrypted:", result[:3])
}

type SecureMLTraining struct {
    params   ckks.Parameters
    parties  []*MLParty
    protocol *multiparty.Protocol
}

type MLParty struct {
    ID       int
    SK       *rlwe.SecretKey
    DataSet  [][]float64
    LocalModel *ckks.Ciphertext
}

func (sml *SecureMLTraining) FederatedAveraging() (*ckks.Ciphertext, error) {
    // Each party trains locally and encrypts their model
    encryptedModels := make([]*ckks.Ciphertext, len(sml.parties))

    for i, party := range sml.parties {
        // Local training (in plaintext)
        localWeights := party.TrainLocal()

        // Encrypt local model
        pt := ckks.NewPlaintext(sml.params, sml.params.MaxLevel())
        encoder.Encode(localWeights, pt)
        encryptedModels[i], _ = encryptor.EncryptNew(pt)
    }

    // Secure aggregation
    globalModel := sml.SecureAggregate(encryptedModels)

    // Collaborative decryption
    return sml.CollectiveDecrypt(globalModel)
}

func (sml *SecureMLTraining) SecureAggregate(models []*ckks.Ciphertext) *ckks.Ciphertext {
    // Average encrypted models
    result := models[0].CopyNew()

    for i := 1; i < len(models); i++ {
        evaluator.Add(result, models[i], result)
    }

    // Divide by number of parties
    scale := 1.0 / float64(len(models))
    evaluator.MulScalar(result, scale, result)
    evaluator.Rescale(result, result)

    return result
}

type PrivatePortfolio struct {
    holdings  map[string]*ckks.Ciphertext // Encrypted quantities
    evaluator ckks.Evaluator
}

func (p *PrivatePortfolio) CalculateValue(prices map[string]float64) *ckks.Ciphertext {
    var totalValue *ckks.Ciphertext

    for asset, encryptedQuantity := range p.holdings {
        price := prices[asset]

        // value = quantity * price
        value := p.evaluator.MulScalarNew(encryptedQuantity, price)
        p.evaluator.Rescale(value, value)

        if totalValue == nil {
            totalValue = value
        } else {
            p.evaluator.Add(totalValue, value, totalValue)
        }
    }

    return totalValue
}

func (p *PrivatePortfolio) CalculateRisk(
    correlationMatrix [][]float64,
    volatilities []float64,
) *ckks.Ciphertext {
    // Portfolio risk = sqrt(w^T * Σ * w)
    // where Σ is covariance matrix

    assets := make([]string, 0, len(p.holdings))
    weights := make([]*ckks.Ciphertext, 0, len(p.holdings))

    for asset, holding := range p.holdings {
        assets = append(assets, asset)
        weights = append(weights, holding)
    }

    // Compute weighted covariance
    risk := ckks.NewCiphertext(p.evaluator.Parameters(), 1, weights[0].Level())

    for i := range assets {
        for j := range assets {
            // cov[i,j] = correlation[i,j] * vol[i] * vol[j]
            cov := correlationMatrix[i][j] * volatilities[i] * volatilities[j]

            // risk += w[i] * w[j] * cov[i,j]
            term := p.evaluator.MulNew(weights[i], weights[j])
            p.evaluator.Rescale(term, term)
            p.evaluator.MulScalar(term, cov, term)
            p.evaluator.Add(risk, term, risk)
        }
    }

    // Note: Square root would require polynomial approximation
    return risk
}

func bootstrappingExample() {
    // Parameters with bootstrapping support
    params, _ := ckks.NewParametersFromLiteral(
        ckks.ParametersLiteral{
            LogN: 16,
            LogQ: []int{
                // Computation levels
                60, 50, 50, 50, 50,
                // Bootstrapping levels
                60, 60, 60, 60, 60, 60, 60, 60,
            },
            LogP:            []int{61, 61, 61},
            LogDefaultScale: 50,
            H:               192,
        },
    )

    // Setup bootstrapper
    btpParams := bootstrapping.ParametersLiteral{
        LogN:     params.LogN(),
        LogSlots: params.LogSlots(),
        H:        192,
        Sigma:    params.Sigma(),
    }

    bootstrapper, _ := bootstrapping.NewBootstrapper(params, btpParams, evk)

    // Perform deep computation
    ct := encryptInitialData()

    for i := 0; i < 100; i++ {
        // Complex computation
        ct = evaluator.MulRelinNew(ct, ct)
        evaluator.Rescale(ct, ct)

        // Check if bootstrapping needed
        if ct.Level() == 0 {
            fmt.Printf("Bootstrapping at iteration %d\n", i)
            ct, _ = bootstrapper.Bootstrap(ct)
        }
    }

    // Decrypt final result
    result := decrypt(ct)
    fmt.Println("Result after 100 multiplications:", result)
}

func evaluateCustomFunction() {
    // Approximate sin(x) using Chebyshev polynomials
    // sin(x) ≈ x - x³/6 + x⁵/120 - x⁷/5040 for x ∈ [-π, π]

    coeffs := []float64{
        0,           // x⁰
        1,           // x¹
        0,           // x²
        -1.0/6,      // x³
        0,           // x⁴
        1.0/120,     // x⁵
        0,           // x⁶
        -1.0/5040,   // x⁷
    }

    // Create polynomial evaluator
    polyEval := ckks.NewPolynomialEvaluator(params, evaluator)

    // Input values
    x := []complex128{0, 0.5, 1.0, 1.5}
    ptX := encode(x)
    ctX := encrypt(ptX)

    // Evaluate polynomial
    ctResult, _ := polyEval.Evaluate(ctX, coeffs)

    // Compare with actual sin(x)
    result := decrypt(ctResult)
    fmt.Println("Approximate sin(x):", result)
    fmt.Println("Actual sin(x):", []float64{
        math.Sin(0), math.Sin(0.5), math.Sin(1.0), math.Sin(1.5),
    })
}

func efficientBatching() {
    // Process multiple data items in single ciphertext
    batchSize := params.Slots() // Use all available slots

    // Pack data efficiently
    data := make([]complex128, batchSize)
    for i := 0; i < actualDataSize; i++ {
        data[i] = actualData[i]
    }
    // Remaining slots stay zero

    // Single encryption for entire batch
    pt := encode(data)
    ct := encrypt(pt)

    // Operations apply to all slots simultaneously
    ctSquared := evaluator.MulRelinNew(ct, ct)
    evaluator.Rescale(ctSquared, ctSquared)

    // Decrypt entire batch
    result := decrypt(ctSquared)
    // Use first actualDataSize elements
}

func memoryEfficientComputation() {
    // Reuse ciphertexts to reduce memory
    ct1 := encrypt(data1)
    ct2 := encrypt(data2)

    // In-place operations
    evaluator.Add(ct1, ct2, ct1) // Result overwrites ct1
    evaluator.MulRelin(ct1, ct1, ct1) // Square in-place
    evaluator.Rescale(ct1, ct1)

    // Manual memory cleanup for large computations
    runtime.GC()

    // Use object pools for repeated allocations
    ctPool := sync.Pool{
        New: func() interface{} {
            return ckks.NewCiphertext(params, 1, params.MaxLevel())
        },
    }

    // Get from pool
    temp := ctPool.Get().(*ckks.Ciphertext)
    // Use temp...
    // Return to pool
    ctPool.Put(temp)
}

func TestHomomorphicAddition(t *testing.T) {
    // Setup
    params, _ := ckks.NewParametersFromLiteral(testParams)
    // ... initialize crypto components

    // Test data
    a := []complex128{1.0, 2.0, 3.0, 4.0}
    b := []complex128{5.0, 6.0, 7.0, 8.0}
    expected := []complex128{6.0, 8.0, 10.0, 12.0}

    // Encrypt
    ctA := encryptVector(a)
    ctB := encryptVector(b)

    // Homomorphic addition
    ctSum := evaluator.AddNew(ctA, ctB)

    // Decrypt and verify
    result := decryptVector(ctSum)

    for i := range expected {
        if math.Abs(real(result[i] - expected[i])) > 0.01 {
            t.Errorf("Mismatch at %d: got %v, want %v",
                i, result[i], expected[i])
        }
    }
}

func BenchmarkHomomorphicMultiplication(b *testing.B) {
    // Setup
    params, _ := ckks.NewParametersFromLiteral(ckks.PN14QP438)
    // ... initialize

    ct1 := encrypt(randomData())
    ct2 := encrypt(randomData())

    b.ResetTimer()
    for i := 0; i < b.N; i++ {
        result := evaluator.MulRelinNew(ct1, ct2)
        evaluator.Rescale(result, result)
    }
}