markov.go

package main

import (
	"bufio"
	"fmt"
	"math/rand"
	"os"
	"strings"
	"time"
)

const (
	NPref  = 2
	MaxGen = 10000
	NoWord = ""
)

type Prefix struct {
	items  [NPref]string
	length int
}

func NewPrefix(words ...string) *Prefix {
	if len(words) > NPref {
		panic("Failed to initialize a new prefix: given words exceed capacity")
	}
	pref := &Prefix{}
	for _, w := range words {
		pref.Push(w)
	}
	return pref
}

func (p *Prefix) Push(s string) {
	ix := p.length
	if p.length == NPref {
		for i := 1; i < NPref; i++ {
			p.items[i-1] = p.items[i]
		}
		ix--
	}
	if p.length < NPref {
		p.length++
	}
	p.items[ix] = s
}

func build(scanner *bufio.Scanner, prefixes map[Prefix][]string) {
	pref := NewPrefix(NoWord, NoWord)

	// Building up prefixes
	for scanner.Scan() {
		w := scanner.Text()
		prefixes[*pref] = append(prefixes[*pref], w)
		pref.Push(w)
	}
	pref.Push(NoWord)
	prefixes[*pref] = append(prefixes[*pref], NoWord)
}

func generate(prefixes map[Prefix][]string, nwords int, builder *strings.Builder) {
	// Generate the output
	pref := NewPrefix(NoWord, NoWord)
	for i := 0; i < nwords; i++ {
		if words, ok := prefixes[*pref]; ok {
			w := words[rand.Intn(len(words))]
			if builder.Len() != 0 {
				builder.WriteRune(' ')
			}
			builder.WriteString(w)
			pref.Push(w)
		} else {
			break
		}
	}
}

func main() {

	rand.Seed(time.Now().UTC().UnixNano())

	reader := bufio.NewReader(os.Stdin)
	scanner := bufio.NewScanner(reader)
	scanner.Split(bufio.ScanWords)

	prefixes := make(map[Prefix][]string)
	build(scanner, prefixes)

	buf := &strings.Builder{}
	generate(prefixes, MaxGen, buf)

	fmt.Println(buf)
}