predictor/internal/datasets/gefs/source.go

// Package gefs implements datasets.Source for NOAA GEFS (Global Ensemble
// Forecast System) forecasts.
//
// Each ensemble member is treated as its own dataset, selected via
// DatasetID.Subset.Members. The download skeleton (HTTP, idx parsing,
// parallel blit) lives in internal/datasets/grib; this package only
// supplies GEFS-specific URL templating and member resolution.
package gefs

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"go.uber.org/zap"

	"predictor-refactored/internal/datasets"
	"predictor-refactored/internal/datasets/grib"
	"predictor-refactored/internal/weather"
	wgfs "predictor-refactored/internal/weather/gfs"
)

// Source is the GEFS implementation of datasets.Source.
type Source struct {
	Variant  *wgfs.Variant
	Parallel int
	Client   *http.Client
	Log      *zap.Logger
}

// NewSource returns a default Source over variant. If variant is nil,
// GEFS 0.5° 3-hour is used.
func NewSource(variant *wgfs.Variant, log *zap.Logger) *Source {
	if variant == nil {
		variant = wgfs.GEFS0p50_3h
	}
	return &Source{
		Variant:  variant,
		Parallel: 8,
		Client:   &http.Client{Timeout: 2 * time.Minute},
		Log:      log,
	}
}

func (s *Source) ID() string { return s.Variant.ID }

func (s *Source) downloader() *grib.Downloader {
	return &grib.Downloader{
		Variant:  s.Variant,
		URLs:     s.url,
		Parallel: s.Parallel,
		Client:   s.Client,
		Log:      s.Log,
	}
}

// url generates the GEFS URL for (date, runHour, member, step, levelSet).
func (s *Source) url(date string, runHour, member, step int, ls wgfs.LevelSet) string {
	if ls == wgfs.LevelSetB {
		return wgfs.GefsGribURLB(date, runHour, member, step, s.Variant.ResToken)
	}
	return wgfs.GefsGribURL(date, runHour, member, step, s.Variant.ResToken)
}

// LatestEpoch HEAD-checks the control member's final forecast hour.
func (s *Source) LatestEpoch(ctx context.Context) (time.Time, error) {
	now := time.Now().UTC()
	hour := now.Hour() - (now.Hour() % 6)
	current := time.Date(now.Year(), now.Month(), now.Day(), hour, 0, 0, 0, time.UTC)

	client := s.Client
	if client == nil {
		client = &http.Client{Timeout: 2 * time.Minute}
	}
	log := s.Log
	if log == nil {
		log = zap.NewNop()
	}

	for range 8 {
		date := current.Format("20060102")
		url := wgfs.GefsGribURL(date, current.Hour(), 0, s.Variant.MaxHour, s.Variant.ResToken) + ".idx"
		req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
		if err == nil {
			resp, err := client.Do(req)
			if err == nil {
				resp.Body.Close()
				if resp.StatusCode == http.StatusOK {
					log.Info("latest GEFS run discovered",
						zap.Time("run", current),
						zap.String("verified_url", url))
					return current, nil
				}
			}
		}
		current = current.Add(-6 * time.Hour)
	}
	return time.Time{}, fmt.Errorf("no recent GEFS run found")
}

// Coverage returns the extent of id.
func (s *Source) Coverage(id datasets.DatasetID) datasets.Coverage {
	v := s.Variant
	cov := datasets.Coverage{
		Region:    datasets.Region{MinLat: -90, MaxLat: 90, MinLng: 0, MaxLng: 360},
		StartTime: id.Epoch,
		EndTime:   id.Epoch.Add(time.Duration(v.MaxHour) * time.Hour),
	}
	if r := id.Subset.Region; r != nil {
		cov.Region = *r
	}
	if h := id.Subset.HourRange; h != nil {
		cov.StartTime = id.Epoch.Add(time.Duration(h.MinHour) * time.Hour)
		cov.EndTime = id.Epoch.Add(time.Duration(h.MaxHour) * time.Hour)
	}
	return cov
}

// Open loads a stored GEFS dataset as a WindField.
func (s *Source) Open(_ context.Context, id datasets.DatasetID, store datasets.Storage) (weather.WindField, error) {
	if !store.Exists(id) {
		return nil, fmt.Errorf("dataset %s not found", id.Filename())
	}
	file, err := wgfs.Open(store.Path(id), s.Variant, id.Epoch.UTC())
	if err != nil {
		return nil, err
	}
	return wgfs.NewWind(file), nil
}

// memberOf extracts the single member index encoded by id.Subset.Members.
func memberOf(id datasets.DatasetID) (int, error) {
	if len(id.Subset.Members) != 1 {
		return 0, fmt.Errorf("gefs dataset id must specify exactly one member (got %v)", id.Subset.Members)
	}
	m := id.Subset.Members[0]
	if m < 0 || m >= wgfs.GEFSMembers {
		return 0, fmt.Errorf("gefs member %d out of range [0, %d)", m, wgfs.GEFSMembers)
	}
	return m, nil
}

// Download fetches one ensemble member's dataset.
func (s *Source) Download(ctx context.Context, id datasets.DatasetID, store datasets.Storage, prog datasets.ProgressSink, throttle datasets.Throttle) error {
	member, err := memberOf(id)
	if err != nil {
		return err
	}
	return s.downloader().Run(ctx, id, member, store, prog, throttle)
}