diff --git a/plugins/babbler/babbler.go b/plugins/babbler/babbler.go index 222df25..01d97af 100644 --- a/plugins/babbler/babbler.go +++ b/plugins/babbler/babbler.go @@ -122,11 +122,48 @@ func (p *BabblerPlugin) Message(message msg.Message) bool { p.Bot.SendMessage(message.Channel, "Phew that was tiring.") return true + } else if len(tokens) == 5 && strings.Index(lowercase, "merge babbler") == 0 { + if tokens[3] != "into" { + p.Bot.SendMessage(message.Channel, "try using 'merge babbler [x] into [y]'") + return true + } + + who := tokens[2] + into := tokens[4] + + if who == into { + p.Bot.SendMessage(message.Channel, "Fuck off") + return true + } + + var whoBabbler *babbler + ok := false + if whoBabbler, ok = p.babblers[who]; !ok { + babbler, err := getMarkovChain(p.db, who) + if err == nil { + whoBabbler = babbler + } else { + whoBabbler = newBabbler() + } + } + + if _, ok := p.babblers[into]; !ok { + babbler, err := getMarkovChain(p.db, into) + if err == nil { + p.babblers[into] = babbler + } else { + p.babblers[into] = newBabbler() + } + } + + p.babblers[into].merge(whoBabbler, into, who) + + p.Bot.SendMessage(message.Channel, "mooooiggged") + return true } else { addToMarkovChain(p.babblers[message.User.Name], lowercase) } - return false } @@ -247,3 +284,75 @@ func (p *BabblerPlugin) babble(who string) string { return fmt.Sprintf("could not find babbler: %s", who) } + +func (into *babbler) merge(other *babbler, intoName, otherName string) { + intoID := "<" + intoName + ">" + otherID := "<" + otherName + ">" + + for nodeWord, myNode := range other.lookup { + if nodeWord == otherID { + nodeWord = intoID + } + + //does this nodeWord exist yet? + if _, ok := into.lookup[nodeWord]; !ok { + into.lookup[nodeWord] = &node{ + wordFrequency: myNode.wordFrequency, + arcs: map[string]*arc{}, + } + } else { + into.lookup[nodeWord].wordFrequency += myNode.wordFrequency + } + + for arcWord, myArc := range myNode.arcs { + if arcWord == otherID { + arcWord = intoID + } + + if myArc.next == other.end { + if _, ok := into.lookup[nodeWord].arcs[arcWord]; !ok { + into.lookup[nodeWord].arcs[arcWord] = &arc{ + transitionFrequency: myArc.transitionFrequency, + next: into.end, + } + } else { + into.lookup[nodeWord].arcs[arcWord].transitionFrequency += myArc.transitionFrequency + } + continue + } + + //does the arcWord exist yet? + if _, ok := into.lookup[arcWord]; !ok { + into.lookup[arcWord] = &node{ + wordFrequency: 0, + arcs: map[string]*arc{}, + } + } + + if _, ok := into.lookup[nodeWord].arcs[arcWord]; !ok { + into.lookup[nodeWord].arcs[arcWord] = &arc{ + transitionFrequency: myArc.transitionFrequency, + next: into.lookup[arcWord], + } + } else { + into.lookup[nodeWord].arcs[arcWord].transitionFrequency += myArc.transitionFrequency + } + } + } + + into.start.wordFrequency += other.start.wordFrequency + + for startWord, startArc := range other.start.arcs { + if startWord == otherID { + startWord = intoID + } + if _, ok := into.start.arcs[startWord]; !ok { + into.start.arcs[startWord] = &arc{ + transitionFrequency: startArc.transitionFrequency, + next: into.lookup[startWord], + } + } else { + into.start.arcs[startWord].transitionFrequency += startArc.transitionFrequency + } + } +} diff --git a/plugins/babbler/babbler_test.go b/plugins/babbler/babbler_test.go index 62b517b..5ebc9df 100644 --- a/plugins/babbler/babbler_test.go +++ b/plugins/babbler/babbler_test.go @@ -61,6 +61,37 @@ func TestBabblerBatch(t *testing.T) { assert.Contains(t, mb.Messages[1], "message") } +func TestBabblerMerge(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + c.config.Babbler.DefaultUsers = []string{"seabass"} + assert.NotNil(t, c) + + seabass := makeMessage(" This is a message") + seabass.User = &user.User{Name: "seabass"} + res := c.Message(seabass) + assert.Len(t, c.babblers, 1) + assert.Len(t, mb.Messages, 0) + + seabass.Body = " This is another message" + res = c.Message(seabass) + + seabass.Body = " This is a long message" + res = c.Message(seabass) + + res = c.Message(makeMessage("!merge babbler seabass into seabass2")) + assert.True(t, res) + assert.Len(t, mb.Messages, 1) + assert.Contains(t, mb.Messages[0], "mooooiggged") + + res = c.Message(makeMessage("!seabass2 says")) + assert.True(t, res) + assert.Len(t, mb.Messages, 2) + + assert.Contains(t, mb.Messages[1], " this is") + assert.Contains(t, mb.Messages[1], "message") +} + func TestHelp(t *testing.T) { mb := bot.NewMockBot() c := New(mb)