diff --git a/internal/core/engine.go b/internal/core/engine.go
index dc79195..8f19577 100644
--- a/internal/core/engine.go
+++ b/internal/core/engine.go
@@ -7,7 +7,7 @@ import (
 	"tz-snipe/internal/geodata"
 )
 
-func getStatsForTZs(tzs []string) ([]Prediction, error) {
+func GetStatsForTZs(tzs []string) ([]Prediction, error) {
 	// assumes tzs at this stage are uniq
 	// 1. get geodata
 	var predictions []Prediction
@@ -25,13 +25,15 @@ func getStatsForTZs(tzs []string) ([]Prediction, error) {
 	// but append handles init in case of slice
 	// while maps are more strict
 	// so here I need a map
-	seenCountry := make(map[string]int)
+	seenCountry := make(map[string]struct{})
+	// another go quirk: structs take up 0 space
+	// hence above
 	var totalPop int64
-	for _, countries := range tzMap {
-		for _, country := range countries {
+	for _, curTz := range tzs {
+		for _, country := range tzMap[curTz] {
 			if _, exists := seenCountry[country.Name]; !exists {
 				totalPop += country.Population
-				seenCountry[country.Name] = 1
+				seenCountry[country.Name] = struct{}{}
 			}
 		}
 	}
diff --git a/internal/core/engine_test.go b/internal/core/engine_test.go
index e3c14ca..0da33f6 100644
--- a/internal/core/engine_test.go
+++ b/internal/core/engine_test.go
@@ -1,9 +1,12 @@
 package core
 
-import "testing"
+import (
+	"math"
+	"testing"
+)
 
 func TestGetStatsForTZs(t *testing.T) {
-	preds, err := getStatsForTZs([]string{"+0100", "+0000"})
+	preds, err := GetStatsForTZs([]string{"+0100", "+0000"})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -12,3 +15,20 @@ func TestGetStatsForTZs(t *testing.T) {
 		t.Fatalf("failed to get predictions")
 	}
 }
+
+func TestPredSum(t *testing.T) {
+	preds, err := GetStatsForTZs([]string{"+0100"})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// sum must be 1
+	sum := 0.0
+	for _, pred := range preds {
+		sum += pred.Probability
+	}
+
+	if math.Abs(sum-1) > 1e-6 {
+		t.Fatalf("Probabilities did not sum to 1. Got %.4f instead", sum)
+	}
+}