Moved deck loading to the deck object, toying with probability

analyze.py

@@ -1,8 +1,19 @@
-import re
 import sys
+import math
 from mtg import *
 import database
 
+def combinations(total, n):
+    return math.factorial(total) / (math.factorial(n) * math.factorial(total - n))
+def h(n, z, x, y):
+    return (combinations(x, n) * combinations(y - x, z - n)) / float(combinations(y, z))
+
+def draw_probability(available, deck, count=1, turn=1, used=0):
+    if count == 0:
+        return h(0, 7 + turn - 1 - used, available, len(deck))
+    total = sum([h(n, 7 + turn - 1 - used, available, len(deck)) for n in range(count)])
+    return 1 - total
+
 def graph(data, title, hideEmpty=True):
     categories = sorted(data.keys())
     cjust = sorted([len(str(k)) for k in data.keys()]).pop()
@@ -61,30 +72,36 @@ def rarities(deck):
         data[Card.rarities[card.rarity]] += 1
     return data
 
-db = database.TextDB('db.txt')
+if __name__ == '__main__':
+    db = database.TextDB('db.txt')
 
-if len(sys.argv) < 2:
-    print 'No decks specified.'
-    sys.exit(1)
-for filename in sys.argv[1:]:
-    print 'Loading deck in "{0}"...'.format(filename)
-    deck = Deck()
-    with open(filename, 'r') as f:
-        for line in f:
-            line = line.strip()
-            if not re.match('^\d+ .+$', line):
-                continue
-            line = line.split(' ')
-            count = int(line.pop(0))
-            name = ' '.join(line)
-            card = db.getCard(name)
-            if not card:
-                print 'Invalid card: ', name
-                continue
-            deck.extend(card * count)
-    
-    print 'Loaded', len(deck), 'cards.'
-    graph(mana_curve(deck), '\nMana Curve', hideEmpty=False)
-    graph(types(deck), '\nCard Types')
-    graph(colors(deck), '\nColors')
-    graph(rarities(deck), '\nRarities')
+    if len(sys.argv) < 2:
+        print 'No decks specified.'
+        sys.exit(1)
+    for filename in sys.argv[1:]:
+        print 'Loading deck in "{0}"...'.format(filename)
+        deck = Deck()
+        """
+        with open(filename, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if not re.match('^\d+ .+$', line):
+                    continue
+                line = line.split(' ')
+                count = int(line.pop(0))
+                name = ' '.join(line)
+                card = db.getCard(name)
+                if not card:
+                    print 'Invalid card: ', name
+                    continue
+                deck.extend(card * count)
+        """
+        deck.load(filename, db)
+        print 'Loaded', len(deck), 'cards.'
+        graph(mana_curve(deck), '\nMana Curve', hideEmpty=False)
+        graph(types(deck), '\nCard Types')
+        graph(colors(deck), '\nColors')
+        graph(rarities(deck), '\nRarities')
+
+        print '\nProbability of drawing 4 mana in 4 turns: {0:.3f}'.format(100 * draw_probability(len([c for c in deck if 'land' in c.type.split(' ')]), deck, 4, 4))
+        print 'Probability of drawing a creature costing 4 or fewer mana in 4 turns: {0:.3f}'.format(100 * draw_probability(len([c for c in deck if 'creature' in c.type.split(' ') and c.cost.converted() <= 4]), deck, 1, 4, 4))

mtg.py

@@ -1,3 +1,4 @@
+import re
 import copy
 import random
 import logging
@@ -257,3 +258,17 @@ class Deck(CardList):
         random.shuffle(self)
     def cards(self):
         return self.__cards
+    def load(self, filename, db):
+        with open(filename, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if not re.match('^\d+ .+$', line):
+                    continue
+                line = line.split(' ')
+                count = int(line.pop(0))
+                name = ' '.join(line)
+                card = db.getCard(name)
+                if not card:
+                    # TODO: The database should log an error
+                    continue
+                self.extend(card * count)