yahtzee-engine 0.1.0

//! The exact solver: expected values for every state, keep, and category
//! under optimal play.

use crate::state::max_upper;
use crate::tables::{DiceTables, KEEPS, ROLLS};
use crate::{Categories, Category, Dice, Keep, State, Strategy, TurnAction, View};

/// An exact expectimax solver over the full game.
///
/// The value table is keyed by [`State::index`] and filled by backward
/// induction: a state's expected further score is the mean over the
/// turn's roll/keep/roll/keep/roll/score lattice of the best reward plus
/// the value of the resulting state.  From the empty card this comes to
/// an optimal solo expectation of 254.5877 points under the official
/// forced-joker rules this crate implements.  The widely cited 254.5896
/// (Verhoeff, 1999) uses a laxer convention — an extra Yahtzee may be
/// scored in any open box, with joker values gated on the matching
/// upper box being filled — which this engine reproduces exactly when
/// the forcing in [`State::legal_categories`] is lifted.
///
/// Construction is cheap; values are computed on demand.  Querying a
/// mid-game state solves only the states reachable from it, so "what is
/// best here" costs a fraction of a full [`solve`](Self::solve).  All
/// queries break ties deterministically (first category in card order,
/// then first keep in canonical order), so replays are stable.
#[derive(Debug, Clone)]
pub struct Solver {
    values: Vec<f64>,
    tables: DiceTables,
}

/// Per-roll menu of legal categories and box values for one mask.
///
/// Legality and box values depend only on the scored set and the dice —
/// not on the upper subtotal or the Yahtzee-50 flag — so the menu is
/// computed once per mask and shared by all its table slots.  Entries
/// come straight from [`State::apply`], keeping the solver on the same
/// rules as the game.
type Menu = Vec<Vec<(Category, u8)>>;

fn menu_for(tables: &DiceTables, mask: Categories) -> Menu {
    let probe = State::new(mask, 0, false);
    (0..ROLLS)
        .map(|r| {
            let dice = tables.roll(r);
            probe
                .legal_categories(dice)
                .iter()
                .map(|c| (c, probe.apply(c, dice).expect("legal").value))
                .collect()
        })
        .collect()
}

/// Expected score of each roll when the turn must end now: the best
/// category's reward plus the value of the state it leads to.
fn final_layer(values: &[f64], tables: &DiceTables, state: State, menu: &Menu) -> [f64; ROLLS] {
    let scored = state.scored().bits() as usize;
    let upper = state.upper();
    let flag = state.yahtzee_50();
    let mut e_roll = [f64::NEG_INFINITY; ROLLS];
    for (r, options) in menu.iter().enumerate() {
        let bonus_eligible = flag && tables.roll(r).yahtzee_face().is_some();
        let mut best = f64::NEG_INFINITY;
        for &(category, value) in options {
            // The same bookkeeping as `State::apply`, on table indices.
            let (next_upper, upper_bonus) = if (category as u8) < 6 {
                let subtotal = upper + value;
                (subtotal.min(63), upper < 63 && subtotal >= 63)
            } else {
                (upper, false)
            };
            let next_flag = flag || (category == Category::Yahtzee && value == 50);
            let next = scored
                | 1 << category as usize
                | (next_upper as usize) << 13
                | usize::from(next_flag) << 19;
            let reward = u16::from(value)
                + if bonus_eligible { 100 } else { 0 }
                + if upper_bonus { 35 } else { 0 };
            let v = f64::from(reward) + values[next];
            debug_assert!(!v.is_nan(), "unsolved successor of {state:?}");
            if v > best {
                best = v;
            }
        }
        e_roll[r] = best;
    }
    e_roll
}

/// One reroll of backward induction: value every keep as the mean over
/// its rolls, then every roll as its best keep.  The full five-die keep
/// carries the "stop rolling" option through unchanged.
fn sweep(tables: &DiceTables, e_roll: &[f64; ROLLS]) -> [f64; ROLLS] {
    let mut e_keep = [0.0; KEEPS];
    for (k, slot) in e_keep.iter_mut().enumerate() {
        *slot = tables
            .keep_successors(k)
            .iter()
            .map(|&(r, p)| p * e_roll[usize::from(r)])
            .sum();
    }
    let mut out = [f64::NEG_INFINITY; ROLLS];
    for (r, slot) in out.iter_mut().enumerate() {
        *slot = tables
            .roll_keeps(r)
            .iter()
            .map(|&k| e_keep[usize::from(k)])
            .fold(f64::NEG_INFINITY, f64::max);
    }
    out
}

