package learn

import (
	"math"
	"testing"
)

func approx(a, b float64) bool { return math.Abs(a-b) < 1e-9 }

func TestEmpiricalHazardBasic(t *testing.T) {
	// Four requirements over a 3 month horizon:
	//   filled at month 1, filled at month 2, censored (still open) at month 2, filled at month 3.
	// At risk: month1 = 4 (all), month2 = 3 (the three that survived month1), month3 = 1 (only the last).
	// Events: month1 = 1, month2 = 1, month3 = 1.
	// Hazard: 1/4, 1/3, 1/1.
	d := []Duration{
		{Periods: 1, Filled: true},
		{Periods: 2, Filled: true},
		{Periods: 2, Filled: false}, // censored at month 2
		{Periods: 3, Filled: true},
	}
	h := EmpiricalHazard(d, 3)
	if !approx(h[0], 0.25) || !approx(h[1], 1.0/3.0) || !approx(h[2], 1.0) {
		t.Fatalf("hazard wrong: %v", h)
	}
}

func TestEmpiricalHazardCensoredNeverAnEvent(t *testing.T) {
	// Every observation is censored (nobody filled): the hazard is zero in every period, the survival stays 1.
	d := []Duration{{Periods: 3, Filled: false}, {Periods: 2, Filled: false}}
	h := EmpiricalHazard(d, 3)
	for _, v := range h {
		if v != 0 {
			t.Fatalf("a censored only set must have zero hazard, got %v", h)
		}
	}
}

func TestEmpiricalHazardClampsMaxPeriod(t *testing.T) {
	// A duration beyond maxPeriod is folded into the last period, never dropped.
	d := []Duration{{Periods: 99, Filled: true}}
	h := EmpiricalHazard(d, 3)
	if len(h) != 3 {
		t.Fatalf("hazard length must be maxPeriod, got %d", len(h))
	}
	if !approx(h[2], 1.0) {
		t.Fatalf("the out of range fill must land in the last period, got %v", h)
	}
}

func TestSurvivalIsCumulativeProduct(t *testing.T) {
	// hazards 0.25, 0.5 => S1 = 0.75, S2 = 0.75*0.5 = 0.375.
	s := Survival([]float64{0.25, 0.5})
	if !approx(s[0], 0.75) || !approx(s[1], 0.375) {
		t.Fatalf("survival wrong: %v", s)
	}
}

func TestSurvivalClampsHazard(t *testing.T) {
	// An out of range hazard is clamped so survival stays a valid probability.
	s := Survival([]float64{1.5, -0.2})
	if s[0] != 0 || s[1] != 0 {
		t.Fatalf("a clamped hazard of 1 must drop survival to 0 and keep it there, got %v", s)
	}
}

func TestFillWithin(t *testing.T) {
	hz := []float64{0.25, 0.25, 0.25}
	// P(fill within 1) = 1 - 0.75 = 0.25.
	if !approx(FillWithin(hz, 1), 0.25) {
		t.Fatalf("fill within 1 wrong: %v", FillWithin(hz, 1))
	}
	// P(fill within 3) = 1 - 0.75^3 = 1 - 0.421875 = 0.578125.
	if !approx(FillWithin(hz, 3), 0.578125) {
		t.Fatalf("fill within 3 wrong: %v", FillWithin(hz, 3))
	}
	// A horizon beyond the vector uses the full vector.
	if !approx(FillWithin(hz, 99), 0.578125) {
		t.Fatalf("fill within a large horizon must use the full vector, got %v", FillWithin(hz, 99))
	}
	// A zero horizon fills nothing.
	if FillWithin(hz, 0) != 0 {
		t.Fatalf("fill within 0 must be 0, got %v", FillWithin(hz, 0))
	}
}