Calculate the probability of stabilization with treatment

probability_swn_stabilize.rake

desc "Calculate the probability of stabilization with treatment"
task :probability_swn_stabilize do
  # This class encapsulates the probability calculations based on
  # the given :distribution.
  class Universe
    def initialize(label:, distribution:)
      @label = label
      @distribution = distribution
      @max = distribution.keys.max
      @size = distribution.values.sum.to_f
    end
    attr_reader :label

    # Given the :modified_target, what is the chance of a result
    # equal to or greater than that target?
    #
    # @param modified_target [Integer] What is the roll we're looking
    # for?
    #
    # @return [Float] the chance, with a range between 0 and 1.
    def chance_of_gt_or_eq_to(modified_target)
      @distribution.slice(*(modified_target..@max)).values.sum /
        @size
    end

    # Given the :modified_target, what is the chance of a result
    # equal to exactly the modified target?
    #
    # @param modified_target [Integer] What is the roll we're looking
    # for?
    #
    # @return [Float] the chance, with a range between 0 and 1.
    def chance_of_exactly(modified_target)
      @distribution.fetch(modified_target, 0) / @size
    end
  end

  # Think to your Settlers of Catan board game.  The "dot" on each
  # of the number chits is the chance in 36 of rolling that number
  # on 2d6.
  universe_2d6 = Universe.new(
    label: "2d6",
    distribution: {
      2 => 1, 3 => 2, 4 => 3,
      5 => 4, 6 => 5, 7 => 6,
      8 => 5, 9 => 4, 10 => 3,
      11 => 2, 12 => 1
    })

  # This universe is the distribution of all possible rolls in which
  # we throw 3 six-sided dice and keep the best 2 values.
  universe_3d6 = Universe.new(
    label: "3d6",
    distribution: {
      2  => 1, 3  => 3, 4  => 7,
      5  => 12, 6  => 19, 7  => 27,
      8  => 34, 9  => 36, 10 => 34,
      11 => 27, 12 => 16
    })

  # This class encapsulates the probability
  #
  # A key assumption in helping is that all helpers are helping with
  # their preferred skill check, and thus the modified_difficulty is
  # the same as the person performing the primary test.
  #
  # @note I tested this using a coin toss universe distribution
  # (e.g. heads or tails)
  class Helper
    # @param number [Integer] the number of helpers
    #
    # @param universe [Universe] the universe of possible dice
    # results for each of the helpers
    def initialize(number:, universe:)
      @number = number
      @universe = universe
    end
    attr_reader :number

    # @param modified_difficulty [Integer] the modified dice roll that
    # the helpers are trying to achieve.
    #
    # @return [Float] the probability that one of the helpers
    # succeeds.
    def chance_someone_succeeds_at(modified_difficulty)
      success = @universe.chance_of_gt_or_eq_to(modified_difficulty)
      accumulate(success: success)
    end

    private

    def accumulate(success:, accumulator: 0, number: @number)
      return accumulator if number <= 0

      chance_of_next_helper = 1 - accumulator
      accumulator += chance_of_next_helper * success
      accumulate(
        success: success,
        number: number - 1,
        accumulator: accumulator)
    end
  end

  # Originally extracted as a parameter object (because the
  # `success_probability_for_given_check_or_later` method had too
  # many parameters), this object helps encapsulate the data used
  # to calculate each round's probability of success.
  class Check
    PARAMETERS = [
      :universe,
      :modified_difficulty,
      :reroll,
      :chance_we_need_this_round,
      :round,
      :helpers
    ]

    def initialize(**kwargs)
      PARAMETERS.each do |param|
        instance_variable_set("@#{param}", kwargs.fetch(param))
      end
    end

    # @return [Universe] The universe of possible modified dice rolls.
    attr_reader :universe

    # @return [Integer] The check's Difficulty Class (DC) plus all of
    # the modifiers (excluding those for rounds since dropping to 0
    # HP) affecting the 2d6 roll.
    #
    # @note for a Heal-2 medic with a Dex of +1 using a Lazurus
    # Patch (DC 6) would have a modified_difficulty of
    # 3. (e.g. 6 - 2 - 1 = 3)
    #
    # @note for an untrained medic with a Dex of -1 using a Lazurus
    # Patch (DC 6) would have a modified_difficulty of
    # 3. (e.g. 6 - (-1) - (-1) = 9)
    attr_reader :modified_difficulty

    # @return [Boolean] True if we allow a reroll.
    attr_reader :reroll

    # @return [Float] The probability that we need this round, range
    # between 0.0 and 1.0.
    attr_reader :chance_we_need_this_round

    # @return [Integer] The round in which we made a check
    attr_reader :round

    # @return [Helper] Who are the helpers for this task
    attr_reader :helpers

    def next(**kwargs)
      new_kwargs = {}
      PARAMETERS.each do |param|
        new_kwargs[param] = kwargs.fetch(param) do
          instance_variable_get("@#{param}")
        end
      end
      self.class.new(**new_kwargs)
    end

    # @return [Float]
    def probability_of_success_this_round
      universe.chance_of_gt_or_eq_to(modified_difficulty + round)
    end

    # @return [Float]
    def chance_of_help_making_the_difference
      universe.chance_of_exactly(modified_difficulty + round - 1) *
        helpers.chance_someone_succeeds_at(modified_difficulty + round)
    end

    # @return [Float]
    def probability_of_reroll_makes_difference(prob:)
      return 0.0 unless reroll
      (1 - prob) * prob
    end
  end


  # @param check [Check] The current round's Check context.
  #
  # @param ttl [Integer] "Time to Live" - After this many rounds
  # without treatment, a character dies.
  #
  # @return [Float] The probability of success in this round or all
  # future rounds. Range between 0.0 and 1.0
  #
  # @note Per SWN rules, characters die after 6 rounds without
  # treatment.
  def success_probability_for_given_check_or_later(check, ttl: 6)
    return 0 if check.round >= ttl

    prob = check.probability_of_success_this_round

    # Using the assumption that the best character is making the
    # check, and everyone that could be helping is using their best
    # check to help.  With the assumption that everyone has the same
    # modifier as the base check.
    prob += check.chance_of_help_making_the_difference

    # I believe the interaction of helpers and reroll is correct.
    # We only attempt a re-roll if the helpers didn't succeed.  And
    # the re-roll has the same probability of success
    prob += check.probability_of_reroll_makes_difference(prob: prob)

    next_check = check.next(
      round: check.round + 1,
      reroll: false,
      chance_we_need_this_round: (1 - prob)
    )

    check.chance_we_need_this_round * (
      prob +
      success_probability_for_given_check_or_later(next_check)
    )
  end

  all_helpers = [
    Helper.new(number: 0, universe: universe_2d6),
    Helper.new(number: 1, universe: universe_2d6),
    Helper.new(number: 2, universe: universe_2d6)
  ]

  number_of_rounds = (0..5)
  header_template = "| %-10s | %-4s | %7s | %6s" +
                    " | %5s" * number_of_rounds.size + " |"
  divider_template = "|------------+------+---------+--------" +
                     ("+-------" * number_of_rounds.size) + "|"
  line_template = "| %-10s | %-4s | %-7d | %6s" +
                  " | %.3f" * number_of_rounds.size + " |"
  header = sprintf(header_template, "Mod Diff", "Dice", "Helpers",
                   "Reroll",
                   *number_of_rounds.to_a.map {|i| "Rnd #{i}"})
  puts header
  puts divider_template

  rows = []
  (0..12).each do |modified_difficulty|
    [universe_2d6, universe_3d6].each do |universe|
      all_helpers.each do |helpers|
        [false, true].each do |reroll|
          rounds = number_of_rounds.map do |round|
            check = Check.new(
              reroll: reroll,
              universe: universe,
              modified_difficulty: modified_difficulty,
              chance_we_need_this_round: 1.0,
              helpers: helpers,
              round: round
            )
            success_probability_for_given_check_or_later(check)
          end
          rows << {
            "Modified Difficulty" => modified_difficulty,
            "Dice Rolled" => universe.label,
            "Helpers" => helpers.number,
            "Reroll" => reroll,
            "Round 0" => sprintf("%.3f", rounds[0]),
            "Round 1" => sprintf("%.3f", rounds[1]),
            "Round 2" => sprintf("%.3f", rounds[2]),
            "Round 3" => sprintf("%.3f", rounds[3]),
            "Round 4" => sprintf("%.3f", rounds[4]),
            "Round 5" => sprintf("%.3f", rounds[5])
          }
          puts sprintf(line_template,
                       modified_difficulty,
                       universe.label,
                       helpers.number,
                       reroll,
                       *rounds)
        end
      end
    end
  end

  # Below is used for writing a YAML file that I use to auto-generate
  # the table for a blog post.
  pwd = if defined?(TakeOnRules::PROJECT_PATH)
          File.join(TakeOnRules::PROJECT_PATH, "data/probabilities")
        else
          Dir.pwd
        end

  column_names = [
    "Modified Difficulty",
    "Dice Rolled",
    "Helpers",
    "Reroll",
    "Round 0",
    "Round 1",
    "Round 2",
    "Round 3",
    "Round 4",
    "Round 5"
  ]

  columns = []
  column_names.each do |name|
    columns << { "label" => name, "key" => name }
  end

  require 'psych'
  File.open(File.join(pwd, "swn_stabilize.yml"), "w+") do |file|
    file.puts(
      Psych.dump(
        "name" => 'Probability of Stabilization in SWN',
        "columns" => columns,
        "rows" => rows
      )
    )
  end
end