Added capture chance formula.
[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(current_hp, max_hp, capture_rate,
48 ball_bonus=1, status_bonus=1, heavy_modifier=0):
49 """Calculates the chance that a Pokémon will be caught.
50
51 Returns five values: the chance of a capture, then the chance of the ball
52 shaking three, two, one, or zero times. Each of these is a float such that
53 0.0 <= n <= 1.0. Feel free to ignore all but the first.
54 """
55
56 if heavy_modifier:
57 # Only used by Heavy Ball. Changes the target's capture rate outright
58 capture_rate += heavy_modifier
59 if capture_rate <= 1:
60 capture_rate = 1
61
62 # This should really be integer math, right? But the formula uses FOURTH
63 # ROOTS in a moment, so it can't possibly be. It probably doesn't matter
64 # either way, so whatever; use regular ol' division. ball_bonus and
65 # status_bonus can be 1.5, anyway.
66 base_chance = ((3 * max_hp - 2 * current_hp) * capture_rate * ball_bonus) \
67 / (3 * max_hp) \
68 * status_bonus
69
70 shake_index = (base_chance / 255) ** 0.25 * (2**16 - 1)
71
72 # Iff base_chance < 255, then shake_index < 65535.
73 # The game now picks four random uwords. However many of them are <=
74 # shake_index is the number of times the ball will shake. If all four are
75 # <= shake_index, the Pokémon is caught.
76
77 # The RNG tends to work with integers, so integer math likely kicks in now.
78 shake_index = int(shake_index)
79
80 # If shake_index >= 65535, all four randoms must be <= it, and the Pokémon
81 # will be caught. Skip hard math
82 if shake_index >= 65535:
83 return (1.0, 0.0, 0.0, 0.0, 0.0)
84
85 # This brings up an interesting invariant: sum(return_value) == 1.0.
86 # Something is guaranteed to happen.
87
88 # Alrighty. Here's some probability.
89 # The chance that a single random number will be <= shake_index is:
90 p = (shake_index + 1) / 65536
91 # Now, the chance that two random numbers will be <= shake_index is p**2.
92 # And the chance that neither will be is (1 - p)**2.
93 # With me so far?
94 # The chance that one will be and one will NOT be is p * (1 - p) * 2.
95 # The 2 is because they can go in any order: the first could be less, or
96 # the second could be less. That 2 is actually nCr(2, 1); the number of
97 # ways of picking one item in any order from a group of two.
98 # Try it yourself add up those three values and you'll get 1.
99
100 # Right. Hopefully, the following now makes sense.
101 # There are five cases: four randoms are <= shake_index (which means
102 # capture), or three are, etc.
103 return [
104 p**i * (1 - p)**(4 - i) * nCr(4, i)
105 for i in reversed(range(5))
106 ]