Add stew command

src/commands/calc/drop.rs

@@ -96,10 +96,7 @@ fn parse_target_range(s: &str) -> anyhow::Result<TargetRange> {
 #[derive(Copy, Clone, Debug)]
 enum TargetRange {
     Eq(usize),
-    Neq(usize),
-    Lt(usize),
     Lte(usize),
-    Gt(usize),
     Gte(usize),
 }
 
@@ -109,38 +106,17 @@ impl TargetRange {
     fn to_values(&self, iterations: usize) -> Box<dyn Iterator<Item = usize>> {
         match self {
             Self::Eq(k) => Box::new(*k..=*k),
-            Self::Neq(k) => Box::new((0..*k).chain(k + 1..=iterations)),
-            Self::Lt(k) => Box::new(0..=*k - 1),
             Self::Lte(k) => Box::new(0..=*k),
-            Self::Gt(k) => Box::new(*k + 1..=iterations),
             Self::Gte(k) => Box::new(*k..=iterations),
         }
     }
-
-    /// Invert this range, to create a range starting at the same point but
-    /// going in the opposite direction (and inclusion/exclusion of the point
-    /// is inverted). A range and its inversion together cover the entire
-    /// number scale without any overlap.
-    fn invert(&self) -> Self {
-        match self {
-            Self::Eq(k) => Self::Neq(*k),
-            Self::Neq(k) => Self::Eq(*k),
-            Self::Lt(k) => Self::Gte(*k),
-            Self::Lte(k) => Self::Gt(*k),
-            Self::Gt(k) => Self::Lte(*k),
-            Self::Gte(k) => Self::Lt(*k),
-        }
-    }
 }
 
 impl Display for TargetRange {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Eq(k) => write!(f, "{}", k),
-            Self::Neq(k) => write!(f, "≠{}", k),
-            Self::Lt(k) => write!(f, "<{}", k),
             Self::Lte(k) => write!(f, "≤{}", k),
-            Self::Gt(k) => write!(f, ">{}", k),
             Self::Gte(k) => write!(f, "≥{}", k),
         }
     }
