# creature script that simulates a conversation with a creature

# configuration:

## name attached to the creature's messages
NAME =   'summer'
## name attached to the user's messages
MYNAME = 'you'
## "corpus file" containing examples to learn from
CORPUS = 'summer.corpus.txt'

# end of configuration


import random
import re

filters = list()
bigrams = dict()
trigrams = dict()

with open(CORPUS, 'r') as f:
    corpus = '\n' + f.read() + '\nUSERTEXT'

tokens = ['\nUSERTEXT', ' CREATURETEXT', ' ENDCONVO']

def tokenise(tokens, text):
    result = list()
    for word in re.split('(?=[^A-Za-z:0-9<>])', text):
        if word not in tokens:
            tokens.append(word)
        result.append(tokens.index(word))
    return result

def detokenise(tokens, stuff):
    return ''.join([tokens[t] for t in stuff])

def add_filters(filters, material, window):
    start = 0
    while start + window < len(material):
        filt = set()
        after = set()
        for t in material[start:start+window]:
            if t > 2: filt.add(t)
        for t in material[start+window:start+int(window*1.5)]:
            after.add(t)
        if len(filt) > 0: filters.append((filt, after))
        start += window // 2

def add_ngrams(ngrams, material, window):
    for i in range(len(material) - window):
        key = tuple(material[i:i+window])
        if key not in ngrams:
            ngrams[key] = dict()
        gram = material[i+window]
        if gram not in ngrams[key]:
            ngrams[key][gram] = 0
        ngrams[key][gram] += 1

def add_new_to_corpus(question):
    print(f" * {NAME} didn't know what to say. Please suggest something appropriate:")
    newanswer = input('> ')
    newstuff = tokenise(tokens, newanswer)
    with open(CORPUS, 'a') as f:
        f.write(f'USERTEXT {question} CREATURETEXT {newanswer}\n')
    return newstuff
    

def infer(filters, bigrams, trigrams, context):
    scores = dict()
    nkey = tuple(context[-2:])
    if nkey in trigrams:
        for a, m in trigrams[nkey].items():
            if a not in scores:
                scores[a] = 0.001
            #scores[a] += m * 0.2 * (n+1)
    permit_bullshit = len(scores) <= 0
    bigram = bigrams[(context[-1],)]
    for f, a in filters:
        m = 0
        for t in context[-len(f):]:
            if t in f:
                m += 0.1 if t <= 2 else 1
        for n in a:
            if n in scores:
                scores[n] += m / len(f)
            elif permit_bullshit:
                scores[n] = m / len(f) + (bigram[n] if n in bigram else 0)
    choices = random.choices(list(scores.items()), list(scores.values()), k=1)
    choice = max(choices, key=lambda c: c[1])[0]
    maxv = max([c[1] for c in choices])
    #print(f'confidence: {maxv}')
    if maxv < 1: return -1
    return choice

material = tokenise(tokens, corpus)
add_filters(filters, material, 2)
add_filters(filters, material, 4)
add_filters(filters, material, 8)
add_filters(filters, material, 16)
add_filters(filters, material, 32)

add_ngrams(bigrams, material, 1)
add_ngrams(trigrams, material, 2)

newstuff = []

shouldask = True
while True:
    if shouldask:
        try:
            question = input(f'<{MYNAME}> ')
        except EOFError:
            exit()
        newstuff += tokenise(tokens, '\nUSERTEXT ' + question + ' CREATURETEXT')
    promptlength = len(newstuff)
    done = False
    shouldquit = False
    maxn = 5000
    while not done and maxn > 0:
        nexttoken = infer(filters, bigrams, trigrams, newstuff)
        if (nexttoken == -1):
            answertokens = add_new_to_corpus(question)
            newstuff += tokenise(tokens, ' ') + answertokens + [0]
            nexttoken = 0
            done = True
            question = detokenise(tokens, answertokens)
            continue
        done = nexttoken <= 2
        shouldquit = nexttoken == 2
        shouldask = nexttoken == 0
        newstuff.append(nexttoken)
        maxn -= 1
        
    answer = detokenise(tokens, newstuff[promptlength:-1])[1:].strip()
    if len(answer) > 0:
        print(f'<{NAME}> ', end='')
        print(answer)
        

    if shouldquit: break