reasoning-gym/reasoning_gym/probability/probability_problems.py

import math
import random
from dataclasses import dataclass, field
from fractions import Fraction
from typing import Any, Optional

from ..coaching import BaseCurriculum, RangeAttributeDefinition
from ..factory import ProceduralDataset, register_dataset

DATASET_NAME = "probability_problems"

TASK_TYPES = (
    "independent_events",
    "compound_events",
    "total_probability",
    "bayes_theorem",
    "binomial_probability",
    "binomial_stats",
    "geometric_series",
    "geometric_region",
    "expectation_variance",
)


@dataclass
class ProbabilityProblemsConfig:
    min_n: int = 3
    max_n: int = 10
    task_types: tuple[str, ...] = TASK_TYPES
    task_weights: list[float] = field(
        default_factory=lambda: [
            0.12,
            0.11,
            0.12,
            0.12,
            0.11,
            0.10,
            0.11,
            0.10,
            0.11,
        ]
    )
    seed: Optional[int] = None
    size: int = 500

    def validate(self) -> None:
        assert self.size > 0, "size must be positive"
        assert self.min_n >= 2, "min_n must be >= 2"
        assert self.max_n >= self.min_n, "max_n must be >= min_n"
        assert len(self.task_types) > 0, "must have at least one task type"
        assert all(t in TASK_TYPES for t in self.task_types), "invalid task type"
        assert len(self.task_weights) == len(self.task_types), "weights must match types"