@@ -180,28 +156,11 @@ impl Command for CalcDropCommand {
         // up the probability of all the values in the
         // range. https://en.wikipedia.org/wiki/Binomial_distribution#Cumulative_distribution_function
 
-        // There's an optimization here to minimize the number of
-        // binomial functions we have to run (minimize the domain of the
-        // cdf we compute). If k is less than half of n, then compute
-        // the odds of hitting the opposite result (e.g. <k instead of
-        // >=k), then subtract from 1.
-        let result_prob: f64 = if self.target.to_values(self.iterations).count()
-            <= self.iterations / 2
-        {
-            // normal case
-            math::binomial_cdf(
-                self.probability,
-                self.iterations,
-                &mut self.target.to_values(self.iterations),
-            )
-        } else {
-            // inverted case
-            1.0 - math::binomial_cdf(
-                self.probability,
-                self.iterations,
-                &mut self.target.invert().to_values(self.iterations),
-            )
-        };
+        let result_prob: f64 = math::binomial_cdf(
+            self.probability,
+            self.iterations,
+            &mut self.target.to_values(self.iterations),
+        );
 
         println!(
             "{:.4}% chance of {} successes in {} attempts",

src/commands/calc/mod.rs

@@ -2,11 +2,14 @@
 //! calculations.
 
 mod drop;
+mod stew;
 mod xp;
 
 use crate::{
     commands::{
-        calc::{drop::CalcDropCommand, xp::CalcXpCommand},
+        calc::{
+            drop::CalcDropCommand, stew::CalcStewCommand, xp::CalcXpCommand,
+        },
         Command, CommandType,
     },
     utils::context::CommandContext,
@@ -16,13 +19,15 @@ use structopt::StructOpt;
 #[derive(Debug, StructOpt)]
 pub enum CalcCommandType {
     Drop(CalcDropCommand),
+    Stew(CalcStewCommand),
     Xp(CalcXpCommand),
 }
 
 impl CommandType for CalcCommandType {
     fn command(&self) -> &dyn Command {
         match &self {
             Self::Drop(cmd) => cmd,
+            Self::Stew(cmd) => cmd,
             Self::Xp(cmd) => cmd,
         }
     }

src/commands/calc/stew.rs

@@ -0,0 +1,109 @@
+use std::iter;
+
+use crate::{
+    commands::Command,
+    error::OsrsError,
+    utils::{context::CommandContext, math},
+};
+use prettytable::{color, format::Alignment, Attr, Cell, Row, Table};
+use structopt::StructOpt;
+
+/// Maximum number of doses per stew
+const MAX_DOSES: usize = 3;
+
+/// Highest possible boost
+const MAX_BOOST: usize = 5;
+
+/// TODO
+const CUMULATIVE_PROBS: [[f64; MAX_BOOST + 1]; MAX_DOSES] = [
+    [0.750, 0.250, 0.000, 0.000, 0.000, 0.000], // 1 dose
+    [0.625, 0.375, 0.250, 0.125, 0.000, 0.000], // 2 doses
+    [0.583, 0.417, 0.333, 0.250, 0.167, 0.083], // 3 doses
+];
+
+fn prob_for_stews(
+    boost: usize,
+    total_doses: usize,
+    doses_per_stew: usize,
+) -> f64 {
+    let total_stews = total_doses / doses_per_stew; // rounded down
+    let prob_per_stew = CUMULATIVE_PROBS[doses_per_stew - 1][boost];
+    math::binomial_cdf(prob_per_stew, total_stews, &mut (1..=total_stews))
+}
+
+/// Calculate probabilities related to spicy stew level boosts.
+#[derive(Debug, StructOpt)]
+pub struct CalcStewCommand {
+    /// The MINIMUM number of levels to boost. Boosting more levels that this
+    /// value WILL be included in the output probabilities.
+    #[structopt(short, long)]
+    boost: usize,
+    /// The total number of doses you have of the relevant spice.
+    #[structopt(short = "d", long = "doses")]
+    total_doses: usize,
+}
+
+impl Command for CalcStewCommand {
+    fn execute(&self, context: &CommandContext) -> anyhow::Result<()> {
+        if self.boost > MAX_BOOST {
+            return Err(OsrsError::ArgsError(format!(
+                "Maximum boost is {} levels",
+                MAX_BOOST
+            ))
+            .into());
+        }
+
+        let mut table = Table::new();
+        table.set_format(
+            *prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE,
+        );
+
+        table.set_titles(Row::new(
+            iter::once(Cell::new_align(&"Doses/Stew", Alignment::RIGHT))
+                // Add one col for each boost number (1-5)
+                .chain((1..=MAX_BOOST).map(|boost| {
+                    let mut cell = Cell::new_align(
+                        &format!("≥+{}", boost),
+                        Alignment::RIGHT,
+                    );
+                    if boost == self.boost {
+                        cell.style(Attr::Bold);
+                    }
+                    cell
+                }))
+                .collect(),
+        ));
+
+        for doses_per_stew in 1..=MAX_DOSES {
+            table.add_row(Row::new(
+                iter::once(Cell::new_align(
+                    &context.fmt_num(&doses_per_stew),
+                    Alignment::RIGHT,
+                ))
+                // Calculate prob for hitting each boost value (1-5)
+                .chain((1..=MAX_BOOST).map(|boost| {
+                    let prob =
+                        prob_for_stews(boost, self.total_doses, doses_per_stew);
+                    let mut cell = Cell::new_align(
+                        &format!("{:.1}%", prob * 100.0),
+                        Alignment::RIGHT,
+                    );
+
+                    // Add some extra conditional styling
+                    if boost == self.boost {
+                        cell.style(Attr::Bold);
+                        if prob > 0.0 {
+                            cell.style(Attr::ForegroundColor(color::GREEN));
+                        }
+                    }
+
+                    cell
+                }))
+                .collect(),
+            ));
+        }
+        table.printstd();
+
+        Ok(())
+    }
+}

src/utils/math.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 /// Calculate numerical combination `nCk` (n choose k).
 pub fn combination(n: usize, k: usize) -> usize {
     // safety check to prevent overflow/underflow
@@ -43,7 +45,23 @@ pub fn binomial_cdf(
     n: usize,
     k_values: &mut dyn Iterator<Item = usize>,
 ) -> f64 {
-    k_values.map(|k_i| binomial(p, n, k_i)).sum()
+    let helper = |iter: &mut dyn Iterator<Item = usize>| {
+        iter.map(|k_i| binomial(p, n, k_i)).sum()
+    };
+
+    // There's an optimization here to minimize the number of
+    // binomial functions we have to run (minimize the domain of the
+    // cdf we compute). If k is less than half of n, then compute
+    // the odds of hitting the opposite result (e.g. <k instead of
+    // >=k), then subtract from 1.
+    let k_owned: HashSet<usize> = k_values.collect();
+    if k_owned.len() <= n / 2 {
+        // normal case
+        helper(&mut k_owned.into_iter())
+    } else {
+        // inverted case
+        1.0 - helper(&mut (0..=n).filter(|k_i| !k_owned.contains(k_i)))
+    }
 }
 
 #[cfg(test)]
@@ -96,5 +114,6 @@ mod tests {
         assert_approx_eq!(binomial_cdf(0.5, 2, &mut (0..=0)), 0.25);
         assert_approx_eq!(binomial_cdf(0.5, 2, &mut (0..=1)), 0.75);
         assert_approx_eq!(binomial_cdf(0.5, 2, &mut (0..=2)), 1.0);
+        assert_approx_eq!(binomial_cdf(0.5, 10, &mut (0..=2)), 0.0546875);
     }
 }