package digest

// P5b task 01: BM25, the lexical ranking used while a long document's embeddings are still
// being built (and inside the hybrid score after).

import (
	"math"
	"strings"
	"testing"
)

func TestBM25RanksTheRareTerm(t *testing.T) {
	docs := []string{
		"the ship sailed from helsinki to tallinn in the morning",
		"the ship stayed in port all day",
		"a zanzibarite sample was found near the harbour",
		"the harbour master logged the ship",
	}
	s := NewBM25(docs).Scores("where was the zanzibarite found")
	if len(s) != len(docs) {
		t.Fatalf("%d scores for %d docs", len(s), len(docs))
	}
	best := 0
	for i := range s {
		if s[i] > s[best] {
			best = i
		}
	}
	if best != 2 {
		t.Errorf("best = doc %d (%v), want doc 2, the only one with the rare term", best, s)
	}
}

func TestBM25Properties(t *testing.T) {
	filler := strings.Repeat("word ", 50)
	docs := []string{
		"alpha " + filler,                     // 0: tf 1
		strings.Repeat("alpha ", 10) + filler, // 1: tf 10, same filler
		"alpha common",                        // 2: tf 1, short
		"common " + filler,                    // 3: no alpha
	}
	idx := NewBM25(docs)
	s := idx.Scores("alpha")
	if !(s[1] > s[0]) {
		t.Errorf("more occurrences must score higher: tf10 %v vs tf1 %v", s[1], s[0])
	}
	if !(s[1] < 10*s[0]) {
		t.Errorf("term frequency must saturate: tf10 %v, 10 × tf1 %v", s[1], 10*s[0])
	}
	if !(s[2] > s[0]) {
		t.Errorf("the same term in a shorter document must score higher: short %v, long %v", s[2], s[0])
	}
	if s[3] != 0 {
		t.Errorf("a document without the term scores %v, want 0", s[3])
	}
	// IDF: a term in every document is worth less than a rare one.
	all := NewBM25([]string{"xx yy", "xx zz", "xx ww"}).Scores("xx")
	rare := NewBM25([]string{"xx yy", "qq zz", "qq ww"}).Scores("xx")
	if !(rare[0] > all[0]) {
		t.Errorf("idf: a term in 1 of 3 docs must outscore one in 3 of 3: %v vs %v", rare[0], all[0])
	}
}

func TestBM25Degenerate(t *testing.T) {
	for _, docs := range [][]string{nil, {}, {""}, {"", "a b"}} {
		s := NewBM25(docs).Scores("a")
		if len(s) != len(docs) {
			t.Errorf("%q: %d scores", docs, len(s))
		}
		for _, v := range s {
			if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
				t.Errorf("%q: score %v", docs, v)
			}
		}
	}
	for _, v := range NewBM25([]string{"a b", "c d"}).Scores("") {
		if v != 0 {
			t.Errorf("empty query scored %v", v)
		}
	}
	// Deterministic.
	a := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
	b := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
	for i := range a {
		if a[i] != b[i] {
			t.Fatal("scores differ between identical runs")
		}
	}
}