"""Unit tests for :mod:`tmlt.core.measures`."""

# SPDX-License-Identifier: Apache-2.0
# Copyright Tumult Labs 2025
import itertools
import re
from typing import Any, Tuple, Union
from unittest.case import TestCase

import sympy as sp
from parameterized import parameterized
from typeguard import TypeCheckError

from tmlt.core.measures import (
    ApproxDP,
    ApproxDPBudget,
    InsufficientBudgetError,
    PureDP,
    PureDPBudget,
    RhoZCDP,
    RhoZCDPBudget,
)
from tmlt.core.utils.exact_number import ExactNumber, ExactNumberInput

VALID_PRIMARY_BUDGET_INPUTS = [
    0,
    10,
    float("inf"),
    "3",
    "32",
    sp.Integer(0),
    sp.Integer(1),
    sp.Rational("42.17"),
    sp.oo,
]
INVALID_PRIMARY_BUDGET_INPUTS = [-1, 2.0, sp.Float(2), "wat", {}]
VALID_DELTA_INPUTS = [0, sp.Integer(1), "1", sp.Rational("0.5")]
INVALID_DELTA_INPUTS = [-1, 0.5, 2, sp.Float(1), "wat", {}]
VALID_APPROX_DP_INPUTS = list(
    itertools.product(VALID_PRIMARY_BUDGET_INPUTS, VALID_DELTA_INPUTS)
)
INVALID_APPROX_DP_INPUTS = (
    list(itertools.product(VALID_PRIMARY_BUDGET_INPUTS, INVALID_DELTA_INPUTS))
    + list(itertools.product(INVALID_PRIMARY_BUDGET_INPUTS, VALID_DELTA_INPUTS))
    + VALID_PRIMARY_BUDGET_INPUTS
    + [(1, 1, 1), [1, 1]]
)

to_singletons = lambda x: map(lambda y: (y,), x)


class TestPureDP(TestCase):
    """TestCase for PureDP."""

    def setUp(self):
        """Setup."""
        self.pureDP = PureDP()

    @parameterized.expand(to_singletons(VALID_PRIMARY_BUDGET_INPUTS))
    def test_valid(self, value: ExactNumberInput):
        """Tests for valid values of epsilon."""
        self.pureDP.validate(value)

    @parameterized.expand(
        [
            (sp.Integer(0), sp.Integer(1), True),
            (sp.Rational("42.17"), sp.Rational("42.17"), True),
            (sp.Integer(0), sp.oo, True),
            (sp.oo, sp.oo, True),
            (sp.Integer(1), sp.Integer(0), False),
            (sp.Integer(1), sp.Rational("0.5"), False),
            (sp.oo, sp.Integer(1000), False),
        ]
    )
    def test_compare(
        self, value1: ExactNumberInput, value2: ExactNumberInput, expected: bool
    ):
        """Tests that compare returns the expected result."""
        self.assertEqual(self.pureDP.compare(value1, value2), expected)

    @parameterized.expand(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS))
    def test_invalid(self, val: Any):
        """Only valid ExactNumberInput's should be allowed."""
        with self.assertRaises((TypeCheckError, ValueError)):
            self.pureDP.validate(val)


class TestApproxDP(TestCase):
    """TestCase for ApproxDP."""

    def setUp(self):
        """Setup."""
        self.approxDP = ApproxDP()

    @parameterized.expand(to_singletons(VALID_APPROX_DP_INPUTS))
    def test_valid(self, value: Any):
        """Tests for valid values of epsilon and delta."""
        self.approxDP.validate(value)

    @parameterized.expand(
        [
            (
                (epsilon1, delta1),
                (epsilon2, delta2),
                (
                    (ExactNumber(epsilon1) == sp.oo or ExactNumber(delta1) == 1)
                    and (ExactNumber(epsilon2) == sp.oo or ExactNumber(delta2) == 1)
                )
                or (
                    ExactNumber(epsilon1) <= ExactNumber(epsilon2)
                    and ExactNumber(delta1) <= ExactNumber(delta2)
                ),
            )
            for epsilon1, epsilon2 in itertools.combinations(
                ["0", 1, sp.Rational("1.3"), sp.oo], 2
            )
            for delta1, delta2 in itertools.combinations([0, "0.5", sp.Integer(1)], 2)
        ]
    )
    def test_compare(
        self,
        value1: Tuple[ExactNumberInput, ExactNumberInput],
        value2: Tuple[ExactNumberInput, ExactNumberInput],
        expected: bool,
    ):
        """Tests that compare returns the expected result."""
        self.assertEqual(self.approxDP.compare(value1, value2), expected)

    @parameterized.expand(to_singletons(INVALID_APPROX_DP_INPUTS))
    def test_invalid(self, value: Any):
        """Only valid budgets should be allowed."""
        with self.assertRaises((TypeCheckError, ValueError)):
            self.approxDP.validate(value)


