From 87b239a53c9e856e4147032f5b642a751deef5cf Mon Sep 17 00:00:00 2001
From: Valentin Ochs <a@0au.de>
Date: Sun, 9 Feb 2020 21:55:55 +0100
Subject: [PATCH] Add generation methods to Markov

---
 markov.hpp | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/markov.hpp b/markov.hpp
index 3a5d7c3..8f6ad98 100644
--- a/markov.hpp
+++ b/markov.hpp
@@ -1,6 +1,7 @@
 #pragma once
 #include <limits>
 #include <map>
+#include <random>
 #include <stdlib.h>
 #include <vector>
 
@@ -62,7 +63,37 @@ public:
     return false;
   }
 
+  template <typename IT, typename Generator>
+  T gen(IT iter, size_t size, Generator &g) const {
+    while (size > order()) {
+      iter++;
+      size--;
+    }
+
+    if (size == 0) {
+      return random_child(g);
+    }
+
+    return gen_impl(iter, size, g, std::bool_constant<(ORDER > 0)>{});
+  }
+
 protected:
+  template <typename IT, typename Generator>
+  T gen_impl(IT iter, size_t size, Generator &g,
+             std::bool_constant<true>) const {
+    auto child = children.find(*iter);
+    if (child == children.end()) {
+      return random_child(g);
+    }
+
+    return child->second.gen(iter + 1, size - 1, g);
+  }
+  template <typename IT, typename Generator>
+  T gen_impl(IT iter, size_t size, Generator &g,
+             std::bool_constant<false>) const {
+    return random_child(g);
+  }
+
   /* Finds the key or adds it */
   typename Children::iterator find(const T &key) {
     return children.try_emplace(key).first;
@@ -104,6 +135,19 @@ public:
       return 0.5 / double(total);
     }
   }
+
+  template <typename Generator> T random_child(Generator &g) const {
+    std::uniform_int_distribution<size_t> dist(0u, total - 1);
+    auto val = dist(g);
+    for (auto &child : children) {
+      if (val < child.second.total) {
+        return child.first;
+      }
+
+      val -= child.second.total;
+    }
+    return T{};
+  }
 };
 
 /* Leaf of the recursion with just a count */