/// The expected further score from the start of a turn in `state`.
fn turn_value(values: &[f64], tables: &DiceTables, state: State, menu: &Menu) -> f64 {
    let mut e_roll = final_layer(values, tables, state, menu);
    e_roll = sweep(tables, &e_roll);
    e_roll = sweep(tables, &e_roll);
    tables
        .roll_prob()
        .iter()
        .zip(e_roll)
        .map(|(p, v)| p * v)
        .sum()
}

/// Values for every table slot of one mask, assuming every strict
/// superset mask is already solved.
fn solve_mask(values: &[f64], tables: &DiceTables, mask: Categories) -> Vec<(usize, f64)> {
    let variants = |mask: Categories| {
        [false, true]
            .into_iter()
            .filter(move |&flag| !flag || mask.contains(Category::Yahtzee))
            .flat_map(move |flag| {
                (0..=max_upper(mask)).map(move |upper| State::new(mask, upper, flag))
            })
    };
    if mask == Categories::ALL {
        return variants(mask).map(|s| (s.index(), 0.0)).collect();
    }
    let menu = menu_for(tables, mask);
    variants(mask)
        .map(|s| (s.index(), turn_value(values, tables, s, &menu)))
        .collect()
}

impl Solver {
    /// Builds the dice tables; no states are solved yet.
    #[must_use]
    pub fn new() -> Self {
        let mut values = vec![f64::NAN; 1 << 20];
        for upper in 0..=63 {
            for flag in [false, true] {
                values[State::new(Categories::ALL, upper, flag).index()] = 0.0;
            }
        }
        Self {
            values,
            tables: DiceTables::new(),
        }
    }

    /// Whether [`value`](Self::value) for this state is already computed.
    #[must_use]
    pub fn is_solved(&self, state: State) -> bool {
        !self.values[state.index()].is_nan()
    }

    /// Solves every state with exactly `filled` categories scored.
    ///
    /// Tiers must be solved from 13 down to 0; [`solve`](Self::solve)
    /// does exactly that.  Exposed so incremental front ends can report
    /// progress between tiers.  With the `parallel` feature the tier is
    /// solved across threads, bit-identical to the serial build.
    ///
    /// # Panics
    ///
    /// Panics if `filled > 13`, or if the tiers above `filled` have not
    /// been solved yet — solving out of order would silently poison the
    /// table otherwise.
    pub fn solve_tier(&mut self, filled: u32) {
        assert!(filled <= 13, "a card has 13 categories");
        let masks: Vec<Categories> = (0..=Categories::ALL.bits())
            .filter(|bits| bits.count_ones() == filled)
            .map(|bits| Categories::from_bits(bits).expect("13-bit masks"))
            .collect();
        let (values, tables) = (&self.values, &self.tables);
        #[cfg(feature = "parallel")]
        let results: Vec<(usize, f64)> = {
            use rayon::prelude::*;
            masks
                .par_iter()
                .flat_map_iter(|&mask| solve_mask(values, tables, mask))
                .collect()
        };
        #[cfg(not(feature = "parallel"))]
        let results: Vec<(usize, f64)> = masks
            .iter()
            .flat_map(|&mask| solve_mask(values, tables, mask))
            .collect();
        for (index, value) in results {
            // A finite value needs every successor solved; an unsolved
            // successor's NaN turns the whole state's value non-finite.
            assert!(
                value.is_finite(),
                "solve_tier({filled}) before its successor tiers: solve from 13 down to 0"
            );
            self.values[index] = value;
        }
    }

    /// Solves the whole game, all reachable states, bottom-up.
    ///
    /// Takes seconds in a release build (use the `parallel` feature to
    /// spread it across cores); afterwards every query is a table lookup
    /// plus one turn evaluation.
    pub fn solve(&mut self) {
        for filled in (0..=13).rev() {
            self.solve_tier(filled);
        }
    }

