feat: add core func to get predictions for tz
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package core
|
||||
|
||||
type Prediction struct {
|
||||
Country string
|
||||
Probability float64
|
||||
}
|
||||
Reference in New Issue
Block a user