test_squid_game.py

# test_squid_game.py

import pytest
from math import isclose
from functools import lru_cache

# Import from our main module
from squid_game_core import (
    parse_tier_map,
    tierValue,
    is_terminal,
    compute_final_payout,
    get_expected_value,
    next_squid_gain_for_nonzero,
    hypothetical_next_round_gain
)

def nearly_equal(a, b, tol=1e-9):
    return isclose(a, b, abs_tol=tol)

@pytest.fixture
def tier_map_example():
    """
    We'll define a bracket:
      1-4:1 => if you have 1..4 squids => total = k * 1
      5-6:3 => if you have 5..6 squids => total = k * 3
    0 => always 0
    """
    # We'll store it as a tuple for use in DP
    # parse_tier_map would do the same, but let's define it directly
    return (
        (0,0,0.0),
        (1,4,1.0),
        (5,6,3.0)
    )

def test_multiple_losers_pay_individually(tier_map_example):
    """
    Scenario: 3 players => final distribution=(0,0,4).
    According to bracket (1-4:1 => 4=>4*1=4).
    - 2 losers => each pays 4 => each has payoff=-4
    - 1 winner => receives 2 * 4=8.
    """
    dist = (0,0,4)
    payoffs = compute_final_payout(dist, tier_map_example)
    # payoffs => [ -4, -4, 8 ]
    assert nearly_equal(payoffs[0], -4)
    assert nearly_equal(payoffs[1], -4)
    assert nearly_equal(payoffs[2],  8)

def test_single_loser(tier_map_example):
    """
    Scenario: 2 players => final distribution=(1,0).
    => 1 zero => that zero pays sum_of_winners= tierValue(1)=1
    => payoff=(1, -1) because:
    - Winner keeps their tierValue (1)
    - Loser pays that amount (-1)
    """
    dist = (1,0)
    payoffs = compute_final_payout(dist, tier_map_example)
    assert nearly_equal(payoffs[0], 1)  # Winner gets tierValue(1)=1
    assert nearly_equal(payoffs[1], -1)  # Loser pays -1

def test_no_losers(tier_map_example):
    """
    Scenario: 2 players => final distribution=(2,3).
    => no zero => no payment => payoff=(0,0).
    """
    dist = (2,3)
    payoffs = compute_final_payout(dist, tier_map_example)
    # 2 => bracket(1..4 => *1)=>2
    # 3 => bracket(1..4 => *1)=>3
    # but since no zero => payoff=(0,0).
    assert nearly_equal(payoffs[0], 0)
    assert nearly_equal(payoffs[1], 0)

def test_next_squid_gain_for_nonzero(tier_map_example):
    """
    distribution=(4,0,2) => 
      - Player0=4 => tierValue(4)=4 => tierValue(5)=15 => gain=11
        (since 5 squids => bracket => 5*3=15)
      - Player1=0 => skip
      - Player2=2 => tierValue(2)=2 => tierValue(3)=3 => gain=1
    """
    dist = (4,0,2)
    gains = next_squid_gain_for_nonzero(dist, tier_map_example)
    assert 0 in gains
    assert gains[0] == 11  # 15-4
    assert 2 in gains
    assert gains[2] == 1   # 3-2
    assert 1 not in gains  # because that player has 0