class TestRhoZCDP(TestCase):
    """Test cases for RhoZCDP."""

    def setUp(self):
        """Setup."""
        self.rhoZCDP = RhoZCDP()

    @parameterized.expand(to_singletons(VALID_PRIMARY_BUDGET_INPUTS))
    def test_valid(self, expr: ExactNumberInput):
        """Tests for valid values of rho."""
        self.rhoZCDP.validate(expr)

    @parameterized.expand(
        [
            (sp.Integer(0), sp.Integer(1), True),
            (sp.Rational("42.17"), sp.Rational("42.17"), True),
            (sp.Integer(0), sp.oo, True),
            (sp.oo, sp.oo, True),
            (sp.Integer(1), sp.Integer(0), False),
            (sp.Integer(1), sp.Rational("0.5"), False),
            (sp.oo, sp.Integer(1000), False),
        ]
    )
    def test_compare(
        self, value1: ExactNumberInput, value2: ExactNumberInput, expected: bool
    ):
        """Tests that compare returns the expected result."""
        self.assertEqual(self.rhoZCDP.compare(value1, value2), expected)

    @parameterized.expand(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS))
    def test_invalid(self, val: Any):
        """Only valid budgets should be allowed."""
        with self.assertRaises((TypeCheckError, ValueError)):
            self.rhoZCDP.validate(val)


class TestPureDPBudget(TestCase):
    """Test cases for PureDPBudget."""

    @parameterized.expand(to_singletons(VALID_PRIMARY_BUDGET_INPUTS))
    def test_init_valid(self, value):
        """Tests that the budget can be created with a valid budget input."""
        PureDPBudget(value)

    @parameterized.expand(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS))
    def test_init_invalid(self, value):
        """Throw an error when created with an invalid budget."""
        with self.assertRaises((TypeCheckError, ValueError)):
            PureDPBudget(value)

    @parameterized.expand(
        [
            (1, True),
            (sp.Rational("76.4"), True),
            (sp.Integer(5), True),
            (sp.oo, False),
            (float("inf"), False),
        ]
    )
    def test_is_finite(self, value, expected: bool):
        """Test that is_finite returns the expected value."""
        self.assertEqual(PureDPBudget(value).is_finite(), expected)

    @parameterized.expand(
        [
            (sp.Integer(0), sp.Integer(1), False),
            (sp.Rational("42.17"), sp.Rational("42.17"), True),
            (sp.Integer(0), sp.oo, False),
            (sp.oo, sp.oo, True),
            (sp.Integer(1), sp.Integer(0), True),
            (sp.Integer(1), sp.Rational("0.5"), True),
            (sp.oo, sp.Integer(1000), True),
        ]
    )
    def test_assert_can_spend_budget(self, value1, value2, expected: bool):
        """Test that assert_can_spend_budget returns the expected value."""
        if expected:
            PureDPBudget(value1).assert_can_spend_budget(value2)
        else:
            with self.assertRaises(ValueError):
                PureDPBudget(value1).assert_can_spend_budget(value2)

    @parameterized.expand(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS))
    def test_assert_can_spend_budget_invalid(self, value):
        """Test that assert_can_spend_budget raises an error for invalid values."""
        with self.assertRaises((TypeCheckError, ValueError)):
            PureDPBudget(1).assert_can_spend_budget(value)

    @parameterized.expand(
        [
            (5, "3", 2),
            (sp.Integer(5), sp.Rational("1.5"), sp.Rational("3.5")),
            (sp.oo, 10, sp.oo),
            (sp.oo, sp.oo, sp.oo),
            ("4", sp.Rational("0"), 4),
            ("4", sp.Integer(4), 0),
        ]
    )
    def test_subtract(self, value1, value2, expected):
        """Test that subtract returns the expected value."""
        self.assertEqual(PureDPBudget(value1).subtract(value2), PureDPBudget(expected))

    @parameterized.expand(
        list(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS)) + [(sp.Rational("1.1"),)]
    )
    def test_subtract_invalid(self, value):
        """Test that subtract raises an error when passed an invalid value."""
        with self.assertRaises((TypeCheckError, ValueError)):
            PureDPBudget(1).subtract(value)


