stroopwafel/Chatbot.py

import asyncio
import discord
from discord.ext import commands
import numpy as np
import random
from datetime import datetime
import re

class Chatbot(commands.Cog):
	"""Een heuse chatbot. Under construction..."""
	
	def __init__(self, bot):
		self.bot = bot
		self.voice_states = {}
		self.END_OF_MESSAGE = "<EOM>"
		self.START_OF_MESSAGE = "<SOM>"
		self.pmLast = datetime(2018, 1, 1, 1, 1, 1)
		self.pmCount = 0
		try:
			np.load('markov_dict.npy', allow_pickle=True)
		except FileNotFoundError:
			#If there's no initial file, make a new one:
			initFile = {'pairs':{self.START_OF_MESSAGE:{'Hoi':1},'Hoi':{self.END_OF_MESSAGE:1}}, 'trairs':{}, 'responsePairs':{}}
			np.save("markov_dict.npy", initFile)
		except:
			print("ERROR! - Could not open or create Markov chat file")
			
			
		#Do stuff when a message comes in:
		@self.bot.event
		@asyncio.coroutine
		def on_message(message):
			if message.author.bot: 
				return #ignore your own messages, Stroop
			
			if message.content.strip() == "" or message.content.strip()[0] != '!': #don't respond to a command
				if message.content.lower() == "nihao":
					yield from message.channel.send("kankerlauw")
					return
				if message.content.lower() == "lol":
					yield from message.channel.send("laffing aut laut")
					return
				if self.bot.user.mentioned_in(message) and "@everyone" not in message.content:
					MENTIONED_RESPONSES = [
						"Hoi",
						"hmm?",
						"JA WAT?",
						"Godve...Godver-fucking-domme! Wie de fuck heeft me zojuist gepingt? Wie de fuck heeft de ballen om mij te pingen?! Als ik jou vind hè, ik ram je helemaal de tyfus in. Ik ga je kapot maken. Welke breincel in jouw hoofd dacht dat het grappig was om mij te pingen? Noem jij dat 'grappig'? Ik zal je godverdomme iets laten zien wat grappig is, ja. Ik neuk je moeder zo hard dat ze van me gaat houden, godverdomme. Dat krijg je er van als je me pingt, vuile teringleijer. Jouw soort mag wat mij betreft uitsterven. Weet je, ik ga je helemaal geen aandacht meer geven. Je mag de tering krijgen, en daar blijft het bij. Kutjong.",
						self.getResponseForMessage("stroopwafel")
					]
					try:
						MENTIONED_RESPONSES.append("Wat mot je, " + random.choice(self.bot.get_cog("Shitpost").insult) + "?") 
					except:
						print("ERROR! - Could not import insults")
					yield from message.channel.send(random.choice(MENTIONED_RESPONSES))
					return
				if re.match("^[0-9]+([,.]?[0-9]+)*?$", message.content.strip()) and random.randint(0,9) == 0:
					yield from message.channel.send("That's Numberwang!") # https://www.youtube.com/watch?v=0obMRztklqU
					return
				if "DMChannel" in type(message.channel).__name__: #respond to messages in PM channels
					if (self.isAllowedToLearnFromPm(message)): #avoid spam-learning
						self.pmCount += 1
						self.pmLast = datetime.now()
						yield from self.learnFromMessagesInChannel(message.channel, 1)
					yield from self.respondToMessage(message)
				elif message.channel.id == self.bot.hdChannels["belangrijk"]: #don't abuse the "belangrijk" channel
					pass
				else:
					a = random.randint(0,999) # probablity
					if a < 40:
						yield from self.learnFromMessagesInChannel(message.channel, 10)
					if "stroop" in message.content.lower():
						a = a / 20 # significantly increase the chance of responding if somebody talks about him.
					if message.channel.id == self.bot.hdChannels["shitpost"] and a < 40:  #shitpost channel
						yield from self.respondToMessage(message)
					elif a < 10:                                             #all other channels
						yield from self.respondToMessage(message)

			#process commands:
			yield from self.bot.process_commands(message)
			
	def __check(self, ctx):
		return True
		
	
	#The "Guus filter"
	#Don't allow to learn from very large personal messages
	#Only learn from 10 personal messages per hour at most
	def isAllowedToLearnFromPm(self, message):
		if len(message.content.strip()) > 150:
			return False
		diff = datetime.now() - self.pmLast
		if diff.days > 0 or diff.seconds > 3600:
			self.pmCount = 0
			return True
		if self.pmCount > 10:
			return False
		return True
	
	#Make word pairs:
	def makePairs(self, corpus):
		for i in range(len(corpus)-1):
			yield (corpus[i], corpus[i+1])
		
	#Make word trairs. Yes, I completely made up "trair"
	#It's like pairs, but with three, but not really, because there are two. OK?
	#[1 0 1]
	def makeTrairs(self, corpus):
		for i in range(1, len(corpus)-1):
			yield (corpus[i-1], corpus[i+1])
		
	#Learn from a human message
	#Returns an updated dict
	def learnFrom(self, wordDictPairs, wordDictTrairs, message: str):
		newPairDict = wordDictPairs
		newTrairDict = wordDictTrairs
		
		#If the message is only a tag, don't learn from it. return
		if message.count(' ') == 0 and message[0:2] == "<@":
			return newPairDict, newTrairDict;
		
		message = self.START_OF_MESSAGE + " " + message + " " + self.END_OF_MESSAGE
		corpus = message.split()
		for n, word in enumerate(corpus):
			if word == "@everyone":
				corpus[n] = "@iedereen"
		pairs = self.makePairs(corpus)
		trairs = self.makeTrairs(corpus)

		#Add wordcounts to pair dict
		for word_1, word_2 in pairs:
			if word_1 in newPairDict.keys():
				if word_2 in newPairDict[word_1].keys():
					newPairDict[word_1][word_2] += 1
				else:
					newPairDict[word_1][word_2] = 1
			else:
				newPairDict[word_1] = {word_2 : 1}		
		#Add wordcounts to trair dict
		for word_1, word_3 in trairs:
			if word_1 in newTrairDict.keys():
				if word_3 in newTrairDict[word_1].keys():
					newTrairDict[word_1][word_3] += 1
				else:
					newTrairDict[word_1][word_3] = 1
			else:
				newTrairDict[word_1] = {word_3 : 1}
				
		return newPairDict, newTrairDict;
	
	#Learn to respond to messages
	#Returns a new dict with response pairs
	def learnResponse(self, wordDictResp, message1: str, message2: str):
		STOP_WORDS = ["aan", "achter", "al", "ben", "dan", "dat", "de", "die", "dit", "een", "en", "er", "gehad", "had", "heb", 
			"hen", "het", "hun", "in", "is", "kon", "met", "na", "of", "om", "ook", "op", "tot", "van", "was", "wat", "zo", 
			"a", "an", "the", "am", "was", "were", "be", "not", "at"]
			
		newRespPairDict = wordDictResp
		
		#Split messages into words and remove stopwords:
		message1List = []
		message2List = []
		for word in message1.split():
			if word not in STOP_WORDS:
				message1List.append(word)
		for word in message2.split():
			if word not in STOP_WORDS:
				message2List.append(word)
		
		#link words of response to words of initial message
		for word_1 in message1List:
			for word_2 in message2List:
				if word_1 in newRespPairDict.keys():
					if word_2 in newRespPairDict[word_1].keys():
						newRespPairDict[word_1][word_2] += 1
					else:
						newRespPairDict[word_1][word_2] = 1
				else:
					newRespPairDict[word_1] = {word_2 : 1}	

		return newRespPairDict

			

	@commands.command(pass_context=True, hidden=True)
	@asyncio.coroutine	
	def trainmarkovchat(self, ctx, channelId=409751773595566081, nrOfMessages=50):
		"""Train that chain, jermaine
		param channelId: use this channel to train
		param nrOfMessages: use the x most recent messages
		"""
		
		channel=self.bot.get_channel(channelId)
		yield from self.learnFromMessagesInChannel(channel, nrOfMessages)
		
		yield from ctx.channel.send("Ik heb weer wat nieuwe woordjes geleerd. :slight_smile:")

	
	#Learns from an x number of messages from a certain channel
	@asyncio.coroutine
	def learnFromMessagesInChannel(self, channel, nrOfMessages):
		
		#Open current chain:
		wordDict = np.load('markov_dict.npy', allow_pickle=True).item()
		
		#Iterate through the last x messages from the channel:
		messageIterator = channel.history(limit=nrOfMessages, before=None, after=None, oldest_first=False, around=None)
		allMessages = []
		for i in range(nrOfMessages):
			msg = yield from anext(messageIterator)
			allMessages.append(msg)
		
		i = 0
		nextMessageIndex = -10
		#Note: the loop goes from new to old
		for msg in allMessages:
			messageContent = msg.content.replace('(', '').replace(')', '').replace('"', '')
			messageContent = messageContent.strip()

			if messageContent == "":
				pass # Bericht is leeg. Negeer.
			elif msg.author.bot:
				pass # Dit is een bot chat. Negeer.
			elif messageContent[0] == "!":
				pass # Dit is een commando. Negeer.
			elif "```" in messageContent:
				pass # Daar zit een codeblok in. Negeer.
			elif messageContent[0:7] == "http://" or messageContent[0:8] == "https://":
				pass # Een link. Negeer. TODO: post af en toe een random link die geen repost is
			else:
				# Usable message!
				wordDict['pairs'], wordDict['trairs'] = self.learnFrom(wordDict['pairs'], wordDict['trairs'], messageContent)
				
				if (i-nextMessageIndex) == 1 and msg.author != nextMessage.author: #the previous usable messages was the one right after this one and was not from the same author
					timeDelta = nextMessage.created_at - msg.created_at
					#TODO check for different author
					if timeDelta.days == 0 and timeDelta.seconds < 3600: #less than an hour difference
						#this is counted as a response. learn from it
						wordDict['responsePairs'] = self.learnResponse(wordDict['responsePairs'], messageContent, nextMessageContent)
				
				#save for next iteration:
				nextMessage = msg
				nextMessageContent = messageContent
				nextMessageIndex = i
			i += 1
			
		np.save("markov_dict.npy", wordDict)
		
	
	@commands.command(pass_context=True, hidden=False)
	@asyncio.coroutine	
	def reageer(self, ctx):
		"""Laat Stroop ergens op reageren.
		Stroop leert van onze gesprekken. Hij kan zich nu voordoen als een van ons. Soort van.
		"""
		
		#Get previous message to respond to
		messageIterator = ctx.channel.history(limit=3, before=None, after=None, oldest_first=False, around=None)
		messageToRespondTo = yield from anext(messageIterator) #this will be the !reageer command. skip
		messageToRespondTo = yield from anext(messageIterator)
		if messageToRespondTo.content.strip()[0] == '!' or messageToRespondTo.author.bot: #another chance
			messageToRespondTo = yield from anext(messageIterator) #but it's not a huge problem if he responds to himself
		if messageToRespondTo.content.strip()[0] == '!': 
			wordsToRespondTo = []
		else:
			wordsToRespondTo = messageToRespondTo.content.strip().split()
		#print("respond to: ", wordsToRespondTo)

		#Delete command
		try:
			yield from ctx.message.delete()
		except:
			pass
			
		response = self.getResponseForMessage(wordsToRespondTo)
		yield from ctx.channel.send(response)
	
	# Calculate probabilities and pick a next word from a word pair dict:
	# https://stackoverflow.com/questions/835092/python-dictionary-are-keys-and-values-always-the-same-order 
	def pickNextWord(self, pairPossibilities):
		possible_words = list(pairPossibilities.keys())
		weights = np.array(list(pairPossibilities.values()))
		probs = weights / weights.sum()
		return np.random.choice(possible_words, p=probs)
	
	#Respond to a message.
	#Assumes the message isn't a command
	@asyncio.coroutine
	def respondToMessage(self, message):
		wordsToRespondTo = message.content.strip().split()
		
		response = self.getResponseForMessage(wordsToRespondTo)
		yield from message.channel.send(response)
		
	#Gets a response
	def getResponseForMessage(self, wordsToRespondTo):
		
		MAX_N_WORDS = 100 #Don't use more words than this
		
		#Get dictionary/Markov chain
		wordDict = np.load('markov_dict.npy', allow_pickle=True).item()
		wordPairDict = wordDict['pairs']
		wordTrairDict = wordDict['trairs']
		responsePairDict = wordDict['responsePairs']
		#TODO store in global variable to minimise IO?
		
		#Start with <SOM>
		chain = [self.START_OF_MESSAGE]

		#Go through the chain:
		for i in range(1, MAX_N_WORDS):
			
			if chain[-1] == self.END_OF_MESSAGE:
				break
			
			#Find the possibilities:
			pairPossibilities = wordPairDict[chain[-1]].copy()
			
			#Increase score if it would be a good response:
			for userWord in wordsToRespondTo:
				if userWord in responsePairDict:
					for goodResponseWord in responsePairDict[userWord]:
						if goodResponseWord in pairPossibilities.keys():
							if i == 1:
								pairPossibilities[goodResponseWord] += responsePairDict[userWord][goodResponseWord] * 2 #first word? try to respond
							else:
								pairPossibilities[goodResponseWord] += responsePairDict[userWord][goodResponseWord] * 0.1
			#Increase score for trairs:
			if i > 1: #we can only compare trairs if we have at least two words already
				if chain[-2] in wordTrairDict:
					trairPossibilities = wordTrairDict[chain[-2]].copy()
					#If a trair matches, increase the frequency of the pair
					for word2 in pairPossibilities.keys():
						if word2 in trairPossibilities.keys():
							pairPossibilities[word2] += trairPossibilities[word2]
							
			#Pick a next word:
			chain.append(self.pickNextWord(pairPossibilities))
	

		if chain[0] == self.START_OF_MESSAGE:
			chain = chain[1:-1] #remove <SOM> and <EOM>
		else:
			chain = chain[:-1] #remove <EOM>
		
		#remove 90% of tags:
		for n, word in enumerate(chain):
			if word[0:2] == "<@":
				a = random.randint(1,10)
				if a != 1:
					chain[n] = "@jemoeder"
			
		#print(' '.join(chain))
		#yield from ctx.channel.send(' '.join(chain))
		return ' '.join(chain)
		
		
		
		
	#Try to count the number of syllables in a word
	#Used to make Haikus
	def countSyllables(self, word: str):
		KNOWN_ACRONYMS = ["aub", "svp", "ivm", "nr", "zsm", "mbt", "tmi", "fyi", "afk", "omg", "idk", "coc"]
		word = word.lower().strip()
		
		#for acryonyms, just count the letters:
		if word in KNOWN_ACRONYMS or re.match("^([a-z]{1}[.]{1}){2,}$", word) or re.match("^([bcdfghjklmnpqrstvwxz]){2,}$", word):
			return len(word.replace(".", ""))
		
		#count vowels surrounded by consonants
		count = len(re.findall("[bcdfghjklmnpqrstvwxyz]{1}[aeiouyëäïöüéèáà]{1,2}[bcdfghjklmnpqrstvwxz]{0,1}|[aeiouyëäïöüéèáà]{1,2}[bcdfghjklmnpqrstvwxz]{1}", word))
		count += len(re.findall("[ëäï]", word)) #ideeën etc
		
		return max(1, count)
		
	#Append a message change with a numer of syllables
	#If it returns False, it couldn't find anything
	#It's recursive, yo
	def appendMessageWithNrOfSyllables(self, nrOfSyl: int, chain, wordPairDict):	
		#Find the possibilities:
		if chain[-1] == "\n":
			pairPossibilities = wordPairDict[chain[-2]].copy()
		else:
			pairPossibilities = wordPairDict[chain[-1]].copy()
		
		for poss in list(pairPossibilities.keys()):
			if self.countSyllables(poss) > nrOfSyl or poss == self.END_OF_MESSAGE:  #TODO change to constant
				del pairPossibilities[poss]
				
		if len(pairPossibilities) == 0:
			return False
		
		#Pick a next word:
		chain.append(self.pickNextWord(pairPossibilities))
		
		if self.countSyllables(chain[-1]) == nrOfSyl:
			return chain
		#RECURSION!
		return self.appendMessageWithNrOfSyllables( (nrOfSyl-self.countSyllables(chain[-1])) , chain, wordPairDict)
		
	#Make a nice Haiku	
	def getHaiku(self):
		MAX_N_TRIES = 50
		
		#Get dictionary/Markov chain
		wordDict = np.load('markov_dict.npy', allow_pickle=True).item()
		wordPairDict = wordDict['pairs']

		#try:
		for i in range(1, MAX_N_TRIES):
			chain = [self.START_OF_MESSAGE]
			
			chain = self.appendMessageWithNrOfSyllables(5, chain, wordPairDict)
			if chain is not False:
				chain.append("\n")
				chain = self.appendMessageWithNrOfSyllables(7, chain, wordPairDict)
			if chain is not False:
				chain.append("\n")
				chain = self.appendMessageWithNrOfSyllables(5, chain, wordPairDict)
			if chain is not False:
				break
				
		chain = chain[1:] #remove <SOM>
		return ' '.join(chain)
		
	@commands.command(pass_context=True, hidden=False)
	@asyncio.coroutine	
	def haiku(self, ctx):
		"""Stroop is een poëet. Laat hem een Haiku maken.
		"""
			
		response = self.getHaiku()
		yield from ctx.channel.send(response)

	@commands.command(pass_context=True, hidden=False)
	@asyncio.coroutine	
	def uitspraak(self, ctx, numPosts=1):
		"""Random uitspraak uit de hall of fame.
		Optioneel kan je een nummer meegeven voor meerdere uitspraken
		0 voor alle uitspraken."""
		try:
			yield from ctx.message.delete()
		except:
			pass
		try:
			message = ""
			with open("uitspraken.txt", "r") as file:
				uitspraken = file.readlines()
				random.shuffle(uitspraken)
				if numPosts == 0:
					numPosts = len(uitspraken)
				for i in range(numPosts):
					message += uitspraken[i]
				yield from ctx.channel.send(message)
		except:
			pass