Include state legend.

diff --git a/plugin/faction/faction.py b/plugin/faction/faction.py
index 8739792..ee183c1 100644 (file)
--- a/plugin/faction/faction.py
+++ b/plugin/faction/faction.py
@@ -949,6 +949,27 @@ class Faction(DBConnection):
       await cat.yelp(message.channel)
     return result
 
+  def state_marker(self, id):
+    name = eddb.state_name(id)
+    if name in ['Civil War', 'Election', 'War']:
+      return 'x'
+    elif name in ['Bust', 'Civil Unrest']:
+      return '|'
+    elif name in ['Famine', 'Outbreak']:
+      return '+'
+    elif name == 'Lockdown':
+      return 'X'
+    elif name == 'Expansion':
+      return '^'
+    elif name == 'Investment':
+      return 'h'
+    elif name == 'Boom':
+      return '.'
+    elif name == 'Retreat':
+      return 'v'
+    else:
+      return None
+
   async def graph_faction_history(self, destination, *, names = None, systems = None, width = 1024, height = 768):
     faction_ids = []
     system_ids = []
@@ -969,7 +990,7 @@ class Faction(DBConnection):
       log.warning("No factions or systems to graph!")
       return False
 
-    sql = 'select s.name as system_name, i.timestamp as timestamp, i.faction_id as faction_id, n.name as faction_name, i.influence as influence from faction_influence i, faction_systems s, faction_names n where i.faction_id=n.faction_id and i.eddb_id=s.eddb_id'
+    sql = 'select s.name as system_name, i.timestamp as timestamp, i.faction_id as faction_id, n.name as faction_name, i.influence as influence, i.state_id as state_id from faction_influence i, faction_systems s, faction_names n where i.faction_id=n.faction_id and i.eddb_id=s.eddb_id'
     params = []
     if len(system_ids):
       sql += ' and i.eddb_id in ({})'.format(', '.join(['?'] * len(system_ids)))
@@ -1006,10 +1027,11 @@ class Faction(DBConnection):
         dt = datetime.datetime.fromtimestamp(bot.parse_iso8601(day))
       entry = data.get(row[key])
       if entry is None:
-        data[row[key]] = { 'x': [], 'y': [] }
+        data[row[key]] = { 'x': [], 'y': [], 's': [] }
         entry = data[row[key]]
       entry['x'].append(dt)
       entry['y'].append(row['influence'] * 100.0)
+      entry['s'].append(row['state_id'])
     self.close_db()
 
     if not len(data):
@@ -1019,9 +1041,20 @@ class Faction(DBConnection):
     log.info('Graphing: {} in {}'.format(title, destination))
 
     fig, ax = plt.subplots()
+    state_lines = {}
     for label, entry in data.items():
-      plt.plot(entry['x'], entry['y'], label = label)
-    legend = ax.legend(bbox_to_anchor = (1, 0.5), loc = 'center left')
+      line = plt.plot(entry['x'], entry['y'], label = label)[0]
+      colour = line.get_color()
+      for state_id, indices in { s: [i for i, e in enumerate(entry['s']) if e == s] for s in set(entry['s']) }.items():
+        marker = self.state_marker(state_id)
+        if marker is None:
+          continue
+        line = plt.plot([entry['x'][i] for i in indices], [entry['y'][i] for i in indices], c = colour, marker = marker, ls = 'none')[0]
+        if state_id not in state_lines:
+          state_lines[eddb.state_name(state_id)] = line
+    state_legend = plt.legend(state_lines.values(), state_lines.keys(), bbox_to_anchor = (1, 0.5), loc = 'upper left')
+    plt.gca().add_artist(state_legend)
+    legend = ax.legend(bbox_to_anchor = (1, 0.5), loc = 'lower left')
 
     plt.title(title)
     plt.xlabel('Date')
@@ -1035,6 +1068,7 @@ class Faction(DBConnection):
     fig.set_size_inches(float(width) / dpi, float(height) / dpi)
     b = io.BytesIO()
     plt.savefig(b, format = 'png', dpi = dpi, bbox_extra_artists = (legend, ), bbox_inches = 'tight')
+    plt.close('all')
     b.seek(0)
     await bot.say(destination, '', attachment = bot.parse_attachment(b.getvalue(), 'influence.png'))
     return True