package recruiting

import (
	"math/rand"
	"sort"
	"testing"

	"github.com/shopspring/decimal"
)

func obsF(days float64) Observation {
	return Observation{Days: decimal.NewFromFloat(days), Filled: true}
}
func obsC(days float64) Observation {
	return Observation{Days: decimal.NewFromFloat(days), Filled: false}
}

// TestKaplanMeierWorkedExample checks the estimator against a hand computed case.
//
// Five filled at 10, 20, 30, 40, 50 and five still open at 100. At each event the risk set is 10, 9, 8,
// 7, 6, so survival steps 0.9, 0.8, 0.7, 0.6, 0.5 and the median is 50. The naive median over the
// filled only is 30, and the twenty day gap between them IS the survivorship bias.
func TestKaplanMeierWorkedExample(t *testing.T) {
	obs := []Observation{obsF(10), obsF(20), obsF(30), obsF(40), obsF(50),
		obsC(100), obsC(100), obsC(100), obsC(100), obsC(100)}
	curve := kaplanMeier(obs)
	want := []string{"0.9", "0.8", "0.7", "0.6", "0.5"}
	if len(curve) != len(want) {
		t.Fatalf("curve has %d steps, want %d", len(curve), len(want))
	}
	for i, w := range want {
		if !decimal.RequireFromString(curve[i].Surviving).Equal(decimal.RequireFromString(w)) {
			t.Errorf("step %d survival = %s, want %s", i, curve[i].Surviving, w)
		}
	}
	m, ok := kmMedian(curve)
	if !ok || !m.Equal(decimal.NewFromInt(50)) {
		t.Errorf("KM median = %v (reached=%v), want 50", m, ok)
	}
	nm, _ := naiveMedian(obs)
	if !nm.Equal(decimal.NewFromInt(30)) {
		t.Errorf("naive median = %v, want 30", nm)
	}
	// The whole point: the naive figure understates by twenty days on this cohort.
	if !m.Sub(nm).Equal(decimal.NewFromInt(20)) {
		t.Errorf("the bias is %v days, want 20", m.Sub(nm))
	}
}

// TestSurvivalIsMonotoneAndBounded over random cohorts. A survival function that rises, or leaves
// [0,1], is not a survival function, and every number derived from it would be meaningless.
func TestSurvivalIsMonotoneAndBounded(t *testing.T) {
	rng := rand.New(rand.NewSource(20260723))
	one := decimal.NewFromInt(1)
	for trial := 0; trial < 20000; trial++ {
		n := rng.Intn(30)
		obs := make([]Observation, n)
		for i := range obs {
			obs[i] = Observation{
				Days:   decimal.NewFromInt(int64(rng.Intn(200))),
				Filled: rng.Intn(2) == 0,
			}
		}
		curve := kaplanMeier(obs)
		prev := one
		for i, p := range curve {
			s := decimal.RequireFromString(p.Surviving)
			if s.IsNegative() || s.GreaterThan(one) {
				t.Fatalf("survival %s outside [0,1] at step %d", s, i)
			}
			if s.GreaterThan(prev) {
				t.Fatalf("survival rose from %s to %s at step %d", prev, s, i)
			}
			if p.AtRisk <= 0 {
				t.Fatalf("step %d reports %d at risk", i, p.AtRisk)
			}
			if p.Filled <= 0 || p.Filled > p.AtRisk {
				t.Fatalf("step %d has %d events with %d at risk", i, p.Filled, p.AtRisk)
			}
			prev = s
		}
	}
}

// TestWithoutCensoringKMEqualsTheEmpiricalMedian. When nothing is still open there is no censoring, so
// the corrected figure must reduce exactly to the ordinary median. A difference here would mean the
// estimator is wrong in the simplest possible case.
func TestWithoutCensoringKMEqualsTheEmpiricalMedian(t *testing.T) {
	rng := rand.New(rand.NewSource(20260724))
	for trial := 0; trial < 5000; trial++ {
		n := rng.Intn(20) + 1
		obs := make([]Observation, n)
		for i := range obs {
			obs[i] = obsF(float64(rng.Intn(100) + 1))
		}
		curve := kaplanMeier(obs)
		km, ok := kmMedian(curve)
		if !ok {
			t.Fatalf("with no censoring the median must always be reached (n=%d)", n)
		}
		nm, _ := naiveMedian(obs)

		// Sorted fill times, to state the expectation exactly.
		var filled []decimal.Decimal
		for _, o := range obs {
			filled = append(filled, o.Days)
		}
		sortDecimals(filled)

		// The two conventions differ for an even count, and both are standard.
		//
		// Kaplan Meier is a step function, so its median is the smallest observed time at which survival
		// has fallen to one half: with no censoring that is the LOWER of the two central values. The
		// arithmetic median averages them. For an odd count there is a single central value and the two
		// agree exactly.
		//
		// An earlier version of this test asserted that the Kaplan Meier median never falls below the
		// arithmetic one. That is false for every even cohort, and the test caught it. The correct
		// statement is the equality below.
		want := filled[(n-1)/2]
		if !km.Equal(want) {
			t.Fatalf("no censoring, n=%d: KM median %s, want the order statistic %s (arithmetic median %s)",
				n, km, want, nm)
		}
		if n%2 == 1 && !km.Equal(nm) {
			t.Fatalf("no censoring, odd n=%d: KM median %s but arithmetic %s, they must agree", n, km, nm)
		}
		if n%2 == 0 && km.GreaterThan(nm) {
			t.Fatalf("even n=%d: KM median %s above the arithmetic median %s, "+
				"the step function cannot exceed the average of the two central values", n, km, nm)
		}
	}
}

