package eeo

import "testing"

// rowsFor builds candidate rows: considered in the group, selected of them selected. A label of
// "unknown" marks the dimension as genuinely missing (the unknown bucket), matching the SQL flags.
func rowsFor(gender, race string, considered, selected int) []candRow {
	out := make([]candRow, 0, considered)
	for i := 0; i < considered; i++ {
		out = append(out, candRow{
			gender: gender, race: race,
			genderUnknown: gender == "unknown", raceUnknown: race == "unknown",
			selected: i < selected,
		})
	}
	return out
}

func findGroup(rows []ImpactRow, group string) (ImpactRow, bool) {
	for _, r := range rows {
		if r.Group == group {
			return r, true
		}
	}
	return ImpactRow{}, false
}

func TestFourFifthsFlagsDisparateRate(t *testing.T) {
	// male 8/10 (rate 0.80), female 4/10 (rate 0.40). Female ratio = 0.40/0.80 = 0.50, below 0.80, so
	// it must be flagged; male is the reference at 1.0. Two known groups, so the comparison is valid.
	var rows []candRow
	rows = append(rows, rowsFor("male", "groupA", 10, 8)...)
	rows = append(rows, rowsFor("female", "groupA", 10, 4)...)

	sec := analyze("selection", rows, 0)
	female, ok := findGroup(sec.ByGender, "female")
	if !ok {
		t.Fatal("female group missing")
	}
	if !female.Flagged || female.ImpactRatio != "0.5000" {
		t.Fatalf("female should be flagged at 0.5000, got ratio %s flagged %v", female.ImpactRatio, female.Flagged)
	}
	male, _ := findGroup(sec.ByGender, "male")
	if male.Flagged || male.ImpactRatio != "1.0000" {
		t.Fatalf("male should be the reference at 1.0000 unflagged, got %s flagged %v", male.ImpactRatio, male.Flagged)
	}
}

func TestFourFifthsDoesNotFlagBalanced(t *testing.T) {
	// male 8/10 (0.80), female 7/10 (0.70). Ratio 0.70/0.80 = 0.875, at or above 0.80, not flagged.
	var rows []candRow
	rows = append(rows, rowsFor("male", "groupA", 10, 8)...)
	rows = append(rows, rowsFor("female", "groupA", 10, 7)...)
	sec := analyze("selection", rows, 0)
	female, _ := findGroup(sec.ByGender, "female")
	if female.Flagged {
		t.Fatalf("female should not be flagged at ratio %s", female.ImpactRatio)
	}
}

func TestUnknownExcludedFromRatioAndCounted(t *testing.T) {
	// Two known groups (so the comparison is valid) plus an unknown bucket. The unknown low rate must
	// not become the reference and the unknown group is never flagged, but it is counted.
	var rows []candRow
	rows = append(rows, rowsFor("male", "groupA", 10, 9)...)   // reference 0.90
	rows = append(rows, rowsFor("female", "groupA", 10, 8)...) // a second known group
	rows = append(rows, rowsFor("unknown", "groupA", 5, 1)...) // unknown, rate 0.20
	sec := analyze("selection", rows, 0)
	if sec.UnknownCount != 5 {
		t.Fatalf("unknown count = %d, want 5", sec.UnknownCount)
	}
	unk, _ := findGroup(sec.ByGender, "unknown")
	if unk.ImpactRatio != "n/a" || unk.Flagged {
		t.Fatalf("unknown group must have ratio n/a and not be flagged, got %s flagged %v", unk.ImpactRatio, unk.Flagged)
	}
	male, _ := findGroup(sec.ByGender, "male")
	if male.ImpactRatio != "1.0000" {
		t.Fatalf("male reference = %s, want 1.0000 (the unknown low rate did not become the reference)", male.ImpactRatio)
	}
}

func TestLoneKnownGroupIsNotAFalsePass(t *testing.T) {
	// One known group plus an unknown bucket. A four fifths comparison needs at least two known groups,
	// so the lone known group must read n/a, never a self referential 1.0000 that masks a disparity.
	var rows []candRow
	rows = append(rows, rowsFor("white", "white", 10, 4)...)     // the only known group, rate 0.40
	rows = append(rows, rowsFor("unknown", "unknown", 10, 9)...) // the high rate cohort is unknown
	sec := analyze("selection", rows, 0)
	white, _ := findGroup(sec.ByRace, "white")
	if white.ImpactRatio != "n/a" || white.Flagged {
		t.Fatalf("a lone known group must be n/a, not a self referential pass, got ratio %s flagged %v", white.ImpactRatio, white.Flagged)
	}
}

func TestRealValueNamedUnknownIsTreatedAsKnown(t *testing.T) {
	// A candidate who self reports the literal value "unknown" is real data, not missing. With explicit
	// flags it is a known comparison group, so it counts toward the unknown bucket count of zero.
	var rows []candRow
	rows = append(rows, candRow{gender: "unknown", race: "white", genderUnknown: false, raceUnknown: false, selected: true})
	rows = append(rows, candRow{gender: "male", race: "white", genderUnknown: false, raceUnknown: false, selected: false})
	sec := analyze("selection", rows, 0)
	if sec.UnknownCount != 0 {
		t.Fatalf("a real value of unknown must not be counted as missing, unknown count = %d, want 0", sec.UnknownCount)
	}
}

func TestIntersectionalAndCutsPresent(t *testing.T) {
	var rows []candRow
	rows = append(rows, rowsFor("male", "white", 10, 8)...)
	rows = append(rows, rowsFor("female", "black", 10, 4)...)
	sec := analyze("selection", rows, 0)
	if len(sec.ByGender) != 2 || len(sec.ByRace) != 2 || len(sec.Intersectional) != 2 {
		t.Fatalf("expected two groups in each cut, got gender %d race %d intersectional %d",
			len(sec.ByGender), len(sec.ByRace), len(sec.Intersectional))
	}
	if _, ok := findGroup(sec.Intersectional, "female / black"); !ok {
		t.Fatal("intersectional cross tab missing female / black")
	}
}