From 1401d5cedef5ac20e9941708d846911c7ec9ca09 Mon Sep 17 00:00:00 2001
From: Brad Smith <bradster@infinitewarp.com>
Date: Sun, 12 Nov 2017 13:36:15 -0500
Subject: [PATCH] add markov generation mode to cli

---
 trekipsum/cli.py             | 11 +++++++++--
 trekipsum/markov/__init__.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/trekipsum/cli.py b/trekipsum/cli.py
index 1f2962d..3e4985a 100644
--- a/trekipsum/cli.py
+++ b/trekipsum/cli.py
@@ -3,7 +3,7 @@
 import argparse
 import logging
 
-from trekipsum import dialog
+from trekipsum import dialog, markov
 
 logger = logging.getLogger(__name__)
 
@@ -19,6 +19,8 @@ def positive(value):
 def parse_cli_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description='TrekIpsum generator')
+    parser.add_argument('-m', '--markov', action='store_true',
+                        help='use markov chain mode for generation')
     parser.add_argument('--speaker', type=str,
                         help='limit output to this speakers')
     parser.add_argument('-a', '--attribute', action='store_true',
@@ -47,7 +49,12 @@ def main_cli():
     loglevel = logging.DEBUG if args.debug else logging.CRITICAL
     logging.basicConfig(level=loglevel, format='%(asctime)s %(levelname)s: %(message)s')
     logger.setLevel(loglevel)
-    chooser = dialog.SqliteRandomChooser()
+
+    if args.markov is True:
+        chooser = markov.MarkovRandomChooser()
+    else:
+        chooser = dialog.SqliteRandomChooser()
+
     for paragraph in range(args.paragraphs):
         speaker = args.speaker
         lines = []
diff --git a/trekipsum/markov/__init__.py b/trekipsum/markov/__init__.py
index 4393d3d..365d582 100644
--- a/trekipsum/markov/__init__.py
+++ b/trekipsum/markov/__init__.py
@@ -6,6 +6,7 @@
 import six
 
 from ..dialog.sqlite import DEFAULT_SQLITE_PATH
+from ..exceptions import NoDialogFoundException
 
 SENTENCE_DELIMITER = ''  # special value for beginning/ending a sentence
 
@@ -214,3 +215,34 @@ def to_chain(self, context):
         for link in links:
             chain[link[0]].append((link[1], link[2]))
         return chain
+
+
+class MarkovRandomChooser(object):
+    """Walk Markov chains to generate dialog from datastore."""
+
+    def __init__(self):
+        """Initialize a new dialog markov chain chooser for dialog."""
+        self._datastore = DialogChainDatastore()
+        self._speaker_walker = ChainWalker(self._datastore.to_chain('speakers'))
+
+    def random_dialog(self, speaker):
+        """
+        Get random line of dialog, optionally limited to specific speaker.
+
+        Returns:
+            tuple containing (speaker name, line of dialog)
+        """
+        speaker = speaker.upper() if speaker else None
+        if speaker is None:
+            speaker = self.random_speaker()
+        dialog_walker = ChainWalker(self._datastore.to_chain(speaker))
+        try:
+            return speaker, dialog_walker.build_sentence()
+        except KeyError:
+            raise NoDialogFoundException(speaker)
+
+    def random_speaker(self, from_speaker=None):
+        """
+        Get random speaker name, optionally walking from a specific speaker.
+        """
+        return self._speaker_walker.next_word(from_speaker)