    /// The expected further score from the start of a turn in `state`,
    /// under optimal play.
    ///
    /// Solves lazily: only states reachable from `state` are computed,
    /// so early queries from a part-filled card are much cheaper than a
    /// full [`solve`](Self::solve).
    pub fn value(&mut self, state: State) -> f64 {
        if self.is_solved(state) {
            return self.values[state.index()];
        }
        let scored = state.scored().bits();
        let complement = (!state.scored()).bits();
        // Every subset of the open categories, largest masks first, so
        // each mask finds its successors already solved.
        let mut subsets = Vec::with_capacity(1 << complement.count_ones());
        let mut bits = complement;
        loop {
            subsets.push(bits);
            if bits == 0 {
                break;
            }
            bits = (bits - 1) & complement;
        }
        subsets.sort_by_key(|bits| core::cmp::Reverse(bits.count_ones()));
        for bits in subsets {
            let mask = Categories::from_bits(scored | bits).expect("13-bit masks");
            // `solve_mask` fills every slot of a mask, so one probe slot
            // tells whether the whole mask is done.
            if self.values[State::new(mask, 0, false).index()].is_nan() {
                for (index, value) in solve_mask(&self.values, &self.tables, mask) {
                    self.values[index] = value;
                }
            }
        }
        self.values[state.index()]
    }

    /// The expected further score of scoring `dice` in `category` now
    /// and playing on optimally: the write, its bonuses, and everything
    /// after — points already on the card are not included.
    ///
    /// Returns [`None`] if the category is illegal here; see
    /// [`State::legal_categories`].
    pub fn category_ev(&mut self, state: State, dice: Dice, category: Category) -> Option<f64> {
        let delta = state.apply(category, dice)?;
        Some(f64::from(delta.reward()) + self.value(delta.next))
    }

    /// The expected further score of holding `keep` and rerolling the
    /// rest, with `rolls_left` rolls remaining (clamped to the turn's
    /// two rerolls).
    ///
    /// Returns [`None`] if the card is full, no rolls remain, or `keep`
    /// is not part of `dice`.
    pub fn keep_ev(&mut self, state: State, dice: Dice, keep: Keep, rolls_left: u8) -> Option<f64> {
        if state.is_full() || rolls_left == 0 || !dice.contains(keep) {
            return None;
        }
        self.value(state);
        let menu = menu_for(&self.tables, state.scored());
        let mut e_roll = final_layer(&self.values, &self.tables, state, &menu);
        for _ in 1..rolls_left.min(2) {
            e_roll = sweep(&self.tables, &e_roll);
        }
        let keep_id = self.tables.keep_id(keep);
        Some(
            self.tables
                .keep_successors(keep_id)
                .iter()
                .map(|&(r, p)| p * e_roll[usize::from(r)])
                .sum(),
        )
    }

    /// The best category for `dice` when the turn must end now.
    ///
    /// # Panics
    ///
    /// Panics if the card is already full.
    pub fn best_category(&mut self, state: State, dice: Dice) -> Category {
        assert!(!state.is_full(), "no category is open on a full card");
        self.value(state);
        state
            .legal_categories(dice)
            .iter()
            .map(|c| {
                let delta = state.apply(c, dice).expect("legal");
                (
                    c,
                    f64::from(delta.reward()) + self.values[delta.next.index()],
                )
            })
            .fold(None, |best: Option<(Category, f64)>, (c, v)| match best {
                Some((_, bv)) if bv >= v => best,
                _ => Some((c, v)),
            })
            .expect("legal categories are non-empty until the card is full")
            .0
    }

