pf1e-simulator/tests/test_dice.py

"""Tests for dice notation parsing and DiceExpr evaluation."""

from __future__ import annotations

from collections import Counter

import pytest

from pf1e_simulator.dice import DiceExpr, DiceParseError, parse_dice
from pf1e_simulator.rng import ScriptedRng

# ── Parsing: accepted forms ──────────────────────────────────────────────────

class TestParseDiceAccepted:
    """Given: valid dice notation strings
    When: parse_dice is called
    Then: returns a correct DiceExpr."""

    def test_2d6_plus_3(self) -> None:
        expr = parse_dice("2d6+3")
        assert expr.count == 2
        assert expr.sides == 6
        assert expr.bonus == 3

    def test_2d6_minus_1(self) -> None:
        expr = parse_dice("2d6-1")
        assert expr.count == 2
        assert expr.sides == 6
        assert expr.bonus == -1

    def test_1d4_no_bonus(self) -> None:
        expr = parse_dice("1d4")
        assert expr.count == 1
        assert expr.sides == 4
        assert expr.bonus == 0

    def test_d20_shorthand(self) -> None:
        expr = parse_dice("d20")
        assert expr.count == 1
        assert expr.sides == 20
        assert expr.bonus == 0

    def test_flat_integer(self) -> None:
        expr = parse_dice("7")
        assert expr.count == 0
        assert expr.sides == 0
        assert expr.bonus == 7

    def test_flat_negative(self) -> None:
        expr = parse_dice("-5")
        assert expr.count == 0
        assert expr.sides == 0
        assert expr.bonus == -5

    def test_flat_zero(self) -> None:
        expr = parse_dice("0")
        assert expr.count == 0
        assert expr.sides == 0
        assert expr.bonus == 0

    def test_large_dice(self) -> None:
        expr = parse_dice("10d100+50")
        assert expr.count == 10
        assert expr.sides == 100
        assert expr.bonus == 50

    def test_whitespace_stripped(self) -> None:
        expr = parse_dice("  2d6+3  ")
        assert expr.count == 2
        assert expr.sides == 6
        assert expr.bonus == 3


# ── Parsing: rejected forms ──────────────────────────────────────────────────

class TestParseDiceRejected:
    """Given: invalid dice notation strings
    When: parse_dice is called
    Then: raises DiceParseError carrying the offending text."""

    @pytest.mark.parametrize(
        "text",
        [
            "",
            "sizeRoll(1, 8, @size)",
            "2x6",
            "1d",
            "d",
            "2d0",
            "0d6",
            "2d 6",
            "abc",
            "d6+",
        ],
    )
    def test_rejects_invalid(self, text: str) -> None:
        with pytest.raises(DiceParseError) as exc_info:
            parse_dice(text)
        assert text in str(exc_info.value)


# ── Mean calculation ─────────────────────────────────────────────────────────

class TestDiceExprMean:
    """Given: a DiceExpr
    When: mean() is called
    Then: returns count * (sides + 1) / 2 + bonus."""

    def test_2d6_plus_3_mean(self) -> None:
        assert DiceExpr(count=2, sides=6, bonus=3).mean() == 10.0

    def test_1d4_mean(self) -> None:
        assert DiceExpr(count=1, sides=4, bonus=0).mean() == 2.5

    def test_d20_mean(self) -> None:
        assert DiceExpr(count=1, sides=20, bonus=0).mean() == 10.5

    def test_flat_bonus_mean(self) -> None:
        assert DiceExpr(count=0, sides=0, bonus=7).mean() == 7.0

    def test_negative_bonus_mean(self) -> None:
        assert DiceExpr(count=1, sides=6, bonus=-2).mean() == 1.5


# ── Roll: uses RNG correctly ────────────────────────────────────────────────

class TestDiceExprRoll:
    """Given: a DiceExpr and a ScriptedRng
    When: roll(rng) is called
    Then: consumes exactly `count` values from the RNG."""

    def test_2d6_plus_3_uses_two_rolls(self) -> None:
        rng = ScriptedRng([3, 4])
        expr = DiceExpr(count=2, sides=6, bonus=3)
        assert expr.roll(rng) == 3 + 4 + 3  # 10

    def test_1d20_uses_one_roll(self) -> None:
        rng = ScriptedRng([15])
        expr = DiceExpr(count=1, sides=20, bonus=0)
        assert expr.roll(rng) == 15

    def test_flat_bonus_uses_no_rolls(self) -> None:
        rng = ScriptedRng([])
        expr = DiceExpr(count=0, sides=0, bonus=7)
        assert expr.roll(rng) == 7

    def test_negative_bonus(self) -> None:
        rng = ScriptedRng([6, 6])
        expr = DiceExpr(count=2, sides=6, bonus=-1)
        assert expr.roll(rng) == 6 + 6 - 1  # 11


# ── Distribution: exhaustive 2d6+3 ──────────────────────────────────────────

class TestDiceExprDistribution:
    """Given: 2d6+3 enumerated over all 36 face pairs
    When: probabilities are computed
    Then: they match the theoretical distribution."""

    def test_2d6_plus_3_distribution(self) -> None:
        totals: Counter[int] = Counter()
        for d1 in range(1, 7):
            for d2 in range(1, 7):
                rng = ScriptedRng([d1, d2])
                expr = DiceExpr(count=2, sides=6, bonus=3)
                totals[expr.roll(rng)] += 1

        assert totals[5] == 1   # (1,1) → 1+1+3=5
        assert totals[7] == 3   # (1,3),(2,2),(3,1) → sum=4+3=7
        assert totals[10] == 6  # sum=7: (1,6)..(6,1) → 7+3=10
        assert totals[15] == 1  # (6,6) → 12+3=15
        assert sum(totals.values()) == 36


# ── DiceExpr.flat factory ────────────────────────────────────────────────────

class TestDiceExprFlat:
    """Given: DiceExpr.flat(n)
    When: called with any integer
    Then: returns count=0, sides=0, bonus=n."""

    def test_flat_positive(self) -> None:
        expr = DiceExpr.flat(10)
        assert expr.count == 0
        assert expr.sides == 0
        assert expr.bonus == 10
        assert expr.mean() == 10.0

    def test_flat_negative(self) -> None:
        expr = DiceExpr.flat(-3)
        assert expr.bonus == -3

    def test_flat_roll_returns_bonus(self) -> None:
        rng = ScriptedRng([])
        expr = DiceExpr.flat(5)
        assert expr.roll(rng) == 5