package learn

import (
	"math"
	"testing"
)

// a small learnable fixture: y is 1 when feature 0 is high, so the fitted coefficient on feature 0 must be
// strongly positive and the model must classify the fixture.
func learnableFixture() ([][]float64, []int) {
	x := [][]float64{
		{0.1, 0.5}, {0.2, 0.4}, {0.15, 0.6}, {0.25, 0.5}, // low feature 0, label 0
		{0.9, 0.5}, {0.8, 0.4}, {0.85, 0.6}, {0.95, 0.5}, // high feature 0, label 1
	}
	y := []int{0, 0, 0, 0, 1, 1, 1, 1}
	return x, y
}

func TestFitLogisticLearnsTheSignal(t *testing.T) {
	x, y := learnableFixture()
	m, err := FitLogistic(x, y, 0.0)
	if err != nil {
		t.Fatalf("fit: %v", err)
	}
	if m.Coeffs[0] <= 0 {
		t.Fatalf("the coefficient on the predictive feature must be positive, got %v", m.Coeffs[0])
	}
	// the model classifies the fixture: a low feature 0 row predicts < 0.5, a high one predicts > 0.5
	if p := m.Predict([]float64{0.15, 0.5}); p >= 0.5 {
		t.Errorf("a low signal row must predict below 0.5, got %v", p)
	}
	if p := m.Predict([]float64{0.9, 0.5}); p <= 0.5 {
		t.Errorf("a high signal row must predict above 0.5, got %v", p)
	}
	if m.Iterations < 1 || m.Iterations > 100 {
		t.Errorf("must converge in a bounded iteration count, got %d", m.Iterations)
	}
}

func TestFitLogisticRidgeShrinksCoefficients(t *testing.T) {
	x, y := learnableFixture()
	loose, _ := FitLogistic(x, y, 0.0)
	tight, _ := FitLogistic(x, y, 50.0)
	if math.Abs(tight.Coeffs[0]) >= math.Abs(loose.Coeffs[0]) {
		t.Errorf("a stronger ridge must shrink the coefficient toward zero: loose %v tight %v", loose.Coeffs[0], tight.Coeffs[0])
	}
}

func TestFitLogisticIsDeterministic(t *testing.T) {
	x, y := learnableFixture()
	a, _ := FitLogistic(x, y, 1.0)
	b, _ := FitLogistic(x, y, 1.0)
	if a.Intercept != b.Intercept {
		t.Fatal("fit must be deterministic across runs")
	}
	for i := range a.Coeffs {
		if a.Coeffs[i] != b.Coeffs[i] {
			t.Fatal("fit coefficients must be deterministic across runs")
		}
	}
}

func TestFitLogisticConstantFeatureContributesNothing(t *testing.T) {
	// feature 1 is constant, so its standardized value is 0 and it must not blow up the fit.
	x := [][]float64{{0.1, 7}, {0.9, 7}, {0.2, 7}, {0.8, 7}}
	y := []int{0, 1, 0, 1}
	m, err := FitLogistic(x, y, 0.0)
	if err != nil {
		t.Fatalf("a constant feature must not break the fit: %v", err)
	}
	if math.IsNaN(m.Coeffs[1]) || math.IsInf(m.Coeffs[1], 0) {
		t.Errorf("a constant feature must not produce NaN or Inf, got %v", m.Coeffs[1])
	}
}

func TestFitLogisticEmptyIsError(t *testing.T) {
	if _, err := FitLogistic(nil, nil, 0); err != ErrNoRows {
		t.Fatalf("an empty dataset must return ErrNoRows, got %v", err)
	}
}

func TestPredictBoundedProbability(t *testing.T) {
	x, y := learnableFixture()
	m, _ := FitLogistic(x, y, 0.0)
	for _, f := range [][]float64{{0, 0}, {1, 1}, {-5, 5}, {100, -100}} {
		p := m.Predict(f)
		if p < 0 || p > 1 || math.IsNaN(p) {
			t.Errorf("Predict must return a probability in [0,1], got %v for %v", p, f)
		}
	}
}

func TestFitPlattMonotonicCalibration(t *testing.T) {
	// raw scores that rank the labels: higher score, more likely label 1.
	scores := []float64{-2, -1.5, -1, -0.5, 0.5, 1, 1.5, 2}
	y := []int{0, 0, 0, 0, 1, 1, 1, 1}
	a, b, err := FitPlatt(scores, y)
	if err != nil {
		t.Fatalf("platt fit: %v", err)
	}
	if a <= 0 {
		t.Fatalf("a positive score to label relationship must give a positive slope, got a=%v", a)
	}
	// the calibrated map is monotone increasing in the score
	prev := sigmoid(a*(-3) + b)
	for s := -3.0; s <= 3.0; s += 0.5 {
		cur := sigmoid(a*s + b)
		if cur < prev-1e-9 {
			t.Fatalf("the platt calibration must be monotone increasing, dropped at s=%v", s)
		}
		prev = cur
	}
}