pokedex/formulae.py

   1 # encoding: utf8
   2 """Faithful translations of calculations the games make."""
   3 from __future__ import division
   4
   5 from itertools import izip
   6
   7 def nCr(n, r):
   8     """n-choose-r.
   9
  10     Thanks for the "compact" solution go to:
  11     http://stackoverflow.com/questions/2096573/counting-combinations-and-permutations-efficiently
  12     """
  13
  14     return reduce(
  15         lambda x, y: x * y[0] / y[1],
  16         izip(xrange(n - r + 1, n + 1),
  17              xrange(1, r + 1)),
  18         1)
  19
  20
  21 def calculated_stat(base_stat, level, iv, effort):
  22     """Returns the calculated stat -- i.e. the value actually shown in the game
  23     on a Pokémon's status tab.
  24     """
  25
  26     # Remember: this is from C; use floor division!
  27     return (base_stat * 2 + iv + effort // 4) * level // 100 + 5
  28
  29 def calculated_hp(base_hp, level, iv, effort):
  30     """Similar to `calculated_stat`, except with a slightly different formula
  31     used specifically for HP.
  32     """
  33
  34     # Shedinja's base stat of 1 is special; its HP is always 1
  35     if base_hp == 1:
  36         return 1
  37
  38     return (base_hp * 2 + iv + effort // 4) * level // 100 + 10 + level
  39
  40 def earned_exp(base_exp, level):
  41     """Returns the amount of EXP earned when defeating a Pokémon at the given
  42     level.
  43     """
  44
  45     return base_exp * level // 7
  46
  47 def capture_chance(percent_hp, capture_rate,
  48                    ball_bonus=10, status_bonus=1,
  49                    capture_bonus=10, capture_modifier=0):
  50     """Calculates the chance that a Pokémon will be caught, given its capture
  51     rate and the percentage of HP it has remaining.
  52
  53     Bonuses are such that 10 means "unchanged".
  54
  55     Returns five values: the chance of a capture, then the chance of the ball
  56     shaking three, two, one, or zero times.  Each of these is a float such that
  57     0.0 <= n <= 1.0.  Feel free to ignore all but the first.
  58     """
  59
  60     # HG/SS Pokéballs modify capture rate rather than the ball bonus
  61     capture_rate = capture_rate * capture_bonus // 10 + capture_modifier
  62     if capture_rate < 1:
  63         capture_rate = 1
  64     elif capture_rate > 255:
  65         capture_rate = 255
  66
  67     # A slight math note:
  68     # The actual formula uses (3 * max_hp - 2 * curr_hp) / (3 * max_hp)
  69     # This uses (1 - 2/3 * curr_hp/max_hp)
  70     # Integer division is taken into account by flooring immediately
  71     # afterwards, so there should be no appreciable rounding error.
  72     base_chance = int(
  73         capture_rate * ball_bonus // 10 * (1 - 2/3 * percent_hp)
  74     )
  75     base_chance = base_chance * status_bonus // 10
  76
  77     # Shake index involves integer sqrt.  Lovely.
  78     isqrt = lambda x: int(x ** 0.5)
  79     if not base_chance:
  80         # This is very silly.  Due to what must be an oversight, it's possible
  81         # for the above formula to end with a zero chance to catch, which is
  82         # then thrown blindly into the below denominator.  Luckily, the games'
  83         # division function is a no-op with a denominator of zero..  which
  84         # means a base_chance of 0 is effectively a base chance of 1.
  85         base_chance = 1
  86     shake_index = 1048560 // isqrt(isqrt(16711680 // base_chance))
  87
  88     # Iff base_chance < 255, then shake_index < 65535.
  89     # The Pokémon now has four chances to escape.  The game starts picking
  90     # random uint16s.  If such a random number is < shake_index, the Pokémon
  91     # stays in the ball, and it wobbles.  If the number is >= shake_index, the
  92     # ball breaks open then and there, and the capture fails.
  93     # If all four are < shake_index, the Pokémon is caught.
  94
  95     # If shake_index >= 65535, all four randoms must be < it, and the Pokémon
  96     # will be caught.  Skip hard math
  97     if shake_index >= 65535:
  98         return (1.0, 0.0, 0.0, 0.0, 0.0)
  99
 100     # This brings up an interesting invariant: sum(return_value) == 1.0.
 101     # Something is guaranteed to happen.
 102
 103     # Alrighty.  Here's some probability.
 104     # The chance that a single random uint16 will be < shake_index, thus
 105     # keeping the Pokémon in the ball, is:
 106     p = shake_index / 65536
 107
 108     # Now, the chance for n wobbles is the chance that the Pokémon will stay in
 109     # the ball for (n-1) attempts, then break out on the nth.
 110     # The chance of capture is just the chance that the Pokémon stays in the
 111     # ball for all four tries.
 112
 113     # There are five cases: captured, wobbled three times, etc.
 114     return [
 115         p**4,  # capture
 116         p**3 * (1 - p),
 117         p**2 * (1 - p),
 118         p**1 * (1 - p),
 119                (1 - p),
 120     ]