<commit_before>package compute

import (
	"fmt"
	"time"

	cache "github.com/patrickmn/go-cache"
	"github.com/underlx/disturbancesmlx/dataobjects"
)

// VehicleETAHandler aggregates and manages VehicleETAs
type VehicleETAHandler struct {
	etas *cache.Cache
}

// NewVehicleETAHandler returns a new VehicleETAHandler
func NewVehicleETAHandler() *VehicleETAHandler {
	return &VehicleETAHandler{
		etas: cache.New(cache.NoExpiration, 1*time.Minute),
	}
}

// RegisterVehicleETA adds a prediction to the system, replacing any previous predictions for the same vehicle
func (h *VehicleETAHandler) RegisterVehicleETA(eta *dataobjects.VehicleETA) {
	h.etas.Set(
		h.cacheKey(eta.Station, eta.Direction, eta.ArrivalOrder),
		eta,
		eta.ValidFor)
}

// VehicleETAs returns the ETAs of the next `numVehicles` arriving at the specified station, in the specified direction
// Returns an empty slice if no ETA is available
func (h *VehicleETAHandler) VehicleETAs(station *dataobjects.Station, direction *dataobjects.Station, numVehicles int) []*dataobjects.VehicleETA {
	result := []*dataobjects.VehicleETA{}
	for i := 1; i < numVehicles+1; i++ {
		etaIface, ok := h.etas.Get(h.cacheKey(station, direction, i))
		if !ok {
			continue
		}

		result = append(result, etaIface.(*dataobjects.VehicleETA))
	}
	return result
}

func (h *VehicleETAHandler) cacheKey(station, direction *dataobjects.Station, arrivalOrder int) string {
	return fmt.Sprintf("%s#%s#%d", station.ID, direction.ID, arrivalOrder)
}

// TrainPositions returns VehicleETAs containing the closest position for each train in the network.
// The map is indexed by VehicleServiceID
func (h *VehicleETAHandler) TrainPositions() map[string]*dataobjects.VehicleETA {
	// m maps VehicleServiceID to VehicleETAs
	m := make(map[string]*dataobjects.VehicleETA)
	for _, itemIface := range h.etas.Items() {
		item := itemIface.Object.(*dataobjects.VehicleETA)
		min, ok := m[item.VehicleServiceID]
		if !ok || item.LiveETA() < min.LiveETA() {
			m[item.VehicleServiceID] = item
		}
	}
	return m
}<commit_msg>Discard old predictions when computing train positions<commit_after>package compute

import (
	"fmt"
	"time"

	cache "github.com/patrickmn/go-cache"
	"github.com/underlx/disturbancesmlx/dataobjects"
)

// VehicleETAHandler aggregates and manages VehicleETAs
type VehicleETAHandler struct {
	etas *cache.Cache
}

// NewVehicleETAHandler returns a new VehicleETAHandler
func NewVehicleETAHandler() *VehicleETAHandler {
	return &VehicleETAHandler{
		etas: cache.New(cache.NoExpiration, 1*time.Minute),
	}
}

// RegisterVehicleETA adds a prediction to the system, replacing any previous predictions for the same vehicle
func (h *VehicleETAHandler) RegisterVehicleETA(eta *dataobjects.VehicleETA) {
	h.etas.Set(
		h.cacheKey(eta.Station, eta.Direction, eta.ArrivalOrder),
		eta,
		eta.ValidFor)
}

// VehicleETAs returns the ETAs of the next `numVehicles` arriving at the specified station, in the specified direction
// Returns an empty slice if no ETA is available
func (h *VehicleETAHandler) VehicleETAs(station *dataobjects.Station, direction *dataobjects.Station, numVehicles int) []*dataobjects.VehicleETA {
	result := []*dataobjects.VehicleETA{}
	for i := 1; i < numVehicles+1; i++ {
		etaIface, ok := h.etas.Get(h.cacheKey(station, direction, i))
		if !ok {
			continue
		}

		result = append(result, etaIface.(*dataobjects.VehicleETA))
	}
	return result
}

func (h *VehicleETAHandler) cacheKey(station, direction *dataobjects.Station, arrivalOrder int) string {
	return fmt.Sprintf("%s#%s#%d", station.ID, direction.ID, arrivalOrder)
}

// TrainPositions returns VehicleETAs containing the closest position for each train in the network.
// The map is indexed by VehicleServiceID
func (h *VehicleETAHandler) TrainPositions() map[string]*dataobjects.VehicleETA {
	// m maps VehicleServiceID to VehicleETAs
	m := make(map[string]*dataobjects.VehicleETA)
	for _, itemIface := range h.etas.Items() {
		item := itemIface.Object.(*dataobjects.VehicleETA)
		if time.Since(item.Computed) > 2*time.Minute {
			continue
		}
		min, ok := m[item.VehicleServiceID]
		if !ok || item.LiveETA() < min.LiveETA() {
			m[item.VehicleServiceID] = item
		}
	}
	return m
}
{'fraction_non_alphanumeric': 0.11161351737267967, 'fraction_numerical': 0.0016658733936220848, 'mean_word_length': 5.021489971346705, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7518270', 'n_tokens_mistral': 1483, 'n_tokens_neox': 1440, 'n_words': 362}