predictor/internal/datasets/subset.go

package datasets

import (
	"fmt"
	"slices"
	"strings"
	"time"
)

// SubsetSpec describes which portion of a dataset to download.
//
// A zero-value SubsetSpec means "the full dataset". The Region and
// HourRange fields independently restrict what is fetched and stored.
type SubsetSpec struct {
	// Region restricts the geographic extent. nil means global.
	Region *Region `json:"region,omitempty"`

	// HourRange restricts the forecast horizon. nil means the source's
	// full horizon (e.g. 0..192h for GFS 0.5°).
	HourRange *HourRange `json:"hour_range,omitempty"`

	// Members restricts ensemble members for sources that support them (GEFS).
	// nil means all available members.
	Members []int `json:"members,omitempty"`
}

// Region is an axis-aligned geographic bounding box.
//
// Longitudes are in [0, 360); a box crossing the antimeridian has
// MinLng > MaxLng.
type Region struct {
	MinLat float64 `json:"min_lat"`
	MaxLat float64 `json:"max_lat"`
	MinLng float64 `json:"min_lng"`
	MaxLng float64 `json:"max_lng"`
}

// HourRange is an inclusive forecast-hour range.
type HourRange struct {
	MinHour int `json:"min_hour"`
	MaxHour int `json:"max_hour"`
}

// IsGlobal reports whether the spec selects the entire dataset.
func (s SubsetSpec) IsGlobal() bool {
	return s.Region == nil && s.HourRange == nil && len(s.Members) == 0
}

// IncludesLatLng reports whether (lat, lng) lies inside the spec's Region,
// or the spec has no Region.
func (s SubsetSpec) IncludesLatLng(lat, lng float64) bool {
	if s.Region == nil {
		return true
	}
	r := s.Region
	if lat < r.MinLat || lat > r.MaxLat {
		return false
	}
	if r.MinLng <= r.MaxLng {
		return lng >= r.MinLng && lng <= r.MaxLng
	}
	// Wraps the antimeridian.
	return lng >= r.MinLng || lng <= r.MaxLng
}

// IncludesHour reports whether the forecast hour is in range.
func (s SubsetSpec) IncludesHour(h int) bool {
	if s.HourRange == nil {
		return true
	}
	return h >= s.HourRange.MinHour && h <= s.HourRange.MaxHour
}

// IncludesMember reports whether the ensemble member is in range.
func (s SubsetSpec) IncludesMember(m int) bool {
	if len(s.Members) == 0 {
		return true
	}
	return slices.Contains(s.Members, m)
}

// Key returns a deterministic short identifier for the spec. The empty
// string represents the global subset.
func (s SubsetSpec) Key() string {
	if s.IsGlobal() {
		return ""
	}
	var b strings.Builder
	if s.Region != nil {
		fmt.Fprintf(&b, "r%g.%g.%g.%g", s.Region.MinLat, s.Region.MaxLat, s.Region.MinLng, s.Region.MaxLng)
	}
	if s.HourRange != nil {
		if b.Len() > 0 {
			b.WriteByte('_')
		}
		fmt.Fprintf(&b, "h%d.%d", s.HourRange.MinHour, s.HourRange.MaxHour)
	}
	if len(s.Members) > 0 {
		if b.Len() > 0 {
			b.WriteByte('_')
		}
		fmt.Fprintf(&b, "m")
		for i, m := range s.Members {
			if i > 0 {
				b.WriteByte('.')
			}
			fmt.Fprintf(&b, "%d", m)
		}
	}
	return b.String()
}

// DatasetID identifies one storable dataset.
type DatasetID struct {
	Epoch  time.Time
	Subset SubsetSpec
}

// Equals reports whether two DatasetIDs refer to the same dataset.
// DatasetID is not comparable with == because SubsetSpec contains slices.
func (id DatasetID) Equals(other DatasetID) bool {
	return id.Epoch.Equal(other.Epoch) && id.Subset.Key() == other.Subset.Key()
}

// Filename returns the canonical filename stem for the dataset. The
// extension is appended by the Storage implementation.
func (id DatasetID) Filename() string {
	stem := id.Epoch.UTC().Format("20060102T150405Z")
	if k := id.Subset.Key(); k != "" {
		return stem + "_" + k
	}
	return stem
}

// Coverage is the spatial and temporal extent of a loaded dataset, used by
// the Manager to select which dataset can serve a given query.
type Coverage struct {
	Region    Region    `json:"region"`
	StartTime time.Time `json:"start_time"`
	EndTime   time.Time `json:"end_time"`
}

// Covers reports whether (t, lat, lng) lies inside the coverage.
func (c Coverage) Covers(t time.Time, lat, lng float64) bool {
	if t.Before(c.StartTime) || t.After(c.EndTime) {
		return false
	}
	r := c.Region
	if lat < r.MinLat || lat > r.MaxLat {
		return false
	}
	if r.MinLng <= r.MaxLng {
		return lng >= r.MinLng && lng <= r.MaxLng
	}
	return lng >= r.MinLng || lng <= r.MaxLng
}