class TestApproxDPBudget(TestCase):
    """Test cases for ApproxDPBudget."""

    # @parameterized.expand(to_singletons(VALID_APPROX_DP_INPUTS))
    # def test_init_valid(self, value):
    #     """Tests that the budget can be created with a valid budget input."""
    #     ApproxDPBudget(value)

    @parameterized.expand(to_singletons(INVALID_APPROX_DP_INPUTS))
    def test_init_invalid(self, value):
        """Throw an error when created with an invalid budget."""
        with self.assertRaises((TypeCheckError, ValueError)):
            ApproxDPBudget(value)

    @parameterized.expand(
        [
            ((1, 0), True),
            ((sp.Rational("76.4"), sp.Rational("0.5")), True),
            ((sp.Integer(5), sp.Rational("0.3")), True),
            ((sp.oo, 0), False),
            ((float("inf"), sp.Rational("0.5")), False),
            ((1, 1), False),
            ((float("inf"), 1), False),
        ]
    )
    def test_is_finite(self, value, expected: bool):
        """Test that is_finite returns the expected value."""
        self.assertEqual(ApproxDPBudget(value).is_finite(), expected)

    @parameterized.expand(
        [
            ((5, sp.Rational("0.1")), ("6", sp.Rational("0.2")), False),
            (("7", sp.Rational("0.1")), (sp.Integer(6), sp.Rational("0.2")), False),
            (
                (sp.Rational("4.5"), sp.Rational("0.2")),
                (sp.Integer(6), sp.Rational("0.1")),
                False,
            ),
            (("6", sp.Rational("0.1")), (sp.Integer(6), sp.Rational("0.1")), True),
            ((7, sp.Rational("0.1")), (sp.Integer(6), sp.Rational("0.1")), True),
            ((6, sp.Rational("0.2")), (sp.Integer(6), sp.Rational("0.1")), True),
            (
                (sp.Rational("8.0"), sp.Rational("0.3")),
                (sp.Integer(6), sp.Rational("0.1")),
                True,
            ),
            ((6, 1), (6, 1), True),
            ((sp.oo, sp.Rational("0.1")), (sp.oo, sp.Rational("0.1")), True),
            ((sp.oo, 1), (sp.oo, 1), True),
        ]
    )
    def test_assert_can_spend_budget(self, value1, value2, expected: bool):
        """Test that assert_can_spend_budget returns the expected value."""
        if expected:
            ApproxDPBudget(value1).assert_can_spend_budget(value2)
        else:
            with self.assertRaises(ValueError):
                ApproxDPBudget(value1).assert_can_spend_budget(value2)

    @parameterized.expand(to_singletons(INVALID_APPROX_DP_INPUTS))
    def test_assert_can_spend_budget_invalid(self, value):
        """Test that assert_can_spend_budget raises an error for invalid values."""
        with self.assertRaises((TypeCheckError, ValueError)):
            ApproxDPBudget((1, sp.Rational("0.1"))).assert_can_spend_budget(value)

    @parameterized.expand(
        [
            (
                (sp.Rational("5.5"), sp.Rational("0.5")),
                (sp.Integer(2), sp.Rational("0.3")),
                (sp.Rational("3.5"), sp.Rational("0.2")),
            ),
            (
                (sp.Rational("5.5"), sp.Rational("0.5")),
                (0, 0),
                (sp.Rational("5.5"), sp.Rational("0.5")),
            ),
            (("1", sp.Rational("0.1")), ("1", sp.Rational("0.1")), (0, 0)),
            (
                # when subtracting from an infinite budget,
                # the initial budget is returned
                (sp.Rational("1.5"), "1"),
                ("3", sp.Rational("0.9")),
                (sp.Rational("1.5"), "1"),
            ),
            (
                # when subtracting from an infinite budget,
                # the initial budget is returned
                (float("inf"), 0),
                ("1", sp.Rational("0.9")),
                (float("inf"), 0),
            ),
        ]
    )
    def test_subtract(self, value1, value2, expected):
        """Test that subtract returns the expected value."""
        self.assertEqual(
            ApproxDPBudget(value1).subtract(value2), ApproxDPBudget(expected)
        )

    @parameterized.expand(
        list(to_singletons(INVALID_APPROX_DP_INPUTS))
        + [((sp.Rational("1.1"), 0),), (("1", sp.Rational("0.2")),)]
    )
    def test_subtract_invalid(self, value):
        """Test that subtract raises an error when passed an invalid value."""
        with self.assertRaises((TypeCheckError, ValueError)):
            ApproxDPBudget((1, sp.Rational("0.1"))).subtract(value)