// TestMedianNotReachedIsReportedAsAbsent. When fewer than half the cohort has filled there is no median
// yet, and inventing one would be a fabrication. Absent must be distinguishable from zero.
func TestMedianNotReachedIsReportedAsAbsent(t *testing.T) {
	// One filled, nine still open. Survival only ever falls to 0.9.
	obs := []Observation{obsF(5)}
	for i := 0; i < 9; i++ {
		obs = append(obs, obsC(200))
	}
	ttf := timeToFillFrom(obs)
	if ttf.MedianReached {
		t.Error("a median was reported when only one of ten orders has filled")
	}
	if ttf.Median != "" {
		t.Errorf("median = %q, want empty when not reached", ttf.Median)
	}
	// The naive figure, by contrast, cheerfully reports five days, which is the failure being guarded.
	if ttf.NaiveMedian != "5.00" {
		t.Errorf("naive median = %q, want 5.00", ttf.NaiveMedian)
	}
	if ttf.Cohort != 10 || ttf.FilledN != 1 || ttf.OpenN != 9 {
		t.Errorf("cohort accounting wrong: %+v", ttf)
	}
}

// TestCensoringBeyondTheLastEventRaisesTheMedian is the survivorship case that actually occurs: slow
// requisitions are still open, so they are censored at times later than every fill. Adding them must
// never LOWER the reported median, because they are known to be at least that slow.
//
// Stated narrowly on purpose. It is tempting to assert that the Kaplan Meier median always exceeds the
// naive one, and that is FALSE: censoring concentrated early shrinks the risk set, each later event
// drops survival further, and the median can come out below the naive figure. TestEarlyCensoringCanLower
// demonstrates it. Asserting the general inequality would have encoded a wrong belief as a test.
func TestCensoringBeyondTheLastEventRaisesTheMedian(t *testing.T) {
	rng := rand.New(rand.NewSource(20260725))
	for trial := 0; trial < 3000; trial++ {
		n := rng.Intn(10) + 4
		base := make([]Observation, n)
		maxDay := 0
		for i := range base {
			d := rng.Intn(50) + 1
			if d > maxDay {
				maxDay = d
			}
			base[i] = obsF(float64(d))
		}
		before, okB := kmMedian(kaplanMeier(base))
		if !okB {
			t.Fatal("an uncensored cohort must always reach its median")
		}
		// Add censored observations strictly beyond every event.
		withCensor := append(append([]Observation{}, base...),
			obsC(float64(maxDay+10)), obsC(float64(maxDay+20)))
		after, okA := kmMedian(kaplanMeier(withCensor))
		if !okA {
			// Permitted: adding unfilled orders can push the cohort below half filled, in which case
			// there is no median, which is a stronger statement than a larger one.
			continue
		}
		if after.LessThan(before) {
			t.Fatalf("adding orders still open at %d+ lowered the median from %s to %s",
				maxDay, before, after)
		}
	}
}

// TestEarlyCensoringCanLowerTheMedian records the counterexample to the inequality it is tempting to
// assume. This is not a defect; it is the estimator behaving correctly on a cohort where most orders
// left the risk set early, and it is here so nobody later "fixes" the code to satisfy a false law.
func TestEarlyCensoringCanLowerTheMedian(t *testing.T) {
	// Two filled at 10 and 100, eight censored at day 1.
	obs := []Observation{obsF(10), obsF(100)}
	for i := 0; i < 8; i++ {
		obs = append(obs, obsC(1))
	}
	km, ok := kmMedian(kaplanMeier(obs))
	if !ok {
		t.Fatal("the median should be reached here")
	}
	nm, _ := naiveMedian(obs)
	if !km.LessThan(nm) {
		t.Fatalf("expected the KM median %s to fall below the naive %s on this cohort, "+
			"which is why the general inequality is not asserted anywhere", km, nm)
	}
}

