# Copyright 2025 Ant Group Co., Ltd.
#
# 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.

"""Example: Secure Permutation and Homomorphic Aggregation.

This example demonstrates two advanced MPC/HE capabilities:
1. **Secure Permutation Network**: Using OT to shuffle data without revealing the permutation.
2. **Homomorphic Aggregation**: Using BFV Rotation to efficiently aggregate packed data.

"""

import jax.numpy as jnp
import numpy as np

import mplang.v2.backends.bfv_impl  # noqa: F401
import mplang.v2.backends.tensor_impl  # noqa: F401
import mplang.v2.dialects.bfv as bfv
import mplang.v2.dialects.simp as simp
import mplang.v2.dialects.tensor as tensor
from mplang.v2.backends.simp_simulator import SimpSimulator
from mplang.v2.libs.mpc import apply_permutation, rotate_and_sum


def main():
    # Initialize Simulator (2 Parties)
    # Party 0: Data Owner
    # Party 1: Compute Node
    interp = SimpSimulator(world_size=2)

    print("=== Part 1: Secure Permutation Network ===")
    run_secure_permutation(interp)

    print("\n=== Part 2: Homomorphic Aggregation (BFV) ===")
    run_homomorphic_aggregation(interp)


def run_secure_permutation(interp):
    """Demonstrate shuffling data using OT."""

    # Data on Party 0
    # Permutation on Party 1: [2, 0, 1, 3] -> Expect [30, 10, 20, 40]
    # Note: Our simple permutation lib currently only supports N=2 or manual construction.
    # For this demo, we will use N=2 to show it works, or manually chain switches for N=4?
    # The `apply_permutation` in lib currently raises NotImplemented for N=4.
    # Let's stick to N=2 for the automated call, or implement the N=4 logic manually here.

    # Let's use N=2 for simplicity of the library call.
    data_n2 = [10, 20]
    perm_n2 = [1, 0]  # Swap

    print(f"Sender (P0) Data: {data_n2}")
    print(f"Receiver (P1) Permutation: {perm_n2}")

    def protocol(d0, d1, p0, p1):
        # Construct permutation tensor on P1
        def make_perm(a, b):
            return tensor.run_jax(lambda x, y: jnp.array([x, y]), a, b)

        perm = simp.pcall_static((1,), make_perm, p0, p1)

        # Run secure permutation
        return apply_permutation([d0, d1], perm, sender=0, receiver=1)

    with interp:
        d0_obj = simp.constant((0,), data_n2[0])
        d1_obj = simp.constant((0,), data_n2[1])
        p0_obj = simp.constant((1,), perm_n2[0])
        p1_obj = simp.constant((1,), perm_n2[1])

        res = protocol(d0_obj, d1_obj, p0_obj, p1_obj)

    # Result is on P1
    res_vals = [r.runtime_obj.values[1] for r in res]
    print(f"Receiver (P1) Result: {res_vals}")

    expected = [20, 10]
    assert res_vals == expected
    print("Secure Permutation Verified!")


def run_homomorphic_aggregation(interp):
    """Demonstrate O(log N) aggregation using BFV."""

    # Setup BFV context (Keys generated by P0)
    # In this simulation, we treat BFV ops as local or "available to all who have keys".
    # We'll simulate P0 generating keys and P1 computing.

    # 1. Keygen (P0)
    pk, sk = bfv.keygen(poly_modulus_degree=4096)
    gk = bfv.make_galois_keys(sk)
    encoder = bfv.create_encoder(poly_modulus_degree=4096)

    # 2. Data (P0)
    # Vector of size 8: [1, 1, 1, 1, 1, 1, 1, 1] -> Sum = 8
    data = np.array([1] * 8 + [0] * 4088, dtype=np.int64)

    t_data = tensor.constant(data)
    pt = bfv.encode(t_data, encoder)
    ct = bfv.encrypt(pt, pk)

    print(f"Encrypted Data (First 8 slots): {data[:8]}")

    # 3. Aggregation (P1)
    # P1 has ct and gk. P1 computes sum of first 8 elements.

    def compute(ct_in, gk_in):
        return rotate_and_sum(ct_in, k=8, galois_keys=gk_in)

    res_ct = compute(ct, gk)

    # 4. Decrypt (P0)
    res_pt = bfv.decrypt(res_ct, sk)
    res_tensor = bfv.decode(res_pt, encoder)

    # In simulation, we can access the value directly
    result_val = res_tensor.runtime_obj[0]
    print(f"Aggregated Result (Slot 0): {result_val}")

    assert result_val == 8
    print("Homomorphic Aggregation Verified!")


if __name__ == "__main__":
    main()
