From e99442c06780d0c2a3d377afc83410635040fa3f Mon Sep 17 00:00:00 2001
From: Aaron Claydon <claydon.aaron@gmail.com>
Date: Fri, 8 Nov 2024 13:14:39 +0000
Subject: [PATCH] Country stats

---
 pkg/stats/calculator/operators.go        |  2 ++
 pkg/stats/calculator/realtimejourneys.go |  2 ++
 pkg/stats/calculator/servicealerts.go    |  6 +++++-
 pkg/stats/calculator/services.go         |  2 ++
 pkg/stats/calculator/stops.go            |  2 ++
 pkg/stats/calculator/util.go             | 15 +++++++++++++++
 6 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/pkg/stats/calculator/operators.go b/pkg/stats/calculator/operators.go
index b12277c..4d4f637 100644
--- a/pkg/stats/calculator/operators.go
+++ b/pkg/stats/calculator/operators.go
@@ -11,6 +11,7 @@ type OperatorsStats struct {
 	Total int
 
 	Datasources map[string]int
+	Countries   map[string]int
 }
 
 func GetOperators() OperatorsStats {
@@ -20,6 +21,7 @@ func GetOperators() OperatorsStats {
 	stats.Total = int(numberoperators)
 
 	stats.Datasources = CountAggregate(operatorsCollection, "$datasource.datasetid")
+	stats.Countries = CountCountries(stats.Datasources)
 
 	return stats
 }
diff --git a/pkg/stats/calculator/realtimejourneys.go b/pkg/stats/calculator/realtimejourneys.go
index ca98a9e..d96bf92 100644
--- a/pkg/stats/calculator/realtimejourneys.go
+++ b/pkg/stats/calculator/realtimejourneys.go
@@ -21,6 +21,7 @@ type RealtimeJourneyStats struct {
 	TransportTypes map[ctdf.TransportType]int
 	Features       map[string]int
 	Datasources    map[string]int
+	Countries      map[string]int
 }
 
 func GetRealtimeJourneys() RealtimeJourneyStats {
@@ -138,6 +139,7 @@ func GetRealtimeJourneys() RealtimeJourneyStats {
 		TransportTypes:       transportTypes,
 		Features:             features,
 		Datasources:          datasources,
+		Countries:            CountCountries(datasources),
 		NotActivelyTracked:   numberActiveRealtimeJourneysNotActivelyTracked,
 	}
 }
diff --git a/pkg/stats/calculator/servicealerts.go b/pkg/stats/calculator/servicealerts.go
index 594d1cb..62824a1 100644
--- a/pkg/stats/calculator/servicealerts.go
+++ b/pkg/stats/calculator/servicealerts.go
@@ -16,6 +16,7 @@ type ServiceAlertsStats struct {
 	Inactive int
 
 	Datasources map[string]int
+	Countries   map[string]int
 }
 
 func GetServiceAlerts() ServiceAlertsStats {
@@ -40,11 +41,14 @@ func GetServiceAlerts() ServiceAlertsStats {
 		}
 	}
 
+	datasources := CountAggregate(collection, "$datasource.datasetid")
+
 	return ServiceAlertsStats{
 		Total:    numberActiveAlerts + numberInactiveAlerts,
 		Active:   numberActiveAlerts,
 		Inactive: numberInactiveAlerts,
 
-		Datasources: CountAggregate(collection, "$datasource.datasetid"),
+		Datasources: datasources,
+		Countries:   CountCountries(datasources),
 	}
 }
diff --git a/pkg/stats/calculator/services.go b/pkg/stats/calculator/services.go
index 6eb8453..73ecd7d 100644
--- a/pkg/stats/calculator/services.go
+++ b/pkg/stats/calculator/services.go
@@ -12,6 +12,7 @@ type ServicesStats struct {
 
 	TransportTypes map[string]int
 	Datasources    map[string]int
+	Countries      map[string]int
 }
 
 func GetServices() ServicesStats {
@@ -26,6 +27,7 @@ func GetServices() ServicesStats {
 
 	stats.TransportTypes = CountAggregate(servicesCollection, "$transporttype")
 	stats.Datasources = CountAggregate(servicesCollection, "$datasource.datasetid")
+	stats.Countries = CountCountries(stats.Datasources)
 
 	return stats
 }
diff --git a/pkg/stats/calculator/stops.go b/pkg/stats/calculator/stops.go
index 80ff9d0..7fa711f 100644
--- a/pkg/stats/calculator/stops.go
+++ b/pkg/stats/calculator/stops.go
@@ -11,6 +11,7 @@ type StopsStats struct {
 	Total int
 
 	Datasources map[string]int
+	Countries   map[string]int
 }
 
 func GetStops() StopsStats {
@@ -20,6 +21,7 @@ func GetStops() StopsStats {
 	stats.Total = int(numberStops)
 
 	stats.Datasources = CountAggregate(stopsCollection, "$datasource.datasetid")
+	stats.Countries = CountCountries(stats.Datasources)
 
 	return stats
 }
diff --git a/pkg/stats/calculator/util.go b/pkg/stats/calculator/util.go
index 9659ae7..934e9ab 100644
--- a/pkg/stats/calculator/util.go
+++ b/pkg/stats/calculator/util.go
@@ -2,6 +2,7 @@ package calculator
 
 import (
 	"context"
+	"strings"
 
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -30,3 +31,17 @@ func CountAggregate(collection *mongo.Collection, aggregateKey string) map[strin
 
 	return countMap
 }
+
+func CountCountries(datasources map[string]int) map[string]int {
+	countries := map[string]int{}
+
+	for datasource, count := range datasources {
+		// Making a big assumption here
+		datasourceSplit := strings.Split(datasource, "-")
+		country := datasourceSplit[0]
+
+		countries[country] = count
+	}
+
+	return countries
+}