package eval

import (
	"math"
	"testing"
)

// closeEnough compares two floats to a tight tolerance, since these metrics involve log2 divisions whose exact
// decimal expansion is irrational. The tolerance is far tighter than any band a regression could move.
func closeEnough(a, b float64) bool { return math.Abs(a-b) < 1e-9 }

func rel(ids ...string) map[string]bool {
	m := make(map[string]bool, len(ids))
	for _, id := range ids {
		m[id] = true
	}
	return m
}

// ---------- RecallAtK ----------

func TestRecallAtK(t *testing.T) {
	tests := []struct {
		name     string
		ranked   []string
		relevant map[string]bool
		k        int
		want     float64
	}{
		// Perfect ranking: both relevant ids are in the top two, so recall at any k that reaches them is 1.0.
		{"perfect both in top k", []string{"a", "b", "c"}, rel("a", "b"), 2, 1.0},
		// Reversed ranking: the relevant ids sit at the bottom; a small k misses them entirely.
		{"reversed misses at small k", []string{"c", "d", "a", "b"}, rel("a", "b"), 2, 0.0},
		// Reversed but k large enough to reach both relevant ids.
		{"reversed found at large k", []string{"c", "d", "a", "b"}, rel("a", "b"), 4, 1.0},
		// Partial: one of two relevant ids inside the cutoff.
		{"partial one of two", []string{"a", "x", "y"}, rel("a", "b"), 3, 0.5},
		// No relevant id is present anywhere in the ranked list.
		{"none relevant found", []string{"x", "y", "z"}, rel("a", "b"), 3, 0.0},
		// k larger than the list: capped at the list length, both relevant ids still found.
		{"k larger than list", []string{"a", "b"}, rel("a", "b"), 50, 1.0},
		// Empty ranked list: nothing can be recalled.
		{"empty ranked list", []string{}, rel("a"), 5, 0.0},
		// Empty relevant set is defined as 1.0 (nothing was missed, vacuously).
		{"empty relevant set", []string{"a", "b"}, rel(), 5, 1.0},
		// k of zero recalls nothing.
		{"k zero", []string{"a", "b"}, rel("a"), 0, 0.0},
		// Negative k recalls nothing.
		{"k negative", []string{"a", "b"}, rel("a"), -3, 0.0},
		// Ties / duplicates: a repeated relevant id is counted once, so recall cannot exceed 1.0.
		{"duplicate relevant counted once", []string{"a", "a", "a"}, rel("a", "b"), 3, 0.5},
		// Single relevant id at the cutoff edge.
		{"relevant at k boundary", []string{"x", "a"}, rel("a"), 2, 1.0},
		// Single relevant id just past the cutoff edge.
		{"relevant just past k", []string{"x", "a"}, rel("a"), 1, 0.0},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := RecallAtK(tt.ranked, tt.relevant, tt.k)
			if !closeEnough(got, tt.want) {
				t.Fatalf("RecallAtK(%v, k=%d) = %v, want %v", tt.ranked, tt.k, got, tt.want)
			}
		})
	}
}

// ---------- MRR ----------

func TestMRR(t *testing.T) {
	third := 1.0 / 3.0
	tests := []struct {
		name     string
		ranked   []string
		relevant map[string]bool
		want     float64
	}{
		// First relevant id at rank one scores the maximum.
		{"first relevant at rank 1", []string{"a", "b", "c"}, rel("a"), 1.0},
		// First relevant id at rank two scores one half.
		{"first relevant at rank 2", []string{"x", "a", "c"}, rel("a"), 0.5},
		// First relevant id at rank three scores one third.
		{"first relevant at rank 3", []string{"x", "y", "a"}, rel("a"), third},
		// Multiple relevant ids: only the FIRST one counts, the earliest rank wins.
		{"earliest of several", []string{"x", "b", "a"}, rel("a", "b"), 0.5},
		// No relevant id present anywhere.
		{"none relevant", []string{"x", "y", "z"}, rel("a"), 0.0},
		// Reversed ranking pushes the only relevant id to the bottom.
		{"reversed to bottom", []string{"d", "c", "b", "a"}, rel("a"), 0.25},
		// Empty ranked list.
		{"empty ranked", []string{}, rel("a"), 0.0},
		// Empty relevant set: there is no first relevant id, so MRR is 0.0.
		{"empty relevant", []string{"a", "b"}, rel(), 0.0},
		// Perfect: relevant id leads the list.
		{"perfect lead", []string{"a"}, rel("a"), 1.0},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := MRR(tt.ranked, tt.relevant)
			if !closeEnough(got, tt.want) {
				t.Fatalf("MRR(%v) = %v, want %v", tt.ranked, got, tt.want)
			}
		})
	}
}

// ---------- NDCGAtK ----------

