changeset 151:372d886f9e70

New suggest_research() method on Lab.
author Jeremy Thurgood <firxen@gmail.com>
date Fri, 11 May 2012 20:06:36 +0200
parents 0090ecf08544
children 5b58e2150a4f
files gamelib/lab.py gamelib/schematics.py gamelib/tests/test_lab.py
diffstat 3 files changed, 145 insertions(+), 5 deletions(-) [+]
line wrap: on
line diff
--- a/gamelib/lab.py	Fri May 11 17:17:44 2012 +0200
+++ b/gamelib/lab.py	Fri May 11 20:06:36 2012 +0200
@@ -1,11 +1,123 @@
 # -*- test-case-name: gamelib.tests.test_lab -*-
 
-from random import random, choice
+from random import random, choice, sample
+
+import networkx
 
 from gamelib import research, schematics
 from gamelib.game_base import get_subclasses
 
 
+class ScienceGraph(object):
+    def __init__(self, all_science, known_science):
+        self.graph = networkx.DiGraph()
+        self.all_science = all_science
+        self.known_science = known_science
+        self.add_all_science()
+        self.tag_known_science()
+        assert networkx.is_directed_acyclic_graph(self.graph)
+
+    def add_all_science(self):
+        # Add level 0 of everything to the graph.
+        for science in self.all_science:
+            self.graph.add_node((science, 0), known=False)
+
+        # Walk dependencies and fill in intermediate nodes.
+        for science in self.all_science:
+            for dep in science.PREREQUISITES:
+                self.add_node_string(*dep)
+                self.graph.add_edge(dep, (science, 0))
+
+    def add_node_string(self, science, level):
+        node = (science, level)
+        if node in self.graph:
+            return node
+
+        # We prepopulate with level 0 of eveything.
+        assert level >= 0
+
+        self.graph.add_node(node, known=False)
+        parent = self.add_node_string(science, level - 1)
+        self.graph.add_edge(parent, node)
+        return node
+
+    def tag_known_science(self):
+        for science in self.known_science:
+            # We may know more of this than the graph has.
+            self.add_node_string(type(science), science.points + 1)
+            for i in range(science.points + 1):
+                self.graph.node[(type(science), i)]['known'] = True
+
+    def is_known(self, node):
+        return self.graph.node[node]['known']
+
+    def distances_to_known(self, science):
+        nodes = set()
+        nodes_to_check = [(science, 0)]
+
+        while nodes_to_check:
+            node = nodes_to_check.pop()
+            if self.is_known(node) or node in nodes:
+                continue
+            nodes.add(node)
+            nodes_to_check.extend(self.graph.predecessors(node))
+
+        distances = {}
+        for node in nodes:
+            distances.setdefault(node[0], 0)
+            distances[node[0]] += 1
+
+        return distances
+
+    def count_unknown(self, sciences):
+        return len([s for s in sciences if not self.is_known((s, 0))])
+
+    def find_prospects(self):
+        prospects = {'research': [], 'schematic': []}
+        for science in self.all_science:
+            distances = self.distances_to_known(science)
+            if not distances:
+                # We already know this thing.
+                continue
+            # Remove the thing we're trying to get from the distance.
+            distances.pop(science)
+            if self.count_unknown(distances.keys()) > 0:
+                # We only want direct breakthroughs.
+                continue
+            prospects[science.SCIENCE_TYPE].append(
+                (sum(distances.values()), science, distances))
+        return dict((k, sorted(v)) for k, v in prospects.items())
+
+    def find_promising_areas(self, size=3):
+        basic_science = False
+        areas_for_research = set()
+        areas_for_schematics = set()
+        prospects = self.find_prospects()
+
+        for points, target, distances in prospects['schematic']:
+            if points > 0:
+                # We need nonzero points in these things.
+                areas_for_schematics.update(distances.keys())
+            else:
+                # Any of these things qualify us.
+                areas_for_schematics.update(p for p, _ in target.PREREQUISITES)
+
+        for points, target, distances in prospects['research']:
+            if points == 0:
+                basic_science = True
+            else:
+                areas_for_research.update(distances.keys())
+
+        suggestions = []
+        k = min(size, len(areas_for_schematics))
+        suggestions.extend(sample(areas_for_schematics, k))
+        if len(suggestions) < size:
+            k = min(size - len(suggestions), len(areas_for_research))
+            suggestions.extend(sample(areas_for_research, k))
+
+        return basic_science, suggestions
+
+
 class Lab(object):
     BASIC_RESEARCH_SUCCESS_RATE = 0.05
     BASIC_RESEARCH_SUCCESS_MULTIPLIER = 2
