feat: add core func to get predictions for tz

This commit is contained in:
2026-02-09 19:51:43 +00:00
parent 2c74a2d2ba
commit f79c9f2865
3 changed files with 84 additions and 0 deletions
+64
View File
@@ -0,0 +1,64 @@
package core
import (
"cmp"
"fmt"
"slices"
"tz-snipe/internal/geodata"
)
func getStatsForTZs(tzs []string) ([]Prediction, error) {
// assumes tzs at this stage are uniq
// 1. get geodata
var predictions []Prediction
tzMap, err := geodata.Fetch()
if err != nil {
return nil, fmt.Errorf("failed to get geodata")
}
// 2. aggregate pop
// summing over countries pop was an afterthough in impl here
// the model from geodata should've been better
// GO QUIRK
// var init of slice and map are both nil
// but append handles init in case of slice
// while maps are more strict
// so here I need a map
seenCountry := make(map[string]int)
var totalPop int64
for _, countries := range tzMap {
for _, country := range countries {
if _, exists := seenCountry[country.Name]; !exists {
totalPop += country.Population
seenCountry[country.Name] = 1
}
}
}
// 3. build predictions
clear(seenCountry)
for _, tz := range tzs {
countries := tzMap[tz]
// for new countries
for _, country := range countries {
if _, existsCountry := seenCountry[country.Name]; existsCountry {
continue
}
// I like this style of avoiding indents, but is it good?
predictions = append(predictions, Prediction{
Country: country.Name,
Probability: float64(country.Population) / float64(totalPop),
})
}
}
// 4. sort
slices.SortFunc(predictions, func(a, b Prediction) int {
return cmp.Compare(b.Probability, a.Probability) // desc
})
return predictions, nil
}
+14
View File
@@ -0,0 +1,14 @@
package core
import "testing"
func TestGetStatsForTZs(t *testing.T) {
preds, err := getStatsForTZs([]string{"+0100", "+0000"})
if err != nil {
t.Fatal(err)
}
if len(preds) == 0 {
t.Fatalf("failed to get predictions")
}
}
+6
View File
@@ -0,0 +1,6 @@
package core
type Prediction struct {
Country string
Probability float64
}