func TestNDCGAtK(t *testing.T) {
	// Hand computed reference values. The discount for the id at one based rank r is 1/log2(r+1).
	// rank 1 discount = 1/log2(2) = 1.0
	// rank 2 discount = 1/log2(3) = 0.6309297535714574
	// rank 3 discount = 1/log2(4) = 0.5
	d1, d2, d3 := 1.0, 1.0/math.Log2(3), 0.5

	tests := []struct {
		name     string
		ranked   []string
		relevant map[string]bool
		k        int
		want     float64
	}{
		// Perfect ranking: relevant ids fill the top, nDCG is 1.0 regardless of their internal order.
		{"perfect packed top", []string{"a", "b", "c"}, rel("a", "b"), 3, 1.0},
		{"perfect internal order irrelevant", []string{"b", "a", "c"}, rel("a", "b"), 3, 1.0},
		// One relevant id at rank one with one irrelevant: ideal is also one relevant at rank one, so 1.0.
		{"single relevant at top", []string{"a", "x"}, rel("a"), 2, 1.0},
		// One relevant id at rank two: dcg = d2, idcg = d1, so nDCG = d2/d1 = d2.
		{"single relevant rank 2", []string{"x", "a"}, rel("a"), 2, d2 / d1},
		// One relevant id at rank three (k high enough): dcg = d3, idcg = d1.
		{"single relevant rank 3", []string{"x", "y", "a"}, rel("a"), 3, d3 / d1},
		// Two relevant ids at ranks one and three: dcg = d1 + d3, idcg = d1 + d2 (two relevant ideal).
		{"two relevant ranks 1 and 3", []string{"a", "x", "b"}, rel("a", "b"), 3, (d1 + d3) / (d1 + d2)},
		// Reversed: relevant ids at the bottom, both inside k.
		{"reversed both in k", []string{"x", "y", "b", "a"}, rel("a", "b"), 4,
			(1.0/math.Log2(4) + 1.0/math.Log2(5)) / (d1 + d2)},
		// No relevant id in the top k.
		{"none in top k", []string{"x", "y", "z"}, rel("a", "b"), 3, 0.0},
		// k larger than the list: only what is present counts; both relevant ids are at the top, so 1.0.
		{"k larger than list", []string{"a", "b"}, rel("a", "b"), 50, 1.0},
		// k smaller than the relevant count: ideal is bounded by k, one relevant fits, so a single relevant
		// id at rank one is perfect for that cutoff.
		{"k below relevant count perfect", []string{"a", "x"}, rel("a", "b"), 1, 1.0},
		// Empty ranked list.
		{"empty ranked", []string{}, rel("a"), 5, 0.0},
		// Empty relevant set is defined as a perfect, vacuous 1.0.
		{"empty relevant", []string{"a", "b"}, rel(), 5, 1.0},
		// k of zero yields 0.0 when there is something relevant to find.
		{"k zero", []string{"a"}, rel("a"), 0, 0.0},
		// Negative k yields 0.0.
		{"k negative", []string{"a"}, rel("a"), -2, 0.0},
		// Ties / duplicates: a repeated relevant id is credited only on its first appearance, so a second
		// copy at rank two adds nothing; dcg = d1, idcg = d1 (one distinct relevant), nDCG = 1.0.
		{"duplicate credited once", []string{"a", "a", "x"}, rel("a"), 3, 1.0},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := NDCGAtK(tt.ranked, tt.relevant, tt.k)
			if !closeEnough(got, tt.want) {
				t.Fatalf("NDCGAtK(%v, k=%d) = %v, want %v", tt.ranked, tt.k, got, tt.want)
			}
		})
	}
}

// TestNDCGBoundedZeroOne proves nDCG never escapes the unit interval across a sweep of rankings, a structural
// invariant a regression in the formula would violate.
func TestNDCGBoundedZeroOne(t *testing.T) {
	relevant := rel("a", "b", "c")
	rankings := [][]string{
		{"a", "b", "c", "x", "y"},
		{"x", "a", "y", "b", "c"},
		{"y", "x", "c", "b", "a"},
		{"a", "x", "b", "y", "c"},
		{"x", "y", "z", "a", "b"},
	}
	for _, r := range rankings {
		for k := 1; k <= 6; k++ {
			got := NDCGAtK(r, relevant, k)
			if got < 0.0 || got > 1.0+1e-9 {
				t.Fatalf("NDCGAtK(%v, k=%d) = %v escaped [0,1]", r, k, got)
			}
		}
	}
}

// TestCountRelevantIgnoresFalse proves a qrels map that records explicit non relevant judgments (value false)
// does not inflate the relevant denominator, so recall and nDCG stay correct against such a fixture.
func TestCountRelevantIgnoresFalse(t *testing.T) {
	m := map[string]bool{"a": true, "b": false, "c": true}
	// Two relevant (a, c). A ranking that finds only a is recall 0.5, not 0.333 from a count of three.
	got := RecallAtK([]string{"a", "z"}, m, 5)
	if !closeEnough(got, 0.5) {
		t.Fatalf("RecallAtK with explicit false judgments = %v, want 0.5", got)
	}
}