root/imbot/SABBrainWeightedGraph.py

import re, sys, time
import yaml
import pickle
import MontyLingua
import networkx as nx
import matplotlib.pyplot as plot # not officially compatible with Python 2.6
from random import choice
from random import uniform
import math

(i_word, i_tag, i_time) = (0, 1, 2)

class SABBrain:
    def __init__(self):
        self.ml = MontyLingua.MontyLingua()
        self.g = nx.LabeledDiGraph()
        self.struct = nx.LabeledDiGraph()
        self.current_timestamp = dict()
        self.elapsed_timestamp = dict()
        
        self.lastStruct = "start"
        self.struct.add_node("start")
        self.struct.add_node("declarative")
        self.struct.add_node("interrogative")
        self.struct.add_node("exclamatory")
        self.struct.add_node("imperative")
        
        self.update_graph(self.struct, "declarative", "declarative")
        self.update_graph(self.struct, "declarative", "interrogative")
        self.update_graph(self.struct, "declarative", "exclamatory")
        self.update_graph(self.struct, "declarative", "imperative")
        self.update_graph(self.struct, "interrogative", "declarative")
        self.update_graph(self.struct, "interrogative", "interrogative")
        self.update_graph(self.struct, "interrogative", "exclamatory")
        self.update_graph(self.struct, "interrogative", "imperative")
        self.update_graph(self.struct, "exclamatory", "declarative")
        self.update_graph(self.struct, "exclamatory", "interrogative")
        self.update_graph(self.struct, "exclamatory", "exclamatory")
        self.update_graph(self.struct, "exclamatory", "imperative")
        self.update_graph(self.struct, "imperative", "declarative")
        self.update_graph(self.struct, "imperative", "interrogative")
        self.update_graph(self.struct, "imperative", "interrogative")
        self.update_graph(self.struct, "imperative", "imperative")
        
        print 'Loading brain...',
        try:
            self.g = nx.read_gpickle(r"Brain\brain.pickle")
            print 'Done!'
        except IOError:
            print "file not found. This is normal behavior if the bot has never run before."
        
        print '\n'
        
    def halfSig(self, x):
        return 2 * (1/(1+math.e**(-x/10)) - .5)
    
    def read_file(self, filename):
        '''Returns the contents of a new-line-delimited text file as a list'''
        myfile = file(filename, 'r')
        contents = yaml.load(myfile.read())
        myfile.close()
        return contents
    
    def translate_from_html_to_english(self, text, debug=False):
        '''Returns a version of 'text' where HTML codes have been replaced with human-readable text'''
        html = self.read_file('filters/translate html to english.yml')
        for first, second in html:
            text = re.sub(r'(?i)\b%s\b' % first, r'_%s_' % second, text)
            if debug: print '\t', text, ": %s <-> %s" % (first, second)
        return text
    
    def translate_from_internet_lingo(self, text):
        '''Returns a version of 'text' where Intenret lingo has been translated to English'''
        for internet, english in self.read_file('filters/translate from internet lingo.yml'):
            text = re.sub(r'(?i)\b%s\b' % internet, r'_%s_' % english, text)
        return re.sub(r'\b_(.*?)_\b', r'\1', text)

    def fix_spelling_mistakes(self, text):
        '''Returns a version of 'text' where spelling mistakes have been corrected'''
        for key, value in self.read_file('filters/fix spelling mistakes.yml'):
            text = re.sub(r'(?i)\b%s\b' % key, r'_%s_' % value, text)
        return re.sub(r'\b_(.*?)_\b', r'\1', text)

    def swap_persons(self, text, debug=False):
        '''Returns a version of 'text' where persons have been swapped. For example: "I" with "you"'''
        if debug: print "swap_persons():"
        
        for first, second in self.read_file('filters/swap persons.yml'):
            text = re.sub(r'(?i)\b%s\b' % first, r'_%s_' % second, text)
            text = re.sub(r'(?i)\b%s\b' % second, r'_%s_' % first, text)
            if debug: print '\t', text, ": %s <-> %s" % (first, second)
        text = re.sub(r'\b_(.*?)_\b', r'\1', text)
        
        # unswap 'not much, you?' or it will become 'not much, me?'
        return re.sub(r'([\.,] ?)(me|I)( *\?*)$', r'\1you\3', text)
    
    def clean_sentence(self, sentence, debug=False):
        '''Returns a cleaned up version of the string 'sentence' as a list of (word, tag) tuples'''
        # clean it up first
        if debug: print "original input:", sentence
        sentence = self.translate_from_html_to_english(sentence, debug)
        if debug: print "translate_from_html_to_english():", sentence
        sentence = self.translate_from_internet_lingo(sentence)
        if debug: print "translate_from_internet_lingo():", sentence
        sentence = self.fix_spelling_mistakes(sentence)
        if debug: print "fix_spelling_mistakes():", sentence
        # sentence = self.swap_persons(sentence, debug)
        # if debug: print "\n\tSwapped:", sentence, '\n'
        sentence = self.ml.tokenize(sentence, 1)
        if debug: print "ml.tokenize():", sentence
        return sentence
    
    def update_graph(self, graph, node1, node2):
        '''Increases edge weight between node1 and node2'''
        
        if not self.struct.has_edge(node1, node2):
            print "Old edge count = 0"
            graph.add_edge(node1, node2, 1)
        else:
            count = graph.get_edge(node1, node2)
            print "Old edge count ="
            print count
            count = count + 1
            graph.add_edge(node1, node2, count)
        print "New edge count from"
        print node1
        print "to"
        print node2
        print "is"
        print graph.get_edge(node1, node2)
    
    def get_edge_weight(self, graph, node1, node2):
        '''Returns the weight of the edge from node1 to node2'''
        
        if not graph.has_edge(node1, node2):
            return 0
        else:
            count = graph.get_edge(node1, node2)
            return 2*((1/(1 + math.e**(-(count/10)))) - .5)
        
    def get_weighted_choice(self, graph, node):
        choices = graph.successors(node)
        
        sum = 0
        for current in choices:
            sum = sum + self.get_edge_weight(graph, node, current)
        
        selector = uniform(0, sum)
        for current in choices:
            selector = selector - self.get_edge_weight(graph, node, current)
            if (selector <= 0): return current
        
        return choices[len(choices) - 1]
        
        
        
        
    
    def parse_incoming_message(self, sentence, sender, debug=False):
        
        sender = sender.replace(' ', '').lower()
        if debug: print '-' * 79
        
        # update the timestamp
        prev_time = time.time() - 3400 if not self.current_timestamp.has_key(sender) else self.current_timestamp[sender]
        self.current_timestamp[sender] = time.time()
        self.elapsed_timestamp[sender] = self.current_timestamp[sender] - prev_time
        if debug: print "Seconds after previous message:", self.elapsed_timestamp[sender]
        
        # tag the sentence for its parts of speech
        sentence = self.clean_sentence(sentence, debug)
        tagged_sentence = self.ml.tag_tokenized(sentence).split(' ')
        if debug: print "tagged:", tagged_sentence, '\n'
        
        # reformat the tagged sentence: "word/part-of-speech" becomes (word, part-of-speech)
        sentence = []
        for index, pair in enumerate(tagged_sentence):
            (word, tag) = pair.split(r'/')
            word = self.ml.theMontyLemmatiser.lemmatise_word(word).lower()
            sentence.append((word, tag, self.elapsed_timestamp[sender]))
        
        # end every sentence with punctuation
        if ((sentence[-1][i_tag] != '.') and (sentence[-1][i_word] != '@')):
            sentence.append(('.', '.', self.elapsed_timestamp[sender]))
        
        if debug: print "Adding %s to the graph" % sentence
        
        # update the senence structure progression graph and lastStruct
        if (sentence[-1][i_word] == '.'):
            self.update_graph(self.struct, self.lastStruct, "declarative")
            self.lastStruct = "declarative"
        elif (sentence[-1][i_word] == '?'):
            self.update_graph(self.struct, self.lastStruct, "interrogative")
            self.lastStruct = "interrogative"
        elif (sentence[-1][i_word] == '!'):
            self.update_graph(self.struct, self.lastStruct, "exclamatory")
            self.lastStruct = "exclamatory"
        else:
            self.update_graph(self.struct, self.lastStruct, "imperative")
            self.lastStruct = "imperative"
        
        
        
        # for every word in the sentence
        for index1, pair1 in enumerate(sentence):
            # verify pair1 exists
            if (not self.g.has_node(pair1)):
                self.g.add_node(pair1)
                print "Adding new node: (%s, %s, %s)" % (pair1[0], pair1[1], pair1[2])
            
            # compare it to every other word in the sentence
            for index2, pair2 in enumerate(sentence):
                # don't compare a word to itself
                if index1 == index2: continue
                
                # verify pair2 exists
                if (not self.g.has_node(pair2)):
                    self.g.add_node(pair2)
                
                # keep the words in order
                (left, right) = (pair1, pair2) if index1 < index2 else (pair2, pair1)
                (left_i, right_i) = (index1, index2) if index1 < index2 else (index2, index1)
                
                # update the edge data
                if not self.g.has_edge(left, right):
                    data = dict()
                    data[abs(left_i - right_i)] = 1
                    self.g.add_edge(left, right, data)
                else:
                    # get any existing edge information
                    data = self.g.get_edge(left, right)[0]
                    
                    # if it doesn't exist yet, create it
                    if type(data) != dict:
                        data = dict()
                    
                    # store new distance in dict
                    data[abs(left_i - right_i)] = 1
                    self.g.remove_edge(left, right)
                    self.g.add_edge(left, right, data)
        
        # save changes to disk
        nx.write_gpickle(self.g, r"Brain\brain.pickle")
        
        # return the sentence in tokenized form (needed to generate the reply)
        return sentence
    
    def debug_starting_words(self, elapsed_time):
        print '\nFinding a word to start a new sentence (%s seconds after previous message)' % elapsed_time
        for choice_index, choice_pair in enumerate(self.g.nodes()):
            usable = '+' if not self.g.predecessors(choice_pair) and ((1-max(elapsed_time, choice_pair[i_time])/(elapsed_time + choice_pair[i_time]))**2 * 100 * (min(elapsed_time, choice_pair[i_time])/(elapsed_time + choice_pair[i_time])) > 1) else '-'
            if usable or len(self.g.nodes()) < 10:
                print "\t\t\t", '%s "%s"' % (usable, choice_pair[i_word])
    
    def debug_word_compatability(self, reply, reply_index, reply_pair, choices, elapsed_time):
        print '\n\t\t    Verifying compatability with word #%d %s' % (reply_index, reply_pair)
        for choice_index, choice_pair in enumerate(choices):
            if self.g.has_edge(reply_pair, choice_pair):
                usable = '+' if self.g.get_edge(reply_pair, choice_pair)[0].has_key(abs(reply_index - len(reply))) and ((1-max(elapsed_time, choice_pair[i_time])/(elapsed_time + choice_pair[i_time]))**2 * 100 * (min(elapsed_time, choice_pair[i_time])/(elapsed_time + choice_pair[i_time])) > 1) else '-'
                print "\t\t\t", '%s "%s" %s "%s"' % (usable, reply_pair[i_word], self.g.get_edge(reply_pair, choice_pair)[0].keys(), choice_pair[i_word])
            else:
                print "\t\t\t  ", '"%s" has no connection to "%s"' % (reply_pair[i_word], choice_pair[i_word])
    
    def debug_sentence_and_successors(self, reply, choices):
        print "\n\n\tSentence so far:", [group[i_word] for group in reply]
        print "\t\tUnfiltered:\n\t\t\t", '\n\t\t\t'.join([str(group) for group in choices]), '\n'
    
    def debug_problem_detected(self):
        print '\t!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
        print "\t!PROBLEM DETECTED. ENDING THIS SENTENCE IMMEDIATELY.!"
        print '\t!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'
    
    def create_random_sentence(self, recipient, debug=False):
        if (debug): print 'Creating a random response'
        recipient = recipient.replace(' ', '').lower()
        
        # get the elapsed time between this message and the last one we received
        elapsed_time = self.elapsed_timestamp[recipient]
        
        # choose the sentence structure
        self.lastStruct = self.get_weighted_choice(self.struct, self.lastStruct)
        
        print "Structure to be used is"
        print self.lastStruct
        
        
        # start the sentence with a random word that has started a sentence before
        if debug: self.debug_starting_words(elapsed_time)
        next = choice([pair for pair in self.g.nodes() if not self.g.predecessors(pair) and 
            ((1 - max(elapsed_time, pair[i_time])/(elapsed_time + pair[i_time]))**2 * \
            100 * min(elapsed_time, pair[i_time])/(elapsed_time + pair[i_time]) > 1) ])
        reply = [next]
        
        while (True):
            choices = self.g.successors(reply[-1])
            if debug: self.debug_sentence_and_successors(reply, choices)
            
            # if there's only one possible word left, use it
            if (len(choices) == 1 and choices[0][i_tag] == '.'):
                if debug: print "\tUsing the last available word"
                reply.append(choices.pop(0))
            
            # if there are no more words to add to the sentence, stop working on the sentence
            if (not len(choices)): break
            
            # for every word in our reply so far, filter out the choices for the next word that aren't good matches
            for reply_index, reply_pair in enumerate(reply):                    
                if debug: self.debug_word_compatability(reply, reply_index, reply_pair, choices, elapsed_time)
                choices = [choice_pair for choice_pair in choices if self.g.has_edge(reply_pair, choice_pair) and
                    self.g.get_edge(reply_pair, choice_pair)[0].has_key(abs(reply_index - len(reply))) and
                    ((1-max(elapsed_time, choice_pair[i_time])/(elapsed_time + choice_pair[i_time]))**2 * 100 * (min(elapsed_time, choice_pair[i_time])/(elapsed_time + choice_pair[i_time])) > 1) ]
            
            if debug: print "\n\t\tFiltered:\n\t\t\t", '\n\t\t\t'.join([str(group) for group in choices]), '\n'
            
            # saves the program from crashing, although this condition should never occur in the first place
            if (not len(choices)):
                if (debug): self.debug_problem_detected()
                break
            
            # use a random word from our list of valid choices
            reply.append(choice(choices))
        
        # return the sentence we created as a list of tuples
        return reply
    
    def generate_reply(self, sentence, username, debug=False):
        reply = self.create_random_sentence(username, False)
        
        reply_plaintext = ''
        for group in reply:
            space = '' if group[i_tag] == '.' or group[i_tag] == ',' else ' '
            reply_plaintext += '%s%s' % (space, self.ml.theMontyNLGenerator.conjugate_verb(group[i_word], group[i_tag]))
        
        # remove the period from sentences ending with one (allows for the possibility of questions without question marks)
        reply_plaintext = re.sub(r'\.\s*$', '', reply_plaintext)
        return re.sub(r'^\s+', '', reply_plaintext)

