package learn

import "sort"

// Bin is one reliability bin: the mean predicted probability, the empirical outcome rate, and the count of
// points that fell in it. It is the row of a reliability diagram.
type Bin struct {
	MeanPredicted float64
	EmpiricalRate float64
	Count         int
}

// Reliability groups predictions into equal width probability bins and returns the mean predicted
// probability and the empirical outcome rate per non empty bin, the raw material for a calibration curve.
func Reliability(probs []float64, y []int, bins int) []Bin {
	if bins < 1 {
		bins = 10
	}
	sumP := make([]float64, bins)
	sumY := make([]int, bins)
	cnt := make([]int, bins)
	for i := range probs {
		b := int(probs[i] * float64(bins))
		if b >= bins {
			b = bins - 1
		}
		if b < 0 {
			b = 0
		}
		sumP[b] += probs[i]
		sumY[b] += y[i]
		cnt[b]++
	}
	out := make([]Bin, 0, bins)
	for b := 0; b < bins; b++ {
		if cnt[b] == 0 {
			continue
		}
		out = append(out, Bin{
			MeanPredicted: sumP[b] / float64(cnt[b]),
			EmpiricalRate: float64(sumY[b]) / float64(cnt[b]),
			Count:         cnt[b],
		})
	}
	return out
}

// ECE is the expected calibration error: the count weighted mean absolute gap between the predicted
// probability and the empirical rate across bins. 0 is perfect calibration.
func ECE(probs []float64, y []int, bins int) float64 {
	n := len(probs)
	if n == 0 {
		return 0
	}
	var ece float64
	for _, b := range Reliability(probs, y, bins) {
		gap := b.MeanPredicted - b.EmpiricalRate
		if gap < 0 {
			gap = -gap
		}
		ece += gap * float64(b.Count) / float64(n)
	}
	return ece
}

// AUC is the rank area under the ROC curve, the probability that a random positive scores above a random
// negative, computed from the Mann Whitney U statistic with a tie correction. Returns 0.5 when either class
// is empty (no discrimination is measurable).
func AUC(probs []float64, y []int) float64 {
	type ps struct {
		p float64
		y int
	}
	n := len(probs)
	if n == 0 {
		return 0.5
	}
	rows := make([]ps, n)
	var pos, neg int
	for i := range probs {
		rows[i] = ps{probs[i], y[i]}
		if y[i] == 1 {
			pos++
		} else {
			neg++
		}
	}
	if pos == 0 || neg == 0 {
		return 0.5
	}
	sort.SliceStable(rows, func(a, b int) bool { return rows[a].p < rows[b].p })
	// assign average ranks (1 based), sharing ranks across ties
	ranks := make([]float64, n)
	for i := 0; i < n; {
		j := i
		for j < n && rows[j].p == rows[i].p {
			j++
		}
		avg := float64(i+1+j) / 2.0 // average of ranks i+1 .. j
		for k := i; k < j; k++ {
			ranks[k] = avg
		}
		i = j
	}
	var sumRankPos float64
	for i := range rows {
		if rows[i].y == 1 {
			sumRankPos += ranks[i]
		}
	}
	u := sumRankPos - float64(pos)*float64(pos+1)/2.0
	return u / (float64(pos) * float64(neg))
}