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