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