package main

import (
    "math/rand"
    "strings"
    "bytes"
    "io/ioutil"
    "log"
    "strconv"
    "os"
    "fmt"
)

type Link struct {
    target *State
    occur int
}

type State struct {
    word  string
    links []Link
}

var MarkovSet []State

func AddLink(front *State, back *State) {
    // See if front already has a link to back
    for i, v := range front.links {
        if v.target.word == back.word {
            // Increment occurrences
            front.links[i].occur++
            return
        }
    }
    // front does not have a link to back
    front.links = append(front.links, Link{back, 1})
}

// Safe method of adding states to avoid duplicates
func AddState(word string) *State{
    for _, i := range MarkovSet {
        if i.word == word {
            return nil
        }
    }

    MarkovSet = append(MarkovSet, State{word, nil})
    return &MarkovSet[len(MarkovSet)-1]
}

// Returns a State from MarkovSet with word value of word
func StateFromWord(word string) *State {
    for i, v := range MarkovSet {
        if v.word == word {
            return &MarkovSet[i]
        }
    }

    MarkovSet = append(MarkovSet, State{word, nil})
    return &MarkovSet[len(MarkovSet)-1]
}

func GetNextState(current *State) *State {
    // Sum all occurrences
    var sum int
    for _, v := range current.links {
        sum += v.occur
    }
    // Iterate through links
    for _, v := range current.links {
        if rand.Intn(sum) <= v.occur {
            return v.target
        }
    }

    return nil // This should never happen
}

func GenerateSentence() string {
    var temp bytes.Buffer
    current := GetNextState(StateFromWord("FRONT"))
    for current.word != "BACK" {
        temp.WriteString(current.word)
        current = GetNextState(current)
    }
    temp.WriteString(".")
    out := temp.String()
    out = strings.Title(out)
    return out
}

func LoadSet(filename string) {
    file, err := ioutil.ReadFile(filename)
    if err != nil {
        log.Fatal(err)
    }

    if len(file) > 0 {
        // Split file by line
        lines := bytes.Split(file, []byte("\n"))
        // Iterate through each line
        for _, line := range lines {
            // Split line into words
            words := bytes.Split(line, []byte(" "))
            key := words[0]
            stateptr := StateFromWord(string(key))
            // Remove the key
            words = append(words[:0], words[:1]...)
            for index := 0; index < len(words); index += 2 {
                // index = word, index+1 = occurrences
                fmt.Println(words[index])
                occur, err := strconv.Atoi(string(words[index+1]))
                if err != nil {
                    occur = 1
                }
                stateptr.links = append(stateptr.links, Link{
                    StateFromWord(string(words[index])),
                    occur,
                })
            }
        }
    }
}

func SaveSet(filename string) {
    if len(MarkovSet) > 0 {
        file, err := os.Create(filename)
        if err != nil {
            log.Fatal(err)
        }

        for _, curstate := range MarkovSet {
            file.WriteString(curstate.word + " ")
            for _, curlink := range curstate.links {
                file.WriteString(curlink.target.word + " " + strconv.Itoa(curlink.occur) + " ")
            }
            file.WriteString("\n")
        }
        file.Sync()
    }
}