gamble simulator

import numpy as np


class Gamble:
    def __init__(self, value: float | list[tuple[float, "Gamble"]]):
        self.value = value

    def sample(self) -> float:
        if isinstance(self.value, float):
            return self.value
        probs = [v[0] for v in self.value]
        idx = np.random.choice(len(probs), p=probs)
        return self.value[idx][1].sample()

    def samples(self, n: int) -> np.ndarray:
        return np.array([self.sample() for _ in range(n)])

    def __or__(self, other: "Gamble") -> "Gamble":
        return Gamble([(0.5, self), (0.5, other)])

    def __mul__(self, other: float) -> "Gamble":
        if isinstance(other, float):
            return Gamble(self.value * other)
        return Gamble([(v[0] * other, v[1]) for v in self.value])


const = Gamble(1.0)
sure = const * 1.0
risky = (const * 1e-100) | (const * 1e20)
sure_or_sure = sure | sure
risky_or_sure = risky | sure


def U(g: Gamble, k: int, n: int = 10**4) -> float:
    outcomes = 1
    for _ in range(k):
        outcomes += g.samples(n=n)
    outcomes /= k
    return np.log(outcomes).mean()


K = 2

print(U(sure, k=K))
print(U(risky, k=K))

print(U(sure_or_sure, k=K))
print(U(risky_or_sure, k=K))