package digest

import (
	"math"
	"strings"
	"unicode"
)

const (
	bm25K1  = 1.2
	bm25B   = 0.75
	bm25Min = 2
)

// BM25 is a lexical ranker over a fixed set of documents. Terms are the same tokens as
// terms() — lower-case runs of letters and digits, at least bm25Min characters — counted
// with frequency rather than kept as a set.
type BM25 struct {
	nDocs int
	avgDL float64
	df    map[string]int
	toks  []map[string]int
	len   []int
}

// NewBM25 indexes docs with the default BM25 parameters (k1 = 1.2, b = 0.75).
func NewBM25(docs []string) *BM25 {
	b := &BM25{
		nDocs: len(docs),
		df:    map[string]int{},
		toks:  make([]map[string]int, len(docs)),
		len:   make([]int, len(docs)),
	}
	var total int
	for i, doc := range docs {
		counts := docTerms(doc)
		n := 0
		for _, c := range counts {
			n += c
		}
		b.toks[i] = counts
		b.len[i] = n
		total += n
		for t := range counts {
			b.df[t]++
		}
	}
	if len(docs) > 0 {
		b.avgDL = float64(total) / float64(len(docs))
	}
	return b
}

// Scores returns one score per document, in order, each >= 0 and never NaN or Inf. An empty
// query, an empty index or a document without any query term scores 0.
func (b *BM25) Scores(query string) []float64 {
	if b == nil || b.nDocs == 0 {
		return make([]float64, len(b.toks))
	}
	out := make([]float64, len(b.toks))
	qterms := docTerms(query)
	if len(qterms) == 0 {
		return out
	}
	avgDL := b.avgDL
	if avgDL == 0 {
		avgDL = 1
	}
	for i, toks := range b.toks {
		var score float64
		for term := range qterms {
			tf, ok := toks[term]
			if !ok {
				continue
			}
			idf := math.Log(1 + (float64(b.nDocs)-float64(b.df[term])+0.5)/(float64(b.df[term])+0.5))
			denom := float64(tf) + bm25K1*(1-bm25B+bm25B*float64(b.len[i])/avgDL)
			score += idf * (float64(tf) * (bm25K1 + 1)) / denom
		}
		out[i] = score
	}
	return out
}

// docTerms counts occurrences of each BM25 term in s, using the same tokenization as terms().
func docTerms(s string) map[string]int {
	out := map[string]int{}
	var b strings.Builder
	flush := func() {
		t := b.String()
		b.Reset()
		if len(t) < bm25Min {
			return
		}
		out[t]++
	}
	for _, r := range strings.ToLower(s) {
		if unicode.IsLetter(r) || unicode.IsDigit(r) {
			b.WriteRune(r)
		} else {
			flush()
		}
	}
	flush()
	return out
}