// TestTiedEventsAndCensoringsAtTheSameInstant. A censored observation at the same time as a fill is
// counted in the risk set for that fill, the standard convention. Getting it wrong shifts every
// subsequent step.
func TestTiedEventsAndCensorings(t *testing.T) {
	// Two fills and two censorings all at day 10, from a cohort of four.
	obs := []Observation{obsF(10), obsF(10), obsC(10), obsC(10)}
	curve := kaplanMeier(obs)
	if len(curve) != 1 {
		t.Fatalf("%d steps, want 1", len(curve))
	}
	if curve[0].AtRisk != 4 || curve[0].Filled != 2 {
		t.Errorf("step reports %d at risk and %d filled, want 4 and 2", curve[0].AtRisk, curve[0].Filled)
	}
	// 1 - 2/4 = 0.5
	if !decimal.RequireFromString(curve[0].Surviving).Equal(decimal.RequireFromString("0.5")) {
		t.Errorf("survival = %s, want 0.5", curve[0].Surviving)
	}
}

// TestEmptyCohortIsStable. A new client's first login must render a complete, empty report rather than a
// divide by zero or a missing section.
func TestEmptyCohortIsStable(t *testing.T) {
	ttf := timeToFillFrom(nil)
	if ttf.Cohort != 0 || ttf.MedianReached || ttf.Median != "" || ttf.NaiveMedian != "" {
		t.Errorf("empty cohort produced %+v", ttf)
	}
	if ttf.Curve == nil {
		t.Error("the curve must be an empty slice, never nil, so the JSON is [] and not null")
	}
	if _, ok := naiveMedian(nil); ok {
		t.Error("a median was reported for an empty set")
	}
}

// TestFillRateSuppressesSmallCohorts. One filled requisition is a hundred percent fill rate and one lost
// one is zero. Both are arithmetically correct and neither means anything.
func TestFillRateSuppressesSmallCohorts(t *testing.T) {
	small := fillRateFrom(1, 0, 0, 3)
	if !small.Suppressed || small.OfResolved != "" || small.OfDecided != "" {
		t.Errorf("a cohort of one reported a rate: %+v", small)
	}
	if small.Filled != 1 || small.StillOpen != 3 {
		t.Error("suppressing the rate must not suppress the counts")
	}
	// At the floor the rate appears.
	ok := fillRateFrom(3, 1, 1, 2)
	if ok.Suppressed {
		t.Error("a cohort at the floor was suppressed")
	}
	// Three of five resolved is 60 percent; excluding the cancellation it is three of four, 75 percent.
	if ok.OfResolved != "60.00" {
		t.Errorf("rate of resolved = %q, want 60.00", ok.OfResolved)
	}
	if ok.OfDecided != "75.00" {
		t.Errorf("rate of decided = %q, want 75.00", ok.OfDecided)
	}
}

// TestFillRateBothDenominatorsAreReported. The definition genuinely forks on whether a client pulling a
// role counts against the recruiter, and picking one silently is how two reports of the same desk
// disagree with each other.
func TestFillRateBothDenominatorsReported(t *testing.T) {
	f := fillRateFrom(5, 5, 10, 0)
	if f.OfResolved != "25.00" {
		t.Errorf("counting cancellations gives %q, want 25.00", f.OfResolved)
	}
	if f.OfDecided != "50.00" {
		t.Errorf("excluding cancellations gives %q, want 50.00", f.OfDecided)
	}
	// With no decided orders at all the second rate is absent rather than zero.
	g := fillRateFrom(0, 0, 6, 0)
	if g.OfDecided != "" {
		t.Errorf("rate of decided = %q with nothing decided, want absent", g.OfDecided)
	}
	if g.OfResolved != "0.00" {
		t.Errorf("rate of resolved = %q, want a real 0.00", g.OfResolved)
	}
}

// TestAgingBandsCoverEveryAge. The bands must be contiguous and exhaustive, or an open requisition falls
// between two buckets and vanishes from the report entirely.
func TestAgingBandsCoverEveryAge(t *testing.T) {
	for age := 0; age <= 2000; age++ {
		hits := 0
		for _, b := range agingBands {
			if age >= b.Min && (b.Max == -1 || age <= b.Max) {
				hits++
			}
		}
		if hits != 1 {
			t.Fatalf("age %d falls in %d bands, want exactly 1", age, hits)
		}
	}
	// And the totals reconcile: every input lands somewhere.
	ages := []int{0, 3, 7, 13, 14, 29, 30, 59, 60, 500, -4}
	buckets := agingFrom(ages)
	total := 0
	for _, b := range buckets {
		total += b.Count
	}
	if total != len(ages) {
		t.Errorf("bucketed %d of %d ages", total, len(ages))
	}
	// A negative age, which a clock skew could produce, is treated as zero rather than dropped.
	if buckets[0].Count != 3 { // 0, 3 and the clamped -4
		t.Errorf("first band holds %d, want 3 including the clamped negative", buckets[0].Count)
	}
}

func sortDecimals(d []decimal.Decimal) {
	sort.Slice(d, func(i, j int) bool { return d[i].LessThan(d[j]) })
}