if __name__ == "__main__":
    try:
        bot = SABBrain()
        debug = (len(sys.argv) == 2 and sys.argv[1] == '-debug')
        
        while 1:
            sentence = ''
            try:
                sentence = raw_input('> ')
            except:
                raise
            
            sentence = bot.parse_incoming_message(sentence, 'ThePIant', debug)
            
            
            # if you supply any command-line arguments, a graph will display each cycle
            if (len(sys.argv) > 2):
                # positions for all nodes
                pos = nx.spring_layout(bot.g)

                # nodes
                nx.draw_networkx_nodes(bot.g, pos, node_size=2000)

                # edges
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) == 1], width=1, alpha=1.0, edge_color='r')
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) == 2], width=1, alpha=0.6, edge_color='g')
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) == 3], width=1, alpha=0.5, edge_color='b')
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) == 4], width=1, alpha=0.4)
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) == 5], width=1, alpha=0.3)
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) == 6], width=1, alpha=0.2)
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) == 7], width=1, alpha=0.2, style='dashed')
                nx.draw_networkx_edges(bot.g, pos, edgelist=[(u,v) for (u,v,d) in bot.g.edges(data=True) if len(d.keys()) >  7], width=1, alpha=0.2, style='dotted')

                # labels
                nx.draw_networkx_labels(bot.g, pos, font_size=8, font_family='sans-serif')

                # turn off x and y axes labels
                plot.xticks([])
                plot.yticks([])

                # display graph
                plot.show()
            
            time1 = time.time()
            reply = bot.generate_reply(sentence, 'thepiant', debug)
            print '%79s' % reply
            time2 = time.time()
            print 'Computed in', str(round(time2 - time1, 2)), 'seconds.'
            print '=' * 79
            
    except KeyboardInterrupt:
        print "\nSigArtBot is shutting down...\n"
        sys.exit(0)