import unittest
import numpy as np
from static_mechanism_payment_rule import calculate_quantities_for_agent  # Assume your code is in 'your_module.py'

class TestQuantitiesForAgent(unittest.TestCase):
    def setUp(self):
        self.tau = 1.0  # Example temperature parameter
        self.num_agents = 3  # Example number of agents
        self.num_outcomes = 4  # Example number of outcomes

    def test_probabilities_consistency_across_agents(self):
        """Test that the probabilities vector is consistent across all agents for the same rewards matrix."""
        # Generate a random rewards matrix
        all_rewards = np.random.rand(self.num_agents, self.num_outcomes)

        # Calculate the probabilities for the first agent to establish a baseline
        baseline_probabilities, _, _, _ = calculate_quantities_for_agent(all_rewards, 0, None, self.tau)

        # Calculate and compare probabilities for each other agent
        for agent_index in range(1, self.num_agents):
            probabilities, _, _, _ = calculate_quantities_for_agent(all_rewards, agent_index, None, self.tau)
            np.testing.assert_array_almost_equal(baseline_probabilities, probabilities, decimal=5,
                                                 err_msg=f"Probabilities vector differs for agent {agent_index}")

    def test_non_negative_utility(self):
        """Test that the utility is always non-negative."""
        all_rewards = np.random.rand(self.num_agents, self.num_outcomes)
        for agent_index in range(self.num_agents):
            _, u, _, _ = calculate_quantities_for_agent(all_rewards, agent_index, None, self.tau)
            self.assertGreaterEqual(u, 0, "Utility is negative.")

    def test_non_negative_transfers(self):
        """Test that the transfers are always non-negative."""
        all_rewards = np.random.rand(self.num_agents, self.num_outcomes)
        for agent_index in range(self.num_agents):
            _, _, _, t = calculate_quantities_for_agent(all_rewards, agent_index, None, self.tau)
            self.assertGreaterEqual(t, 0, "Transfer is negative.")

if __name__ == "__main__":
    unittest.main()