@@ -14,6 +126,7 @@
         self.science = []
         self.new_research = get_subclasses(research.ResearchArea)
         self.new_schematics = get_subclasses(schematics.Schematic)
+        self.all_science = [s for s in self.new_research + self.new_schematics]
 
         if init_data is not None:
             # Load stored state.
@@ -49,8 +162,8 @@
         new_science = [physics]
 
         # We get two other random sciences with no prerequisites.
-        for _ in range(2):
-            science = choice(self.find_new_research())()
+        for science in sample(self.find_new_research(), 2):
+            science = science()
             self._gain_science(science)
             new_science.append(science)
 
@@ -141,3 +254,13 @@
             self._gain_science(breakthrough)
             breakthroughs = [breakthrough]
         return breakthroughs
+
+    def suggest_research(self):
+        """Suggest research areas to pursue.
+
+        Return value is a tuple of (bool, list), where the first element
+        indicates whether basic research might pay off and the second contains
+        Science classes that can be profitably pursued.
+        """
+        graph = ScienceGraph(self.all_science, self.science)
+        return graph.find_promising_areas()
--- a/gamelib/schematics.py	Fri May 11 17:17:44 2012 +0200
+++ b/gamelib/schematics.py	Fri May 11 20:06:36 2012 +0200
@@ -19,6 +19,7 @@
     'AQUATIC',
     'INTELLIGENCE',
     'AI',
+    'COUNTERMEASURE',
     )
 
 K = 1000
@@ -189,7 +190,7 @@
 
 
 class LaserGun(Schematic):
-    NAME = "Laser Gun"
+    NAME = "laser gun"
     COST = 300
     CATEGORIES = (cat.HAND_WEAPON,)
     PREREQUISITES = (
@@ -197,6 +198,16 @@
         )
 
 
+class EmpMissile(Schematic):
+    NAME = "EMP missile"
+    COST = 1500
+    CATEGORIES = (cat.COUNTERMEASURE,)
+    PREREQUISITES = (
+        (research.Electrickery, 5),
+        (research.Rocketry, 2),
+        )
+
+
 class DoomsdayVirus(Schematic):
     NAME = "doomsday virus"
     COST = 100 * K
--- a/gamelib/tests/test_lab.py	Fri May 11 17:17:44 2012 +0200
+++ b/gamelib/tests/test_lab.py	Fri May 11 20:06:36 2012 +0200
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from gamelib.lab import Lab
+from gamelib.lab import Lab, ScienceGraph
 from gamelib import research, schematics
 
 
@@ -9,6 +9,7 @@
         'research.Physics': 5,
         'research.Rocketry': 2,
         'research.Electrickery': 3,
+        'schematic.MachineGun': 1,
         },
     }
 
@@ -46,3 +47,8 @@
         new_science = dict(
             LAB_DATA['science'].items() + [('schematic.LightningGun', 0)])
         self.assertEqual({'science': new_science}, lab.save_data())
+
+    def test_science_graph(self):
+        lab = Lab(LAB_DATA)
+        self.assertEqual((True, [research.Electrickery, research.Physics]),
+                         lab.suggest_research())