    /// The optimal action for `dice` with `rolls_left` rolls remaining
    /// (clamped to the turn's two rerolls).  Scoring wins ties against
    /// rerolling.
    ///
    /// # Panics
    ///
    /// Panics if the card is already full.
    pub fn best_action(&mut self, state: State, dice: Dice, rolls_left: u8) -> TurnAction {
        let category = self.best_category(state, dice);
        if rolls_left == 0 {
            return TurnAction::Score(category);
        }
        let menu = menu_for(&self.tables, state.scored());
        let mut e_roll = final_layer(&self.values, &self.tables, state, &menu);
        let score_now = e_roll[self.tables.roll_id(dice)];
        for _ in 1..rolls_left.min(2) {
            e_roll = sweep(&self.tables, &e_roll);
        }
        let mut best = TurnAction::Score(category);
        let mut best_value = score_now;
        for keep_id in 0..KEEPS {
            let keep = self.tables.keep(keep_id);
            if !dice.contains(keep) {
                continue;
            }
            let value: f64 = self
                .tables
                .keep_successors(keep_id)
                .iter()
                .map(|&(r, p)| p * e_roll[usize::from(r)])
                .sum();
            if value > best_value {
                best_value = value;
                best = TurnAction::Reroll(keep);
            }
        }
        best
    }
}

impl Default for Solver {
    fn default() -> Self {
        Self::new()
    }
}

/// A [`Strategy`] that plays perfectly, backed by a [`Solver`].
///
/// The default construction solves lazily: the first decision of a fresh
/// game triggers close to a full solve, which takes seconds in a release
/// build.  Use [`presolved`](Self::presolved) to pay that cost up front,
/// or [`from_solver`](Self::from_solver) to share a table you already
/// built.
#[derive(Debug, Clone, Default)]
pub struct OptimalBot {
    solver: Solver,
}

impl OptimalBot {
    /// A bot that solves on demand.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// A bot with the whole game solved up front; see [`Solver::solve`].
    #[must_use]
    pub fn presolved() -> Self {
        let mut solver = Solver::new();
        solver.solve();
        Self { solver }
    }

    /// A bot backed by an existing solver, keeping its solved states.
    #[must_use]
    pub const fn from_solver(solver: Solver) -> Self {
        Self { solver }
    }

    /// The underlying solver, for inspection or further queries.
    pub const fn solver(&mut self) -> &mut Solver {
        &mut self.solver
    }
}

impl Strategy for OptimalBot {
    fn choose_action(&mut self, view: &View<'_>) -> TurnAction {
        self.solver
            .best_action(view.state(), view.dice(), view.rolls_left())
    }

    fn choose_category(&mut self, view: &View<'_>) -> Category {
        self.solver.best_category(view.state(), view.dice())
    }

