dbff83b99f5703251e6e2365fbf518da3de6437f
[pyrisk.git] / risk / base.py
1 #!/usr/bin/python
2 #
3 # A Python engine for Risk-like games
4
5 import random
6
7
8 VERSION='0.1'
9
10
11 class PlayerError (Exception):
12     pass
13
14 class ID_CmpMixin (object):
15     def __str__(self):
16         return self.name
17     def __repr__(self):
18         return self.__str__()
19     def __cmp__(self, other):
20         return cmp(id(self), id(other))
21     def __eq__(self, other):
22         return self.__cmp__(other) == 0
23     def __ne__(self, other):
24         return self.__cmp__(other) != 0
25
26 class Territory (list, ID_CmpMixin):
27     def __init__(self, name, short_name=None, type=-1,
28                  link_names=[], continent=None, player=None):
29         list.__init__(self)
30         ID_CmpMixin.__init__(self)
31         self.name = name
32         self.short_name = short_name
33         if short_name == None:
34             self.short_name = name
35         self._card_type = type
36         self._link_names = list(link_names)
37         self.continent = continent
38         self.player = player
39         self.card = None
40         self.armies = 0
41     def __str__(self):
42         if self.short_name == self.name:
43             return self.name
44         return '%s (%s)' % (self.name, self.short_name)
45     def __repr__(self):
46         return self.__str__()
47     def borders(self, other):
48         for t in self:
49             if id(t) == id(other):
50                 return True
51         return False
52
53 class Continent (list, ID_CmpMixin):
54     def __init__(self, name, production, territories=[]):
55         list.__init__(self, territories)
56         ID_CmpMixin.__init__(self)
57         self.name = name
58         self.production = production
59     def append(self, territory):
60         list.append(self, territory)
61         territory.continent = self
62     def territory_by_name(self, name):
63         for t in self:
64             if name.lower() in [t.short_name.lower(), t.name.lower()]:
65                 #assert self.contains_territory(t), t
66                 return t
67         raise KeyError(name)
68     def contains_territory(self, territory):
69         for t in self:
70             if t == territory:
71                 return True
72         return False
73     def single_player(self):
74         p = self[0].player
75         for territory in self:
76             if territory.player != p:
77                 return False
78         return True
79
80 class World (list, ID_CmpMixin):
81     def __init__(self, name, continents=[]):
82         list.__init__(self, continents)
83         ID_CmpMixin.__init__(self)
84         self.name = name
85         self.initial_armies = { # num_players:num_armies
86             2: 40, 3:35, 4:30, 5:25, 6:20
87                 }
88     def territories(self):
89         for continent in self:
90             for territory in continent:
91                 yield territory
92     def territory_by_name(self, name):
93         for continent in self:
94             try:
95                 return continent.territory_by_name(name)
96             except KeyError:
97                 pass
98         raise KeyError(name)
99     def contains_territory(self, territory):
100         for continent in self:
101             if continent.contains_territory(territory):
102                 return True
103         return False
104     def continent_by_name(self, name):
105         for continent in self:
106             if continent.name.lower() == name.lower():
107                 return continent
108         raise KeyError(name)
109     def _resolve_link_names(self):
110         self._check_short_names()
111         for territory in self.territories():
112             for name in territory._link_names:
113                 other = self.territory_by_name(name)
114                 if not territory.borders(other):
115                     territory.append(other)
116                 if not other.borders(territory):
117                     other.append(territory)
118     def _check_short_names(self):
119         ts = {}
120         for t in self.territories():
121             if t.short_name not in ts:
122                 ts[t.short_name] = t
123             else:
124                 raise ValueError('%s shared by %s and %s'
125                                  % (t.short_name, ts[t.short_name], t))
126     def production(self, player):
127         ts = list(player.territories(world))
128         production = max(3, len(ts) / 3)
129         continents = set([t.continent.name for t in ts])
130         for c_name in continents:
131             c = self.continent_by_name(c_name)
132             if c.single_player() == True:
133                 production += c.production
134         return (production, {})
135     def place_territory_production(self, territory_production):
136         for territory_name,production in territory_production.items():
137             t = self.territory_by_name(territory_name)
138             t.armies += production
139
140 class Card (ID_CmpMixin):
141     def __init__(self, deck, type_, territory=None):
142         ID_CmpMixin.__init__(self)
143         self.deck = deck
144         self.territory = territory
145         self.type = type_
146     def __str__(self):
147         if self.territory == None:
148             return '<Card %s>' % (self.deck.type_names[self.type])
149
150         return '<Card %s %s>' % (self.territory,
151                                  self.deck.type_names[self.type])
152     def __repr__(self):
153         return self.__str__()
154
155 class Deck (list):
156     def __init__(self, territories=[]):
157         list.__init__(self, [Card(self, t._card_type, t) for t in territories])
158         random.shuffle(self)
159         self.type_names = ['Wild', 'Infantry', 'Cavalry', 'Artillery', 'Wild']
160         for i in range(2):
161             self.append(Card(self, 0))
162         self._production_sequence = [4, 6, 8, 10, 12, 15]
163         self._production_index = 0
164     def production_value(self, index):
165         """
166         >>> d = Deck()
167         >>> [d.production_value(i) for i in range(8)]
168         [4, 6, 8, 10, 12, 15, 20, 25]
169         """
170         if index < len(self._production_sequence):
171             return self._production_sequence[index]
172         extra = index - len(self._production_sequence) + 1
173         return self._production_sequence[-1] + 5 * extra
174     def production(self, player, cards=None):
175         """
176         >>> d = Deck()
177         >>> a = Player('Alice')
178         >>> b = Player('Bob')
179         >>> d.production(a, None)
180         (0, {})
181         >>> d.production(a, [Card(d, 1, Territory('a')),
182         ...                  Card(d, 1, Territory('b'))])
183         Traceback (most recent call last):
184           ...
185         PlayerError: You must play cards in groups of 3, not 2
186         ([<Card a Infantry>, <Card b Infantry>])
187         >>> d.production(a, [Card(d, 1, Territory('a', player=a)),
188         ...                  Card(d, 1, Territory('b', player=b)),
189         ...                  Card(d, 1, Territory('c'))])
190         (4, {'a': 1})
191         >>> p,tp = d.production(a, [Card(d, 1, Territory('a', player=a)),
192         ...                     Card(d, 2, Territory('b', player=a)),
193         ...                     Card(d, 0, Territory('c', player=a))])
194         >>> p
195         6
196         >>> sorted(tp.items())
197         [('a', 1), ('b', 1), ('c', 1)]
198         """
199         if cards == None:
200             return (0, {})
201         if len(cards) != 3:
202             raise PlayerError('You must play cards in groups of 3, not %d\n(%s)'
203                               % (len(cards), cards))
204         h = Hand(cards)
205         if h.set() or h.run():
206             p = self.production_value(self._production_index)
207             self._production_index += 1
208             territory_production = {}
209             for c in cards:
210                 if c.territory != None and c.territory.player == player:
211                     territory_production[c.territory.name] = 1
212             return (p, territory_production)
213         raise PlayerError('%s is neither a set nor a run' % cards)
214
215 class Hand (list):
216     def __init__(self, cards=[]):
217         list.__init__(self, cards)
218     def set(self):
219         s = sorted(set([card.type for card in self]))
220         if len(s) == 1 \
221                 or (len(s) == 2 and s[0] == 0):
222             return True
223         return False
224     def run(self):
225         if len(set([card.type for card in self])) == 3:
226             return True
227         return False
228     def possible(self):
229         if len(self) >= 3:
230             for i,c1 in enumerate(self[:-2]):
231                 for j,c2 in enumerate(self[i+1:-1]):
232                     for c3 in self[i+j+2:]:
233                         h = Hand([c1, c2, c3])
234                         if h.set() or h.run():
235                             yield h
236
237 class Player (ID_CmpMixin):
238     def __init__(self, name):
239         self.name = name
240         ID_CmpMixin.__init__(self)
241         self.alive = True
242         self.hand = Hand()
243     def territories(self, world):
244         for t in world.territories():
245             if t.player == self:
246                 yield t
247     def border_territories(self, world):
248         for t in self.territories(world):
249             for neighbor in t:
250                 if neighbor.player != self:
251                     yield t
252                     break
253     def phase_select_territory(self, world):
254         """Return the selected territory
255         """
256         free_territories = [t for t in world.territories() if t.player == None]
257         return random.sample(free_territories, 1)[0]
258     def phase_play_cards(self, world, play_required=True):
259         if play_required == True:
260             return random.sample(list(self.hand.possible()), 1)[0]
261     def phase_place_armies(self, world, remaining=1, this_round=1):
262         """Both during setup and before each turn.
263
264         Return {territory_name: num_armies, ...}
265         """
266         t = random.sample(list(self.border_territories(world)), 1)[0]
267         return {t.name: this_round}
268     def phase_attack(self, world):
269         """Return list of (source, target, armies) tuples.  Place None
270         in the list to end this phase.
271         """
272         possible_attacks = []
273         for t in self.border_territories(world):
274             if t.armies <= 3: #1: # be more conservative, only attack with 3 dice
275                 continue
276             targets = [border_t for border_t in t if border_t.player != self]
277             for tg in targets:
278                 possible_attacks.append((t.name, tg.name, min(3, t.armies-1)))
279         if len(possible_attacks) == 0:
280             return [None]
281         return random.sample(possible_attacks, 1) # + [None]
282     def phase_support_attack(self, world, source, target):
283         return source.armies-1
284     def phase_fortify(self, world):
285         """Return list of (source, target, armies) tuples.  Place None
286         in the list to end this phase.
287         """
288         return [None]
289     def phase_draw(self, cards=[]):
290         """Only called if you earned a new card (or cards)"""
291         self.hand.extend(cards)
292
293 class Engine (ID_CmpMixin):
294     def __init__(self, world, players, deck_class=Deck):
295         ID_CmpMixin.__init__(self)
296         self.world = world
297         self.deck = deck_class(world.territories())
298         self.players = players
299     def __str__(self):
300         return '<engine %s %s>' % (self.world, self.players)
301     def __repr__(self):
302         return self.__str__()
303     def run(self):
304         self.setup()
305         self.play()
306     def setup(self):
307         for p in self.players:
308             p.alive = True
309         random.shuffle(players)
310         self.select_territories()
311         self.place_initial_armies()
312         self.deal()
313     def play(self):
314         turn = 0
315         active_player = 0
316         living = len(self.living_players())
317         while living > 1:
318             self.play_turn(self.players[active_player])
319             living = len(self.living_players())
320             active_player = (active_player + 1) % living
321             turn += 1
322     def play_turn(self, player):
323         print "%s's turn (territory score: %s)" \
324             % (player, [(p,len(list(p.territories(self.world))))
325                         for p in self.players])
326         self.play_cards_and_place_armies(player)
327         captures = self.attack_phase(player)
328         if captures > 0 and len(self.deck) > 0 and len(self.living_players()) > 1:
329             player.phase_draw([self.deck.pop()])
330     def select_territories(self):
331         for t in self.world.territories():
332             t.player = None
333         for i in range(len(list(self.world.territories()))):
334             p = self.players[i % len(self.players)]
335             t = p.phase_select_territory(self.world)
336             if t.player != None:
337                 raise PlayerError('Cannot select %s owned by %s'
338                                   % (t, t.player))
339             print '%s selects %s' % (p, t)
340             t.player = p
341             t.armies = 1
342     def place_initial_armies(self):
343         already_placed = [len(list(p.territories(self.world))) for p in self.players]
344         s = list(set(already_placed))
345         assert len(s) in [1,2], already_placed
346         if len(s) == 2: # catch up the players who are one territory short
347             assert min(s) == max(s)-1, 'Min %d, max %d' % (min(s), max(s))
348             for p,placed in zip(self.players, already_placed):
349                 if placed == min(s):
350                     self.player_place_armies(p, remaining, 1)
351         remaining = self.world.initial_armies[len(self.players)] - max(s)
352         while remaining > 0:
353             for p in self.players:
354                 self.player_place_armies(p, remaining, 1)
355             remaining -= 1
356     def player_place_armies(self, player, remaining=1, this_round=1):
357         placements = player.phase_place_armies(self.world, remaining, this_round)
358         if sum(placements.values()) != this_round:
359             raise PlayerError('Placing more than %d armies' % this_round)
360         for ter_name,armies in placements.items():
361             t = world.territory_by_name(ter_name)
362             if t.player != player:
363                 raise PlayerError('Placing armies in %s owned by %s'
364                                   % (t, t.player))
365             if armies < 0:
366                 raise PlayerError('Placing a negative number of armies (%d) in %s'
367                                   % (armies, t))
368         print '%s places %s' % (player, placements)
369         for ter_name,armies in placements.items():
370             t = world.territory_by_name(ter_name)
371             t.armies += armies
372     def deal(self):
373         for p in self.players:
374             cards = []
375             for i in range(3):
376                 cards.append(self.deck.pop())
377             p.phase_draw(cards)
378         print 'Initial hands dealt'
379     def play_cards_and_place_armies(self, player, additional_armies=0):
380         cards_required = len(player.hand) >= 5
381         cards = player.phase_play_cards(
382             self.world, play_required=cards_required)
383         if cards_required == True and cards == None:
384             raise PlayerError('You have %d >= 5 cards in your hand, you must play'
385                               % len(player.hand))
386         w_prod,w_terr_prod = self.world.production(player)
387         print '%s earned %d armies from territories' % (player, w_prod)
388         c_prod,c_terr_prod = self.deck.production(player, cards)
389         if c_prod > 0:
390             print '%s played %s, earning %d armies' \
391                 % (player, cards, c_prod+sum(c_terr_prod.values()))
392         if cards != None:
393             for c in cards:
394                 player.hand.remove(c)
395         for terr,prod in c_terr_prod.items():
396             if terr in w_terr_prod:
397                 w_terr_prod[terr] += prod
398             else:
399                 w_terr_prod[terr] = prod
400         self.world.place_territory_production(w_terr_prod)
401         if len(w_terr_prod) > 0:
402             print '%s was required to place %s' % (player, w_terr_prod)
403         armies = w_prod + c_prod
404         self.player_place_armies(player, armies, armies)
405     def attack_phase(self, player):
406         captures = 0
407         while True:
408             attacks = player.phase_attack(self.world)
409             for attack in attacks:
410                 if attack == None:
411                     return captures
412                 source_name,target_name,armies = attack
413                 source = self.world.territory_by_name(source_name)
414                 target = self.world.territory_by_name(target_name)
415                 tplayer = target.player
416                 capture = self.attack(source, target, armies)
417                 if capture == True:
418                     captures += 1
419                     if len(list(tplayer.territories(self.world))) == 0:
420                         self.player_killed(tplayer, killer=player)
421     def attack(self, source, target, armies):
422         if source.player == target.player:
423             raise PlayerError('%s attacking %s, but you own both.'
424                               % (source, target))
425         if armies == 0:
426             raise PlayerError('%s attacking %s with 0 armies.'
427                               % (source, target))
428         if armies >= source.armies:
429             raise PlayerError('%s attacking %s with %d armies, but only %d are available.'
430                               % (source, target, armies, source.armies-1))
431         a_dice = sorted([random.randint(1, 6) for i in range(armies)],
432                         reverse=True)
433         t_dice = sorted([random.randint(1, 6) for i in range(min(2, target.armies))],
434                         reverse=True)
435         a_dead = 0
436         t_dead = 0
437         for a,d in zip(a_dice, t_dice):
438             if d >= a:
439                 a_dead += 1
440             else:
441                 t_dead += 1
442         source.armies -= a_dead
443         target.armies -= t_dead
444         if target.armies == 0:
445             self.takeover(source, target, remaining_attackers=armies-a_dead)
446             print '%s conquered %s from %s with %d:%d.  Deaths %d:%d.  Remaining %d:%d' \
447                 % (source.player, target, source, armies, len(t_dice),
448                    a_dead, t_dead, source.armies, target.armies)
449             assert target.armies > 0, target
450             return True
451         print '%s attacked %s from %s with %d:%d.  Deaths %d:%d.  Remaining %d:%d' \
452             % (source.player, target, source, armies, len(t_dice),
453                a_dead, t_dead, source.armies, target.armies)
454         assert target.armies > 0, target
455         return False
456     def takeover(self, source, target, remaining_attackers):
457         source.armies -= remaining_attackers
458         target.armies += remaining_attackers
459         target.player = source.player
460         support = source.player.phase_support_attack(self.world, source, target)
461         if support < 0 or support >= source.armies:
462             raise PlayerError('Cannot support from %s to %s with %d armies, only %d available'
463                               % (source, target, support, source.armies-1))
464         source.armies -= support
465         target.armies += support
466     def player_killed(self, player, killer):
467         player.alive = False
468         killer.hand.extend(player.hand)
469         if len(self.living_players()) > 1:
470             while len(killer.hand) > 5:
471                 self.play_cards_and_place_armies(killer)
472         print '%s killed by %s' % (player, killer)
473     def living_players(self):
474         return [p for p in self.players if p.alive == True]
475
476
477 def generate_earth():
478     w = World('Earth')
479     c = Continent('North America', 5)
480     c.append(Territory('Alaska', 'ala', 1, ['kam', 'nwt']))    
481     c.append(Territory('Northwest Territory', 'nwt', 2, ['alb', 'ont', 'gre']))
482     c.append(Territory('Greenland', 'gre', 3, ['ont', 'que', 'ice']))
483     c.append(Territory('Alberta', 'alb', 1, ['ont', 'wus']))
484     c.append(Territory('Ontario', 'ont', 2, ['wus', 'eus', 'que']))
485     c.append(Territory('Quebec', 'que', 3, ['eus']))
486     c.append(Territory('Western United States', 'wus', 1, ['eus', 'cam']))
487     c.append(Territory('Eastern United States', 'eus', 2, ['cam']))
488     c.append(Territory('Central America', 'cam', 3, ['ven']))
489     w.append(c)
490     
491     c = Continent('Europe', 5)
492     c.append(Territory('Iceland', 'ice', 1, ['gbr', 'sca']))
493     c.append(Territory('Scandanavia', 'sca', 2, ['gbr', 'neu', 'ukr']))
494     c.append(Territory('Ukraine', 'ukr', 3, ['neu', 'seu', 'ura', 'afg', 'mea']))
495     c.append(Territory('Great Britain', 'gbr', 1, ['neu', 'weu']))
496     c.append(Territory('Northern Europe', 'neu', 2, ['weu', 'seu']))
497     c.append(Territory('Western Europe', 'weu', 3, ['naf', 'seu']))
498     c.append(Territory('Southern Europe', 'seu', 1, ['naf', 'egy', 'mea']))
499     w.append(c)
500
501     c = Continent('Asia', 7)
502     c.append(Territory('Urals', 'ura', 2, ['afg', 'chi', 'sib']))
503     c.append(Territory('Siberia', 'sib', 3, ['chi', 'mon', 'irk', 'yak']))
504     c.append(Territory('Yakutsk', 'yak', 1, ['irk', 'kam']))
505     c.append(Territory('Kamchatka', 'kam', 2, ['mon', 'jap']))
506     c.append(Territory('Irkutsk', 'irk', 3, ['mon']))
507     c.append(Territory('Mongolia', 'mon', 1, ['chi', 'jap']))
508     c.append(Territory('Japan', 'jap', 2))
509     c.append(Territory('Afghanistan', 'afg', 3, ['mea', 'indi', 'chi']))
510     c.append(Territory('China', 'chi', 1, ['indi', 'sia']))
511     c.append(Territory('Middle East', 'mea', 2, ['egy', 'eaf', 'indi']))
512     c.append(Territory('India', 'indi', 3, ['sia']))
513     c.append(Territory('Siam', 'sia', 1, ['indo']))
514
515     w.append(c)
516
517     c = Continent('South America', 2)
518     c.append(Territory('Venezuala', 'ven', 2, ['per', 'bra']))
519     c.append(Territory('Peru', 'per', 3, ['arg', 'bra']))
520     c.append(Territory('Brazil', 'bra', 1, ['arg', 'naf']))
521     c.append(Territory('Argentina', 'arg', 2))
522     w.append(c)
523
524     c = Continent('Africa', 3)
525     c.append(Territory('North Africa', 'naf', 3, ['egy', 'eaf', 'con']))
526     c.append(Territory('Egypt', 'egy', 1, ['eaf']))
527     c.append(Territory('East Africa', 'eaf', 2, ['con', 'saf', 'mad']))
528     c.append(Territory('Congo', 'con', 3, ['saf']))
529     c.append(Territory('South Africa', 'saf', 1, ['mad']))
530     c.append(Territory('Madagascar', 'mad', 2))
531     w.append(c)
532
533     c = Continent('Australia', 2)
534     c.append(Territory('Indonesia', 'indo', 3, ['ngu', 'wau']))
535     c.append(Territory('New Guinea', 'ngu', 1, ['wau', 'eau']))
536     c.append(Territory('Western Australia', 'wau', 2, ['eau']))
537     c.append(Territory('Eastern Australia', 'eau', 3))
538     w.append(c)
539
540     w._resolve_link_names()
541     return w
542
543 if __name__ == '__main__':
544     import doctest
545     import sys
546
547     failures,tests = doctest.testmod()
548     if failures > 0:
549         sys.exit(1)
550
551     world = generate_earth()
552     players = [Player('A'), Player('B')]
553     e = Engine(world, players)
554     e.run()