discord: add cadeybot

5 files changed, 293 insertions, 0 deletions

diff --git a/discord/cadeybot/.gitignore b/discord/cadeybot/.gitignore
new file mode 100644
index 0000000..a029fb7
--- /dev/null
+++ b/discord/cadeybot/.gitignore
@@ -0,0 +1,3 @@
+.env
+cadeybot
+*.gob
diff --git a/discord/cadeybot/corpusmake.sh b/discord/cadeybot/corpusmake.sh
new file mode 100755
index 0000000..884f3d0
--- /dev/null
+++ b/discord/cadeybot/corpusmake.sh
@@ -0,0 +1,9 @@
+#!/bin/sh
+
+set -e
+set -x
+
+for dir in $(find -type d)
+do
+    [ "$dir" != . ] && ./importer ./"$dir"/messages.csv >> brain.txt
+done
diff --git a/discord/cadeybot/importer/main.go b/discord/cadeybot/importer/main.go
new file mode 100644
index 0000000..bb38417
--- /dev/null
+++ b/discord/cadeybot/importer/main.go
@@ -0,0 +1,42 @@
+package main
+
+import (
+	"encoding/csv"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+)
+
+func main() {
+	flag.Parse()
+
+	if flag.NArg() != 1 {
+		log.Fatal("usage: importer <messages.csv>")
+	}
+
+	fname := flag.Arg(0)
+	if fname == "" {
+		log.Fatal("usage: importer <messages.csv>")
+	}
+
+	fin, err := os.Open(fname)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	csvReader := csv.NewReader(fin)
+	_, err = csvReader.Read() // ignore the first row, it's the index
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	all, err := csvReader.ReadAll()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	for _, row := range all {
+		fmt.Println(row[2])
+	}
+}
diff --git a/discord/cadeybot/main.go b/discord/cadeybot/main.go
new file mode 100644
index 0000000..6b18fd8
--- /dev/null
+++ b/discord/cadeybot/main.go
@@ -0,0 +1,102 @@
+package main
+
+import (
+	"bufio"
+	"flag"
+	"log"
+	"math/rand"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/bwmarrin/discordgo"
+	"github.com/facebookgo/flagenv"
+	_ "github.com/joho/godotenv"
+)
+
+var (
+	token      = flag.String("token", "", "discord token")
+	brainInput = flag.String("brain", "", "brain file")
+)
+
+func main() {
+	flagenv.Parse()
+	flag.Parse()
+
+	chain := NewChain(3)
+
+	if *brainInput != "" {
+		log.Printf("Opening %s...", *brainInput)
+
+		fin, err := os.Open(*brainInput)
+		if err != nil {
+			panic(err)
+		}
+
+		s := bufio.NewScanner(fin)
+		for s.Scan() {
+			t := s.Text()
+
+			_, err := chain.Write(t)
+			if err != nil {
+				panic(err)
+			}
+		}
+
+		err = chain.Save("cadey.gob")
+		if err != nil {
+			panic(err)
+		}
+	} else {
+		err := chain.Load("cadey.gob")
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	rand.Seed(time.Now().Unix())
+
+	mc := func(s *discordgo.Session, m *discordgo.MessageCreate) {
+		// Ignore all messages created by the bot itself
+		// This isn't required in this specific example but it's a good practice.
+		if m.Author.ID == s.State.User.ID {
+			return
+		}
+
+		mentionsMe := false
+		for _, us := range m.Mentions {
+			if us.ID == s.State.User.ID {
+				mentionsMe = true
+				break
+			}
+		}
+
+		if !mentionsMe {
+			return
+		}
+
+		s.ChannelMessageSend(m.ChannelID, chain.Generate(15))
+	}
+
+	if *token == "" {
+		log.Fatal("set -token or TOKEN")
+	}
+
+	dg, err := discordgo.New("Bot " + *token)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	dg.AddHandler(mc)
+
+	err = dg.Open()
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer dg.Close()
+
+	sc := make(chan os.Signal, 1)
+	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
+	<-sc
+}
diff --git a/discord/cadeybot/markov.go b/discord/cadeybot/markov.go
new file mode 100644
index 0000000..f31c133
--- /dev/null
+++ b/discord/cadeybot/markov.go
@@ -0,0 +1,137 @@
+package main
+
+// This Markov chain code is taken from the "Generating arbitrary text"
+// codewalk: http://golang.org/doc/codewalk/markov/
+//
+// Minor modifications have been made to make it easier to integrate
+// with a webserver and to save/load state
+
+import (
+	"encoding/gob"
+	"fmt"
+	"math/rand"
+	"os"
+	"strings"
+	"sync"
+)
+
+// Prefix is a Markov chain prefix of one or more words.
+type Prefix []string
+
+// String returns the Prefix as a string (for use as a map key).
+func (p Prefix) String() string {
+	return strings.Join(p, " ")
+}
+
+// Shift removes the first word from the Prefix and appends the given word.
+func (p Prefix) Shift(word string) {
+	copy(p, p[1:])
+	p[len(p)-1] = word
+}
+
+// Chain contains a map ("chain") of prefixes to a list of suffixes.
+// A prefix is a string of prefixLen words joined with spaces.
+// A suffix is a single word. A prefix can have multiple suffixes.
+type Chain struct {
+	Chain     map[string][]string
+	prefixLen int
+	mu        sync.Mutex
+}
+
+// NewChain returns a new Chain with prefixes of prefixLen words.
+func NewChain(prefixLen int) *Chain {
+	return &Chain{
+		Chain:     make(map[string][]string),
+		prefixLen: prefixLen,
+	}
+}
+
+// Write parses the bytes into prefixes and suffixes that are stored in Chain.
+func (c *Chain) Write(in string) (int, error) {
+	sr := strings.NewReader(in)
+	p := make(Prefix, c.prefixLen)
+	for {
+		var s string
+		if _, err := fmt.Fscan(sr, &s); err != nil {
+			break
+		}
+		key := p.String()
+		c.mu.Lock()
+		c.Chain[key] = append(c.Chain[key], s)
+		c.mu.Unlock()
+		p.Shift(s)
+	}
+	return len(in), nil
+}
+
+// Generate returns a string of at most n words generated from Chain.
+func (c *Chain) Generate(n int) string {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	p := make(Prefix, c.prefixLen)
+	var words []string
+	for i := 0; i < n; i++ {
+		choices := c.Chain[p.String()]
+		if len(choices) == 0 {
+			break
+		}
+		next := choices[rand.Intn(len(choices))]
+		words = append(words, next)
+		p.Shift(next)
+	}
+	return strings.Join(words, " ")
+}
+
+// Save the chain to a file
+func (c *Chain) Save(fileName string) error {
+	// Open the file for writing
+	fo, err := os.Create(fileName)
+	if err != nil {
+		return err
+	}
+	// close fo on exit and check for its returned error
+	defer func() {
+		if err := fo.Close(); err != nil {
+			panic(err)
+		}
+	}()
+
+	// Create an encoder and dump to it
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	enc := gob.NewEncoder(fo)
+	err = enc.Encode(c)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Load the chain from a file
+func (c *Chain) Load(fileName string) error {
+	// Open the file for reading
+	fi, err := os.Open(fileName)
+	if err != nil {
+		return err
+	}
+	// close fi on exit and check for its returned error
+	defer func() {
+		if err := fi.Close(); err != nil {
+			panic(err)
+		}
+	}()
+
+	// Create a decoder and read from it
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	dec := gob.NewDecoder(fi)
+	err = dec.Decode(c)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}