class ProbabilityProblemsDataset(ProceduralDataset):
    def __init__(self, config: ProbabilityProblemsConfig):
        super().__init__(config=config, seed=config.seed, size=config.size)

    def _rand_prob(self, rng: random.Random) -> Fraction:
        """Generate a random probability as a simple fraction in (0, 1)."""
        denom = rng.choice([2, 3, 4, 5, 6, 8, 10])
        numer = rng.randint(1, denom - 1)
        return Fraction(numer, denom)

    # --- Section 1: Conditional Probability & Multiplication Theorem ---

    def _make_independent_events(self, rng: random.Random) -> dict:
        pa = self._rand_prob(rng)
        pb = self._rand_prob(rng)
        variant = rng.choice(["intersection", "union", "neither"])

        if variant == "intersection":
            answer = pa * pb
            question = (
                f"Events A and B are independent with P(A) = {pa} and P(B) = {pb}. "
                f"What is P(A and B)? Give your answer as a simplified fraction."
            )
        elif variant == "union":
            answer = pa + pb - pa * pb
            question = (
                f"Events A and B are independent with P(A) = {pa} and P(B) = {pb}. "
                f"What is P(A or B)? Give your answer as a simplified fraction."
            )
        else:
            answer = (1 - pa) * (1 - pb)
            question = (
                f"Events A and B are independent with P(A) = {pa} and P(B) = {pb}. "
                f"What is the probability that neither A nor B occurs? "
                f"Give your answer as a simplified fraction."
            )
        return {"question": question, "answer": str(answer), "task_type": "independent_events"}

    def _make_compound_events(self, rng: random.Random) -> dict:
        total = rng.randint(max(4, self.config.min_n), max(4, self.config.max_n))
        color_a_count = rng.randint(2, total - 2)
        color_b_count = total - color_a_count

        colors = ["red", "blue", "green", "white", "black"]
        color_a = rng.choice(colors)
        color_b = rng.choice([c for c in colors if c != color_a])

        seq = rng.choice(["ab", "ba"])
        if seq == "ab":
            prob = Fraction(color_a_count, total) * Fraction(color_b_count, total - 1)
            seq_desc = f"the first is {color_a} and the second is {color_b}"
        else:
            prob = Fraction(color_b_count, total) * Fraction(color_a_count, total - 1)
            seq_desc = f"the first is {color_b} and the second is {color_a}"

        question = (
            f"A bag contains {color_a_count} {color_a} balls and {color_b_count} {color_b} balls. "
            f"You draw 2 balls one after another without replacement. "
            f"What is the probability that {seq_desc}? "
            f"Give your answer as a simplified fraction."
        )
        return {"question": question, "answer": str(prob), "task_type": "compound_events"}

    # --- Section 2: Total Probability & Bayes' Theorem ---

    def _make_total_probability(self, rng: random.Random) -> dict:
        num_bags = rng.randint(2, 3)
        bags = []
        for _ in range(num_bags):
            red = rng.randint(1, self.config.max_n)
            blue = rng.randint(1, self.config.max_n)
            bags.append((red, blue))

        p_bag = Fraction(1, num_bags)
        p_red = Fraction(0)
        for red, blue in bags:
            p_red += p_bag * Fraction(red, red + blue)

        bag_desc = ". ".join(f"Bag {i + 1} contains {r} red and {b} blue balls" for i, (r, b) in enumerate(bags))
        question = (
            f"{bag_desc}. "
            f"One bag is chosen uniformly at random and a ball is drawn from it. "
            f"What is the probability the ball is red? "
            f"Give your answer as a simplified fraction."
        )
        return {"question": question, "answer": str(p_red), "task_type": "total_probability"}

    def _make_bayes_theorem(self, rng: random.Random) -> dict:
        num_bags = rng.randint(2, 3)
        bags = []
        for _ in range(num_bags):
            red = rng.randint(1, self.config.max_n)
            blue = rng.randint(1, self.config.max_n)
            bags.append((red, blue))

        target_bag = rng.randint(0, num_bags - 1)

        p_bag = Fraction(1, num_bags)
        p_red = Fraction(0)
        for red, blue in bags:
            p_red += p_bag * Fraction(red, red + blue)

        red_t, blue_t = bags[target_bag]
        p_red_given_target = Fraction(red_t, red_t + blue_t)
        p_target_given_red = (p_bag * p_red_given_target) / p_red

        bag_desc = ". ".join(f"Bag {i + 1} contains {r} red and {b} blue balls" for i, (r, b) in enumerate(bags))
        question = (
            f"{bag_desc}. "
            f"One bag is chosen uniformly at random and a ball is drawn. The ball is red. "
            f"What is the probability that it came from Bag {target_bag + 1}? "
            f"Give your answer as a simplified fraction."
        )
        return {"question": question, "answer": str(p_target_given_red), "task_type": "bayes_theorem"}

    # --- Section 3: Probability Distributions ---

    def _make_binomial_probability(self, rng: random.Random) -> dict:
        p_choices = [
            Fraction(1, 6),
            Fraction(1, 4),
            Fraction(1, 3),
            Fraction(1, 2),
            Fraction(2, 3),
            Fraction(3, 4),
        ]
        p = rng.choice(p_choices)
        q = 1 - p
        n = rng.randint(self.config.min_n, min(self.config.max_n, 8))
        r = rng.randint(0, n)

        prob = Fraction(math.comb(n, r)) * (p**r) * (q ** (n - r))
        question = (
            f"A biased coin has a probability of heads equal to {p}. "
            f"If it is flipped {n} times, what is the probability of getting exactly {r} heads? "
            f"Give your answer as a simplified fraction."
        )
        return {"question": question, "answer": str(prob), "task_type": "binomial_probability"}

    def _make_binomial_stats(self, rng: random.Random) -> dict:
        p_choices = [
            Fraction(1, 6),
            Fraction(1, 4),
            Fraction(1, 3),
            Fraction(1, 2),
            Fraction(2, 3),
            Fraction(3, 4),
        ]
        p = rng.choice(p_choices)
        q = 1 - p
        n = rng.randint(self.config.min_n, self.config.max_n)
        variant = rng.choice(["mean", "variance"])

        if variant == "mean":
            answer = Fraction(n) * p
            question = (
                f"A random variable X follows a binomial distribution with {n} trials "
                f"and success probability {p}. What is E(X), the expected value? "
                f"Give your answer as a simplified fraction or integer."
            )
        else:
            answer = Fraction(n) * p * q
            question = (
                f"A random variable X follows a binomial distribution with {n} trials "
                f"and success probability {p}. What is Var(X), the variance? "
                f"Give your answer as a simplified fraction or integer."
            )
        return {"question": question, "answer": str(answer), "task_type": "binomial_stats"}

    # --- Section 4: Geometric Probability ---

    def _make_geometric_series(self, rng: random.Random) -> dict:
        p_choices = [
            Fraction(1, 6),
            Fraction(1, 5),
            Fraction(1, 4),
            Fraction(1, 3),
            Fraction(1, 2),
        ]
        p = rng.choice(p_choices)
        q = rng.choice(p_choices)

        # A and B alternate; A goes first.
        # P(A wins) = p / (1 - (1-p)(1-q))  via infinite geometric series
        answer = p / (1 - (1 - p) * (1 - q))

        name_a = rng.choice(["Alice", "Arun", "Alex"])
        name_b = rng.choice(["Bob", "Bala", "Beth"])

        question = (
            f"{name_a} and {name_b} play a game where they take alternate turns, "
            f"with {name_a} going first. On each of her turns, {name_a} has a probability "
            f"of {p} of winning the game. If {name_a} does not win, {name_b} then has a "
            f"probability of {q} of winning on his turn. If neither wins, the process "
            f"repeats. What is the probability that {name_a} wins the game? "
            f"Give your answer as a simplified fraction."
        )
        return {"question": question, "answer": str(answer), "task_type": "geometric_series"}

    def _make_geometric_region(self, rng: random.Random) -> dict:
        a = rng.randint(max(2, self.config.min_n), self.config.max_n)
        variant = rng.choice(["leq", "geq"])

        if variant == "leq":
            c = rng.randint(1, a)
            answer = Fraction(c * c, 2 * a * a)
            question = (
                f"Two numbers x and y are each chosen uniformly at random from [0, {a}]. "
                f"What is the probability that x + y <= {c}? "
                f"Give your answer as a simplified fraction."
            )
        else:
            c = rng.randint(a + 1, 2 * a - 1)
            side = 2 * a - c
            answer = Fraction(side * side, 2 * a * a)
            question = (
                f"Two numbers x and y are each chosen uniformly at random from [0, {a}]. "
                f"What is the probability that x + y >= {c}? "
                f"Give your answer as a simplified fraction."
            )
        return {"question": question, "answer": str(answer), "task_type": "geometric_region"}

    # --- Section 5: Random Variables & Expectation ---

    def _make_expectation_variance(self, rng: random.Random) -> dict:
        k = rng.randint(3, 5)
        outcomes = sorted(rng.sample(range(1, 11), k))
        weights = [rng.randint(1, 10) for _ in range(k)]
        total_weight = sum(weights)
        probs = [Fraction(w, total_weight) for w in weights]

        variant = rng.choice(["expectation", "variance"])

        ex = sum(Fraction(x) * p for x, p in zip(outcomes, probs))

        if variant == "expectation":
            answer = ex
            stat_name = "E(X), the expected value"
        else:
            ex2 = sum(Fraction(x * x) * p for x, p in zip(outcomes, probs))
            answer = ex2 - ex * ex
            stat_name = "Var(X), the variance"

        table_lines = " | ".join(f"P(X={x}) = {p}" for x, p in zip(outcomes, probs))
        question = (
            f"A discrete random variable X has the following probability distribution: "
            f"{table_lines}. "
            f"What is {stat_name}? "
            f"Give your answer as a simplified fraction or integer."
        )
        return {"question": question, "answer": str(answer), "task_type": "expectation_variance"}

    def __getitem__(self, idx: int) -> dict:
        rng = random.Random(self.seed + idx)
        task_type = rng.choices(self.config.task_types, weights=self.config.task_weights, k=1)[0]

        generators = {
            "independent_events": self._make_independent_events,
            "compound_events": self._make_compound_events,
            "total_probability": self._make_total_probability,
            "bayes_theorem": self._make_bayes_theorem,
            "binomial_probability": self._make_binomial_probability,
            "binomial_stats": self._make_binomial_stats,
            "geometric_series": self._make_geometric_series,
            "geometric_region": self._make_geometric_region,
            "expectation_variance": self._make_expectation_variance,
        }
        result = generators[task_type](rng)
        return {
            "question": result["question"],
            "answer": result["answer"],
            "metadata": {
                "source_dataset": DATASET_NAME,
                "source_index": idx,
                "task_type": result["task_type"],
                "difficulty": {"min_n": self.config.min_n, "max_n": self.config.max_n},
            },
        }

    def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
        if answer is None:
            return 0.0
        oracle = entry["answer"]
        if answer.strip() == oracle.strip():
            return 1.0
        try:
            ans_frac = Fraction(answer.strip())
            oracle_frac = Fraction(oracle.strip())
            if ans_frac == oracle_frac:
                return 1.0
            diff = abs(float(ans_frac) - float(oracle_frac))
            if diff < 1e-4:
                return 0.9
            if diff < 1e-2:
                return 0.5
            return 0.0
        except (ValueError, ZeroDivisionError):
            return 0.0


class ProbabilityProblemsCurriculum(BaseCurriculum):
    def __init__(self):
        super().__init__(ProbabilityProblemsCurriculum.__name__, ProbabilityProblemsConfig)
        self._define_attributes(
            RangeAttributeDefinition(
                name="n_range",
                levels=[3, 5, 10, 15, 20],
                lower_field_name="min_n",
                upper_field_name="max_n",
                description="Range for n in probability problems",
            ),
        )


register_dataset(DATASET_NAME, ProbabilityProblemsDataset, ProbabilityProblemsConfig, ProbabilityProblemsCurriculum)