import random
from collections import defaultdict

# NAME GEN MODEL


class NameGenerator:
    def __init__(self, corpus, validate=None):
        self.model = self.build_model(self.pad_names_for_model(corpus))
        self.validate = validate

    def pad_names_for_model(self, names):
        new_names = []
        for name in names:
            new_names.append(f"<<{name}>")
        return new_names

    def build_model(self, words):
        model = defaultdict(list)
        for word in words:
            for pos in range(len(word) - 2):
                model[word[pos] + word[pos + 1]].append(word[pos + 2])
        return model

    def make_name(self):
        result = "<<"
        while True:
            window = result[-2:]
            nxt = random.choice(self.model[window])
            if nxt == ">":
                break
            result += nxt
        return result.removeprefix("<<")

    def generate_name(self):
        while True:
            name = self.make_name()
            if self.validate is None or self.validate(name) is None:
                return name