Tiny fix for CLI help.
[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):
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 ]