Make lookup index respect the moved names
[zzz-pokedex.git] / 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 ]