From f79c9f286568406b2e01ef701c58f857707075c7 Mon Sep 17 00:00:00 2001
From: ruinivist <ruinivist@proton.me>
Date: Mon, 9 Feb 2026 19:51:43 +0000
Subject: [PATCH] feat: add core func to get predictions for tz

---
 internal/core/engine.go      | 64 ++++++++++++++++++++++++++++++++++++
 internal/core/engine_test.go | 14 ++++++++
 internal/core/model.go       |  6 ++++
 3 files changed, 84 insertions(+)
 create mode 100644 internal/core/engine.go
 create mode 100644 internal/core/engine_test.go
 create mode 100644 internal/core/model.go

diff --git a/internal/core/engine.go b/internal/core/engine.go
new file mode 100644
index 0000000..dc79195
--- /dev/null
+++ b/internal/core/engine.go
@@ -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
+}
diff --git a/internal/core/engine_test.go b/internal/core/engine_test.go
new file mode 100644
index 0000000..e3c14ca
--- /dev/null
+++ b/internal/core/engine_test.go
@@ -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")
+	}
+}
diff --git a/internal/core/model.go b/internal/core/model.go
new file mode 100644
index 0000000..207161d
--- /dev/null
+++ b/internal/core/model.go
@@ -0,0 +1,6 @@
+package core
+
+type Prediction struct {
+	Country     string
+	Probability float64
+}