import random


class MadLibs:

    def __init__(self):
        self.woordenschat = {}

    def leren(self, categorie, nieuwe_woorden):
        if isinstance(nieuwe_woorden, str):
            nieuwe_woorden = {nieuwe_woorden.lower()}
        else:
            nieuwe_woorden = set(w.lower() for w in nieuwe_woorden)
        if categorie.lower() in self.woordenschat:
            self.woordenschat[categorie].update(nieuwe_woorden)
        else:
            self.woordenschat.update({categorie.lower(): nieuwe_woorden})

    def suggereren(self, categorie):
        assert categorie.lower() in self.woordenschat, 'onbekende categorie'
        woord = random.choice(list(self.woordenschat[categorie.lower()]))
        if categorie.isupper():
            woord = woord.upper()
        elif categorie[0].isupper():
            woord = woord.capitalize()
        return woord

    def invullen(self, zin):
        zin = zin.split('_')
        for index, woord in enumerate(zin):
            if index % 2:
                assert woord.lower() in self.woordenschat, 'onbekende categorie'
                invulwoord = self.suggereren(woord)
                zin[index] = invulwoord
        return ''.join(zin)