dag.go

// Copyright 2016 ego authors
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package gse

import (
	"math"
	"regexp"
	"strings"
)

const (
	// RatioWord ratio words and letters
	RatioWord float32 = 1.5
	// RatioWordFull full ratio words and letters
	RatioWordFull float32 = 1
)

var reEng = regexp.MustCompile(`[[:alnum:]]`)

type route struct {
	frequency float64
	index     int
}

// Find find word in dictionary return word's frequency and existence
func (seg *Segmenter) Find(str string) (float64, string, bool) {
	return seg.Dict.Find([]byte(str))
}

// Value find word in dictionary return word's value
func (seg *Segmenter) Value(str string) (int, int, error) {
	return seg.Dict.Value([]byte(str))
}

// Analyze analyze the token segment info
func (seg *Segmenter) Analyze(text []string) (az []AnalyzeToken) {
	if len(text) <= 0 {
		return
	}

	start := 0
	end := len([]rune(text[0]))
	for k, v := range text {
		if k > 0 {
			start = az[k-1].End
			end = az[k-1].End + len([]rune(v))
		}

		freq, pos, _ := seg.Find(v)
		az = append(az, AnalyzeToken{
			Position: k,
			Start:    start,
			End:      end,
			Text:     v,
			Freq:     freq,
			Pos:      pos,
		})
	}

	return
}

func (seg *Segmenter) getDag(runes []rune) map[int][]int {
	dag := make(map[int][]int)
	n := len(runes)

	var (
		frag []rune
		i    int
	)

	for k := 0; k < n; k++ {
		dag[k] = make([]int, 0)
		i = k
		frag = runes[k : k+1]

		for {
			freq, _, ok := seg.Find(string(frag))
			if !ok {
				break
			}

			if freq > 0 {
				dag[k] = append(dag[k], i)
			}

			i++
			if i >= n {
				break
			}

			frag = runes[k : i+1]
		}

		if len(dag[k]) == 0 {
			dag[k] = append(dag[k], k)
		}
	}

	return dag
}

func (seg *Segmenter) calc(runes []rune) map[int]route {
	dag := seg.getDag(runes)

	n := len(runes)
	rs := make(map[int]route)

	rs[n] = route{frequency: 0.0, index: 0}
	var r route

	logT := math.Log(seg.Dict.totalFrequency)
	for idx := n - 1; idx >= 0; idx-- {
		for _, i := range dag[idx] {
			freq, _, ok := seg.Find(string(runes[idx : i+1]))

			if ok {
				f := math.Log(freq) - logT + rs[i+1].frequency
				r = route{frequency: f, index: i}
			} else {
				f := math.Log(1.0) - logT + rs[i+1].frequency
				r = route{frequency: f, index: i}
			}

			if v, ok := rs[idx]; !ok {
				rs[idx] = r
			} else {
				f := v.frequency == r.frequency && v.index < r.index
				if v.frequency < r.frequency || f {
					rs[idx] = r
				}
			}
		}
	}

	return rs
}

func (seg *Segmenter) hmm(bufString string, buf []rune) (result []string) {

	v, _, ok := seg.Find(bufString)
	if !ok || v == 0 {
		result = append(result, seg.HMMCut(bufString)...)
		return
	}

	for _, elem := range buf {
		result = append(result, string(elem))
	}
	return
}

func (seg *Segmenter) cutDAG(str string) []string {

	mLen := int(float32(len(str))/RatioWord) + 1
	result := make([]string, 0, mLen)

	runes := []rune(str)
	routes := seg.calc(runes)

	var y int
	length := len(runes)
	var buf []rune

	for x := 0; x < length; {
		y = routes[x].index + 1
		frag := runes[x:y]

		if y-x == 1 {
			buf = append(buf, frag...)
		} else {
			if len(buf) > 0 {
				bufString := string(buf)
				if len(buf) == 1 {
					result = append(result, bufString)
				} else {
					result = append(result, seg.hmm(bufString, buf)...)
				}

				buf = make([]rune, 0)
			}

			result = append(result, string(frag))
		}

		x = y
	}

	if len(buf) > 0 {
		bufString := string(buf)

		if len(buf) == 1 {
			result = append(result, bufString)
		} else {
			result = append(result, seg.hmm(bufString, buf)...)
		}
	}

	return result
}

func (seg *Segmenter) cutDAGNoHMM(str string) []string {
	mLen := int(float32(len(str))/RatioWord) + 1
	result := make([]string, 0, mLen)

	runes := []rune(str)
	routes := seg.calc(runes)
	length := len(runes)

	var y int
	var buf []rune

	for x := 0; x < length; {
		y = routes[x].index + 1
		frag := runes[x:y]

		if reEng.MatchString(string(frag)) && len(frag) == 1 {
			buf = append(buf, frag...)
			x = y
			continue
		}

		if len(buf) > 0 {
			result = append(result, string(buf))
			buf = make([]rune, 0)
		}

		result = append(result, string(frag))
		x = y
	}

	if len(buf) > 0 {
		result = append(result, string(buf))
		// buf = make([]rune, 0)
	}

	return result
}

func (seg *Segmenter) cutAll(str string) []string {
	mLen := int(float32(len(str))/RatioWord) + 1
	result := make([]string, 0, mLen)

	runes := []rune(str)
	dag := seg.getDag(runes)
	start := -1
	ks := make([]int, len(dag))

	for k := range dag {
		ks[k] = k
	}

	var l []int
	for k := range ks {
		l = dag[k]

		if len(l) == 1 && k > start {
			result = append(result, string(runes[k:l[0]+1]))
			start = l[0]
			continue
		}

		for _, j := range l {
			if j > k {
				result = append(result, string(runes[k:j+1]))
				start = j
			}
		}
	}

	return result
}

func (seg *Segmenter) cutForSearch(str string, hmm ...bool) []string {

	mLen := int(float32(len(str))/RatioWordFull) + 1
	result := make([]string, 0, mLen)

	ws := seg.Cut(str, hmm...)
	for _, word := range ws {
		runes := []rune(word)
		for _, incr := range []int{2, 3} {
			if len(runes) <= incr {
				continue
			}

			var gram string
			for i := 0; i < len(runes)-incr+1; i++ {
				gram = string(runes[i : i+incr])
				v, _, ok := seg.Find(gram)
				if ok && v > 0 {
					result = append(result, gram)
				}
			}
		}

		result = append(result, word)
	}

	return result
}

// SuggestFreq suggest the words frequency
// returns a suggested frequncy of a word
// cutted into several short words.
func (seg *Segmenter) SuggestFreq(words ...string) float64 {
	frequency := 1.0
	total := seg.Dict.totalFrequency

	if len(words) > 1 {
		for _, word := range words {
			freq, _, ok := seg.Find(word)
			if ok {
				frequency *= freq
			}

			frequency /= total
		}

		frequency, _ = math.Modf(frequency * total)
		wordFreq := 0.0
		freq, _, ok := seg.Find(strings.Join(words, ""))
		if ok {
			wordFreq = freq
		}

		if wordFreq < frequency {
			frequency = wordFreq
		}

		return frequency
	}

	word := words[0]
	for _, segment := range seg.Cut(word, false) {
		freq, _, ok := seg.Find(segment)
		if ok {
			frequency *= freq
		}

		frequency /= total
	}

	frequency, _ = math.Modf(frequency * total)
	frequency += 1.0
	wordFreq := 1.0

	freq, _, ok := seg.Find(word)
	if ok {
		wordFreq = freq
	}

	if wordFreq > frequency {
		frequency = wordFreq
	}

	return frequency
}