def test_ev_with_leftover_multiple_losers(tier_map_example):
    """
    We'll test a DP scenario:
      N=3, X=4 total squids
      Current distribution=(0,1,1)
      => sum=2 => leftover=2 => not terminal.

    Let's see possible final states:
      - They could keep awarding 2 more squids in 2 rounds.
      - It's possible we end with multiple zeros if the 2 additional squids both go to the same non-zero player, 
        or exactly one zero, etc.

    We'll just check the computed EV matches a hand-run or at least we confirm no errors.
    """
    dist = (0,1,1)
    X = 4
    r = X - sum(dist)  # 4-2=2

    # We'll do a quick partial analysis by enumerating the 2 leftover squids:
    # Round 1 => three possibilities: P0, P1, P2
    # But let's just rely on the solver to give us a final result, 
    # and we'll assert that we get a numeric 3-tuple and 
    # the sum of payoffs is near 0 (since it's a zero-sum game).
    
    from squid_game import get_expected_value
    get_expected_value.cache_clear()
    ev = get_expected_value(dist, r, tier_map_example)
    assert len(ev) == 3
    # Because it's zero-sum, the sum of EVs should be very close to 0:
    total_ev = sum(ev)
    assert nearly_equal(total_ev, 0.0), f"Sum of EVs is not near 0, got {total_ev}"

    # We won't do a full hand enumeration here, 
    # but we at least confirm the DP runs and yields a plausible sum=0 result.

def test_ev_multiple_losers_specific(tier_map_example):
    """
    A more direct test for multiple losers via DP:
      N=3, X=2, distribution=(0,0,2)
      => sum=2 => leftover=0 => terminal => multiple zero => each zero pays tierValue(2)=2 => payoff=(-2,-2,4)

    Then we check that the DP logic returns the same final payoff if is_terminal is triggered.
    """
    dist = (0,0,2)
    X = 2
    r = X - sum(dist)  # leftover=0 => terminal immediately
    from squid_game import get_expected_value
    get_expected_value.cache_clear()
    ev = get_expected_value(dist, r, tier_map_example)
    # With bracket => 2 => 2*1=2 => each zero pays 2 => 2 losers => winner gets 2*2=4
    # payoff=( -2, -2, 4 )
    assert nearly_equal(ev[0], -2)
    assert nearly_equal(ev[1], -2)
    assert nearly_equal(ev[2],  4)
    # sum = 0
    assert nearly_equal(sum(ev), 0.0)

def test_hypothetical_next_round_gain_with_penalty():
    """
    Test scenario: 3 players with distribution=(0,0,2)
    Expected values would be (-2,-2,4) as shown in test_ev_multiple_losers_specific
    So penalty should be 2 (abs of -2)
    
    For zero-squid players (0,1):
    - Getting 1 squid = tierValue(1) = 1
    - Avoiding penalty share = 2/2 = 1 (penalty/zero_count)
    - Total gain = 2
    
    For player with 2 squids:
    - Going from 2 to 3 = tierValue(3) - tierValue(2) = 3 - 2 = 1
    """
    tier_map = (
        (0,0,0.0),
        (1,4,1.0),
        (5,6,3.0)
    )
    dist = (0,0,2)
    penalty = 2  # abs of -2 from expected values
    
    gains = hypothetical_next_round_gain(dist, tier_map, penalty)
    
    # Check zero-squid players
    assert nearly_equal(gains[0], 2.0), f"Expected gain 2.0 for player 1, got {gains[0]}"
    assert nearly_equal(gains[1], 2.0), f"Expected gain 2.0 for player 2, got {gains[1]}"
    
    # Check non-zero player
    assert nearly_equal(gains[2], 1.0), f"Expected gain 1.0 for player 3, got {gains[2]}"

def test_hypothetical_next_round_gain_no_zeros():
    """
    Test scenario: 2 players with distribution=(1,2)
    No zero-squid players, so penalty doesn't matter
    
    Player 1: going from 1 to 2 = tierValue(2) - tierValue(1) = 2 - 1 = 1
    Player 2: going from 2 to 3 = tierValue(3) - tierValue(2) = 3 - 2 = 1
    """
    tier_map = (
        (0,0,0.0),
        (1,4,1.0),
        (5,6,3.0)
    )
    dist = (1,2)
    penalty = 0  # doesn't matter since no zero-squid players
    
    gains = hypothetical_next_round_gain(dist, tier_map, penalty)
    
    assert nearly_equal(gains[0], 1.0), f"Expected gain 1.0 for player 1, got {gains[0]}"
    assert nearly_equal(gains[1], 1.0), f"Expected gain 1.0 for player 2, got {gains[1]}"