    fn name(&self) -> &str {
        "optimal"
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn dice(s: &str) -> Dice {
        s.parse().expect("a valid roll")
    }

    /// With only Chance open, optimal play keeps 5-6 on the first two
    /// rolls and 4-6 on the last, worth exactly 70/3 by independence of
    /// the five dice.
    #[test]
    fn chance_alone_is_worth_seventy_thirds() {
        let mut solver = Solver::new();
        let state = State::new(!Category::Chance.bit(), 0, false);
        let value = solver.value(state);
        assert!((value - 70.0 / 3.0).abs() < 1e-9, "got {value}");
    }

    /// With only the Yahtzee box open, the optimal single-turn Yahtzee
    /// probability is 4.6029%, a classic figure.
    #[test]
    fn yahtzee_alone_matches_the_classic_probability() {
        let mut solver = Solver::new();
        let state = State::new(!Category::Yahtzee.bit(), 0, false);
        let value = solver.value(state);
        assert!((value - 50.0 * 0.046_029).abs() < 1e-3, "got {value}");
    }

    /// The lazy path and a fresh tier-by-tier solve agree bit for bit.
    #[test]
    fn lazy_and_tiered_solves_agree() {
        let scored = !(Category::Yahtzee
            .bit()
            .with(Category::Chance)
            .with(Category::Sixes));
        let state = State::new(scored, 40, false);
        let mut lazy = Solver::new();
        let lazy_value = lazy.value(state);
        let mut tiered = Solver::new();
        for filled in (scored.len()..=13).rev() {
            tiered.solve_tier(filled);
        }
        assert!(tiered.is_solved(state));
        let tiered_value = tiered.value(state);
        assert_eq!(lazy_value.to_bits(), tiered_value.to_bits());
    }

    /// Brute force straight from `State::apply`, with no shared tables:
    /// the strongest guard against index or probability bugs in the
    /// solver.  Dice fill one at a time (uniform over faces), so the
    /// probabilities emerge from counting rather than multinomials; the
    /// memo only tames the blowup and stores nothing the solver computes.
    fn brute_force_turn(
        solver: &mut Solver,
        memo: &mut std::collections::HashMap<(usize, u8), f64>,
        state: State,
        rolls_left: u8,
        counts: [u8; 6],
    ) -> f64 {
        let key = counts
            .iter()
            .rev()
            .fold(0, |acc, &c| acc * 6 + usize::from(c));
        if let Some(&value) = memo.get(&(key, rolls_left)) {
            return value;
        }
        let filled = usize::from(counts.iter().sum::<u8>());
        let value = if filled < 5 {
            // Average over the next die of the unordered remainder.
            (0..6)
                .map(|f| {
                    let mut counts = counts;
                    counts[f] += 1;
                    brute_force_turn(solver, memo, state, rolls_left, counts)
                })
                .sum::<f64>()
                / 6.0
        } else {
            let roll = Dice::from_counts(counts).expect("five dice");
            let score_now = state
                .legal_categories(roll)
                .iter()
                .map(|c| {
                    let delta = state.apply(c, roll).expect("legal");
                    f64::from(delta.reward()) + solver.value(delta.next)
                })
                .fold(f64::NEG_INFINITY, f64::max);
            if rolls_left == 0 {
                score_now
            } else {
                roll.keeps()
                    .map(|keep| {
                        brute_force_turn(solver, memo, state, rolls_left - 1, keep.counts())
                    })
                    .fold(score_now, f64::max)
            }
        };
        memo.insert((key, rolls_left), value);
        value
    }

    /// The widget agrees with the brute-force evaluator on joker-loaded
    /// endgames (Yahtzee scored 50, scored 0, and open).
    #[test]
    fn widget_matches_brute_force() {
        let lowers_done = Categories::LOWER.with(Category::Aces).with(Category::Twos);
        let cases = [
            State::new(!Category::Chance.bit(), 21, true),
            State::new(!Category::Chance.bit(), 21, false),
            State::new(lowers_done, 5, true),
            State::new(lowers_done, 5, false),
            State::new(!(Category::Yahtzee.bit().with(Category::Fours)), 60, false),
        ];
        for state in cases {
            let mut solver = Solver::new();
            let mut memo = std::collections::HashMap::new();
            let expected = brute_force_turn(&mut solver, &mut memo, state, 2, [0; 6]);
            let value = solver.value(state);
            assert!(
                (value - expected).abs() < 1e-9,
                "{state:?}: widget {value} vs brute force {expected}"
            );
        }
    }

    /// Query sanity on a late-game state: the best action is among the
    /// legal ones and EV accessors agree with it.
    #[test]
    fn queries_are_consistent() {
        let mut solver = Solver::new();
        let state = State::new(!(Category::Chance.bit().with(Category::Fives)), 30, false);
        let roll = dice("35556");
        let action = solver.best_action(state, roll, 2);
        match action {
            TurnAction::Score(category) => {
                assert!(state.legal_categories(roll).contains(category));
            }
            TurnAction::Reroll(keep) => {
                let best = solver.best_category(state, roll);
                let ev = solver.keep_ev(state, roll, keep, 2).expect("legal keep");
                let stand = solver.category_ev(state, roll, best).expect("legal");
                assert!(ev >= stand);
            }
        }
        assert_eq!(solver.keep_ev(state, roll, Keep::EMPTY, 0), None);
        assert_eq!(
            solver.keep_ev(state, roll, "11".parse().expect("keep"), 1),
            None
        );
        // A full card has nothing to reroll for.
        let full = State::new(Categories::ALL, 63, true);
        assert_eq!(solver.keep_ev(full, roll, Keep::EMPTY, 2), None);
    }

    /// Solving a tier before its successors must fail loudly, not
    /// poison the table with non-finite values.
    #[test]
    #[should_panic]
    fn out_of_order_tiers_are_rejected() {
        Solver::new().solve_tier(11);
    }
}