‘Faster Than Light: Nomad’ Probabilities 2

#!/usr/bin/env python3

import argparse

from functools import cache
from fractions import Fraction
from itertools import combinations_with_replacement, chain, product

import tabulate

NKEEP: int = 2
NBONUS: int = 3
NSIDES: int = 6
TARGET: int = 8

SetHistogram = dict[tuple[int, ...], Fraction]

Histogram = dict[int, Fraction]

Cell = Fraction | int | None

Table = list[list[Cell]]


def new_tuple(t: tuple[int, ...], d: int, low: bool) -> tuple[int, ...]:
    temp_t = tuple(sorted((d,) + t))
    if low:
        return temp_t[:-1]
    return temp_t[1:]


@cache
def set_histogram(ndice: int, nsides: int, nkeep: int, low: bool) -> SetHistogram:
    assert nkeep <= ndice
    result: SetHistogram = {}

    for t in combinations_with_replacement(range(1, nsides + 1), nkeep):
        result[t] = Fraction(0, 1)

    if nkeep == ndice:
        # Create simple set histogram
        p: Fraction = Fraction(1, nsides**nkeep)
        for t in product(range(1, nsides + 1), repeat=nkeep):
            modt: tuple[int, ...] = tuple(sorted(t))
            result[modt] += p

        return result

    # Multiply lesser order set_histogram by one die
    lesser: SetHistogram = set_histogram(ndice - 1, nsides, nkeep, low)
    onedie_p: Fraction = Fraction(1, nsides)
    for d, t in product(range(1, nsides + 1), lesser):
        new_t: tuple[int, ...] = new_tuple(t, d, low)
        result[new_t] += lesser[t] * onedie_p

    return result


def histogram(ndice: int, nsides: int, nkeep: int, low: bool) -> Histogram:
    result: Histogram = {}
    sethist: SetHistogram = set_histogram(ndice, nsides, nkeep, low)

    # Sum the sets to provide a Histogram
    for n in range(nkeep, nsides * nkeep + 1):
        result[n] = Fraction(0, 1)
    for t in sethist:
        sum_n: int = sum(t)
        result[sum_n] += sethist[t]
    return result


def cumulative(h: Histogram) -> Histogram:
    result: Histogram = {}
    counter: Fraction = Fraction(0, 1)
    for key in reversed(sorted(list(h))):
        counter += h[key]
        result[key] = counter
    return result


def format_pct(h: dict, i: int) -> Cell:
    if i not in h:
        return None
    return h[i] * 100


def skill(k: int, target: int, neg: bool = False) -> Cell:
    if neg:
        return target - k
    return target - k if k <= target else None


def collate(hl: list[Histogram], target: int = TARGET, neg: bool = False) -> Table:
    allkeys: list[int] = sorted(set(chain.from_iterable([list(h) for h in hl])))

    return [
        [k, skill(k, target, neg)] + [format_pct(h, k) for h in hl] for k in allkeys
    ]


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate a probability table for roll-and-keep mechanics."
    )
    parser.add_argument(
        "-k",
        "--keep",
        help=f"number of dice to keep (default {NKEEP})",
        default=NKEEP,
        type=int,
    )
    parser.add_argument(
        "-b",
        "--bonus",
        help=f"max number of bonus/penalty dice (default {NBONUS})",
        default=NBONUS,
        type=int,
    )
    parser.add_argument(
        "-d",
        "--sides",
        help=f"number of sides (default {NSIDES})",
        default=NSIDES,
        type=int,
    )
    parser.add_argument(
        "-t",
        "--target",
        help=f"standard target number (default {TARGET})",
        default=TARGET,
        type=int,
    )
    parser.add_argument(
        "-n",
        "--negative-skill",
        help="show negative skill numbers",
        action="store_true",
    )
    args = parser.parse_args()

    keep: int = args.keep
    bonus: int = args.bonus
    sides: int = args.sides

    headers = ["TN", "Skill"] + [f"{n:+}D" for n in range(-bonus, bonus + 1)]

    table = collate(
        [
            cumulative(histogram(ntotal, sides, nkeep=keep, low=True))
            for ntotal in range(keep + bonus, keep, -1)
        ]
        + [
            cumulative(histogram(ntotal, sides, nkeep=keep, low=False))
            for ntotal in range(keep, keep + bonus + 1)
        ],
        target=args.target,
        neg=args.negative_skill,
    )

    print(tabulate.tabulate(table, headers=headers, tablefmt="pipe", floatfmt=".2f"))


if __name__ == "__main__":
    main()