class TestRhoZCDPBudget(TestCase):
    """Test cases for RhoZCDPBudget."""

    @parameterized.expand(to_singletons(VALID_PRIMARY_BUDGET_INPUTS))
    def test_init_valid(self, value):
        """Tests that the budget can be created with a valid budget input."""
        RhoZCDPBudget(value)

    @parameterized.expand(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS))
    def test_init_invalid(self, value):
        """Throw an error when created with an invalid budget."""
        with self.assertRaises((TypeCheckError, ValueError)):
            RhoZCDPBudget(value)

    @parameterized.expand(
        [
            (1, True),
            (sp.Rational("76.4"), True),
            (sp.Integer(5), True),
            (sp.oo, False),
            (float("inf"), False),
        ]
    )
    def test_is_finite(self, value, expected: bool):
        """Test that is_finite returns the expected value."""
        self.assertEqual(RhoZCDPBudget(value).is_finite(), expected)

    @parameterized.expand(
        [
            (sp.Integer(0), sp.Integer(1), False),
            (sp.Rational("42.17"), sp.Rational("42.17"), True),
            (sp.Integer(0), sp.oo, False),
            (sp.oo, sp.oo, True),
            (sp.Integer(1), sp.Integer(0), True),
            (sp.Integer(1), sp.Rational("0.5"), True),
            (sp.oo, sp.Integer(1000), True),
        ]
    )
    def test_assert_can_spend_budget(self, value1, value2, expected: bool):
        """Test that assert_can_spend_budget returns the expected value."""
        if expected:
            RhoZCDPBudget(value1).assert_can_spend_budget(value2)
        else:
            with self.assertRaises(ValueError):
                RhoZCDPBudget(value1).assert_can_spend_budget(value2)

    @parameterized.expand(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS))
    def test_assert_can_spend_budget_invalid(self, value):
        """Test that assert_can_spend_budget raises an error for invalid values."""
        with self.assertRaises((TypeCheckError, ValueError)):
            RhoZCDPBudget(1).assert_can_spend_budget(value)

    @parameterized.expand(
        [
            (5, "3", 2),
            (sp.Integer(5), sp.Rational("1.5"), sp.Rational("3.5")),
            (sp.oo, 10, sp.oo),
            (sp.oo, sp.oo, sp.oo),
            ("4", sp.Rational("0"), 4),
            ("4", sp.Integer(4), 0),
        ]
    )
    def test_subtract(self, value1, value2, expected):
        """Test that subtract returns the expected value."""
        self.assertEqual(
            RhoZCDPBudget(value1).subtract(value2), RhoZCDPBudget(expected)
        )

    @parameterized.expand(
        list(to_singletons(INVALID_PRIMARY_BUDGET_INPUTS)) + [(sp.Rational("1.1"),)]
    )
    def test_subtract_invalid(self, value):
        """Test that subtract raises an error when passed an invalid value."""
        with self.assertRaises((TypeCheckError, ValueError)):
            RhoZCDPBudget(1).subtract(value)


class TestInsufficientBudgetError(TestCase):
    """Test cases for InsufficientBudgetError."""

    @parameterized.expand(
        [
            (
                ApproxDPBudget,
                ((ExactNumber(1), ExactNumber("0.1"))),
                (ExactNumber(1), ExactNumber("0.2")),
            ),
            (
                ApproxDPBudget,
                ((ExactNumber(1), ExactNumber("0.1"))),
                (ExactNumber(2), ExactNumber("0.1")),
            ),
            (PureDPBudget, ExactNumber(1), ExactNumber(2)),
            (RhoZCDPBudget, ExactNumber(1), ExactNumber(2)),
        ]
    )
    def test_InsufficientBudgetError(
        self,
        budget_type: Union[PureDPBudget, ApproxDPBudget, RhoZCDPBudget],
        low_budget: Union[ExactNumber, Tuple[ExactNumber, ExactNumber]],
        high_budget: Union[ExactNumber, Tuple[ExactNumber, ExactNumber]],
    ):
        """Tests that the error message related to InsufficientBudgetError is useful."""

        low = budget_type(low_budget)  # type: ignore
        high = budget_type(high_budget)  # type: ignore

        error_message = re.escape(
            (
                f"The remaining privacy budget is {low}, which "
                f"is insufficient given the requested budget {high}."
            )
        )

        with self.assertRaisesRegex(InsufficientBudgetError, error_message):
            raise InsufficientBudgetError(low, high)
