# 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.

"""Tests for OT library."""

import numpy as np
import pytest

import mplang.v2.backends.crypto_impl  # noqa: F401
import mplang.v2.backends.tensor_impl  # noqa: F401 (registers tensor primitives)
from mplang.v2.backends.simp_simulator import SimpSimulator
from mplang.v2.backends.tensor_impl import TensorValue
from mplang.v2.dialects import simp
from mplang.v2.libs.mpc import ot


def _unwrap(val):
    """Unwrap TensorValue to numpy array."""
    if isinstance(val, TensorValue):
        return val.unwrap()
    return val


class TestOTScalar:
    """Test suite for scalar 1-out-of-2 Oblivious Transfer."""

    def setup_method(self):
        """Initialize simulator for each test."""
        self.interp = SimpSimulator(world_size=2)

    @pytest.mark.parametrize(
        "m0_val, m1_val, choice_val, expected",
        [
            # Basic cases
            (10, 20, 0, 10),
            (10, 20, 1, 20),
            # Zero values
            (0, 100, 0, 0),
            (0, 100, 1, 100),
            # Negative values
            (-42, -100, 0, -42),
            (-42, -100, 1, -100),
            # Mixed signs
            (-42, 42, 0, -42),
            (-42, 42, 1, 42),
            # Large values
            (2**30, 2**31 - 1, 1, 2**31 - 1),
            # Max/Min int64
            (2**63 - 1, 2**63 - 2, 0, 2**63 - 1),
            (-(2**63), -(2**63) + 1, 0, -(2**63)),
            # Equal messages
            (42, 42, 0, 42),
            (42, 42, 1, 42),
            # Powers of two
            (2**10, 2**20, 1, 2**20),
            # Binary values
            (1, 0, 0, 1),
            (1, 0, 1, 0),
        ],
    )
    def test_scalar_ot(self, m0_val, m1_val, choice_val, expected):
        """Test scalar OT with various inputs."""
        with self.interp:
            m0 = simp.constant((0,), m0_val)
            m1 = simp.constant((0,), m1_val)
            choice = simp.constant((1,), choice_val)

            res = ot.transfer(m0, m1, choice, sender=0, receiver=1)

        values = self.interp.fetch(res.runtime_obj)
        result = _unwrap(values[1])
        assert result.item() == expected


class TestOTVector:
    """Test suite for vectorized OT operations."""

    def setup_method(self):
        """Initialize simulator for each test."""
        self.interp = SimpSimulator(world_size=2)

    def test_vector_all_zeros(self):
        """Test vectorized OT with all choices=0."""
        N = 10
        m0_data = np.arange(N, dtype=np.int64)
        m1_data = np.arange(N, N * 2, dtype=np.int64)
        choices = np.zeros(N, dtype=np.int64)

        with self.interp:
            m0 = simp.constant((0,), m0_data)
            m1 = simp.constant((0,), m1_data)
            choice = simp.constant((1,), choices)

            res = ot.transfer(m0, m1, choice, sender=0, receiver=1)

        values = self.interp.fetch(res.runtime_obj)
        np.testing.assert_array_equal(_unwrap(values[1]), m0_data)

    def test_vector_all_ones(self):
        """Test vectorized OT with all choices=1."""
        N = 10
        m0_data = np.arange(N, dtype=np.int64)
        m1_data = np.arange(N, N * 2, dtype=np.int64)
        choices = np.ones(N, dtype=np.int64)

        with self.interp:
            m0 = simp.constant((0,), m0_data)
            m1 = simp.constant((0,), m1_data)
            choice = simp.constant((1,), choices)

            res = ot.transfer(m0, m1, choice, sender=0, receiver=1)

        values = self.interp.fetch(res.runtime_obj)
        np.testing.assert_array_equal(_unwrap(values[1]), m1_data)

    def test_vector_mixed_choices(self):
        """Test vectorized OT with mixed choices (0 and 1)."""
        N = 10
        m0_data = np.full(N, 10, dtype=np.int64)
        m1_data = np.full(N, 20, dtype=np.int64)
        # Alternating 0 and 1: [0, 1, 0, 1, ...]
        choices = np.array([i % 2 for i in range(N)], dtype=np.int64)
        expected = np.where(choices == 0, m0_data, m1_data)

        with self.interp:
            m0 = simp.constant((0,), m0_data)
            m1 = simp.constant((0,), m1_data)
            choice = simp.constant((1,), choices)

            res = ot.transfer(m0, m1, choice, sender=0, receiver=1)

        values = self.interp.fetch(res.runtime_obj)
        np.testing.assert_array_equal(_unwrap(values[1]), expected)
