package gfs

import (
	"time"

	"predictor-refactored/internal/numerics"
	"predictor-refactored/internal/weather"
)

// Wind is a WindField backed by a GFS dataset file.
//
// The cube is addressed in flat element units with fixed strides so the
// sampler can compute the eight horizontal interpolation corners once and
// reach any (level, variable) by adding constant strides — avoiding the
// five-multiply offset computation per corner per evaluation.
type Wind struct {
	file *File

	hourAxis numerics.Axis
	latAxis  numerics.Axis
	lngAxis  numerics.Axis

	hourStride  int64 // elements between successive hours
	levelStride int64 // elements between successive pressure levels
	varStride   int64 // elements between successive variables
	latStride   int64 // elements between successive latitudes
}

// NewWind returns a Wind backed by file. Axes and strides are derived from
// the file's variant geometry.
func NewWind(file *File) *Wind {
	v := file.variant
	nLat := v.NumLatitudes()
	nLng := v.NumLongitudes()
	nLev := v.NumLevels()
	return &Wind{
		file:        file,
		hourAxis:    numerics.Axis{Left: 0, Step: float64(v.HourStep), N: v.NumHours(), Name: "hour"},
		latAxis:     numerics.Axis{Left: LatStart, Step: v.Resolution, N: nLat, Name: "lat"},
		lngAxis:     numerics.Axis{Left: LonStart, Step: v.Resolution, N: nLng, Wrap: true, Name: "lng"},
		hourStride:  int64(nLev) * NumVariables * int64(nLat) * int64(nLng),
		levelStride: NumVariables * int64(nLat) * int64(nLng),
		varStride:   int64(nLat) * int64(nLng),
		latStride:   int64(nLng),
	}
}

// Epoch returns the forecast run time of the underlying file.
func (w *Wind) Epoch() time.Time { return w.file.Epoch }

// Source returns the variant ID (e.g. "gfs-0p50-3h").
func (w *Wind) Source() string { return w.file.variant.ID }

// Close releases the underlying file's resources.
func (w *Wind) Close() error { return w.file.Close() }

// Wind samples the field at the given UNIX time, geographic coordinate, and
// altitude. Vertical interpolation matches Tawhiri: locate the two pressure
// levels whose interpolated geopotential heights bracket alt, then linearly
// interpolate U and V between them.
func (w *Wind) Wind(t, lat, lng, alt float64) (weather.Sample, error) {
	hours := (t - float64(w.file.Epoch.Unix())) / 3600.0

	bh, err := w.hourAxis.Locate(hours)
	if err != nil {
		return weather.Sample{}, err
	}
	bla, err := w.latAxis.Locate(lat)
	if err != nil {
		return weather.Sample{}, err
	}
	bln, err := w.lngAxis.Locate(lng)
	if err != nil {
		return weather.Sample{}, err
	}

	weights := numerics.TrilinearWeights([3]numerics.Bracket{bh, bla, bln})

	// Flat element index of each of the eight horizontal corners, at level 0
	// variable 0, in the canonical TrilinearWeights order (hour outer, lng
	// inner). Reaching a given (level, variable) corner only adds constant
	// strides.
	var base [8]int64
	hours2 := [2]int64{int64(bh.Lo) * w.hourStride, int64(bh.Hi) * w.hourStride}
	lats2 := [2]int64{int64(bla.Lo) * w.latStride, int64(bla.Hi) * w.latStride}
	lngs2 := [2]int64{int64(bln.Lo), int64(bln.Hi)}
	i := 0
	for _, h := range hours2 {
		for _, la := range lats2 {
			for _, ln := range lngs2 {
				base[i] = h + la + ln
				i++
			}
		}
	}

	sample := func(level int, varIdx int64) float64 {
		off := int64(level)*w.levelStride + varIdx*w.varStride
		var vals [8]float64
		for k := range 8 {
			vals[k] = float64(w.file.ValByElem(base[k] + off))
		}
		return numerics.Dot8(&weights, &vals)
	}

	// Largest pressure level whose interpolated geopotential height is below alt.
	levelIdx := numerics.Bisect(0, w.file.variant.NumLevels()-2, alt, func(level int) float64 {
		return sample(level, VarHeight)
	})

	lowerHGT := sample(levelIdx, VarHeight)
	upperHGT := sample(levelIdx+1, VarHeight)

	altFrac := 0.5
	if lowerHGT != upperHGT {
		altFrac = (upperHGT - alt) / (upperHGT - lowerHGT)
	}

	lowerU := sample(levelIdx, VarWindU)
	upperU := sample(levelIdx+1, VarWindU)
	lowerV := sample(levelIdx, VarWindV)
	upperV := sample(levelIdx+1, VarWindV)

	return weather.Sample{
		U:          lowerU*altFrac + upperU*(1-altFrac),
		V:          lowerV*altFrac + upperV*(1-altFrac),
		AboveModel: altFrac < 0,
	}, nil
}