predictor/internal/prediction/interpolate.go

package prediction

import (
	"fmt"

	"predictor-refactored/internal/dataset"
)

// Exact port of the reference interpolation logic (interpolate.pyx).
// 4D interpolation: time, latitude, longitude, altitude (via geopotential height).

// lerp1 holds an index and interpolation weight for one axis.
type lerp1 struct {
	index int
	lerp  float64
}

// lerp3 holds indices and a combined weight for the (hour, lat, lon) axes.
type lerp3 struct {
	hour, lat, lng int
	lerp           float64
}

// RangeError indicates a coordinate is outside the dataset bounds.
type RangeError struct {
	Variable string
	Value    float64
}

func (e *RangeError) Error() string {
	return fmt.Sprintf("%s=%f out of range", e.Variable, e.Value)
}

// pick computes interpolation indices and weights for a single axis.
// left: axis start, step: axis spacing, n: number of points, value: query value.
// Returns two lerp1 values (lower and upper bracket).
func pick(left, step float64, n int, value float64, variableName string) ([2]lerp1, error) {
	a := (value - left) / step
	b := int(a) // truncation toward zero, same as Cython <long> cast
	if b < 0 || b >= n-1 {
		return [2]lerp1{}, &RangeError{Variable: variableName, Value: value}
	}
	l := a - float64(b)
	return [2]lerp1{
		{index: b, lerp: 1 - l},
		{index: b + 1, lerp: l},
	}, nil
}

// pick3 computes 8 trilinear interpolation weights for (hour, lat, lng).
func pick3(hour, lat, lng float64) ([8]lerp3, error) {
	lhour, err := pick(0, 3, 65, hour, "hour")
	if err != nil {
		return [8]lerp3{}, err
	}
	llat, err := pick(-90, 0.5, 361, lat, "lat")
	if err != nil {
		return [8]lerp3{}, err
	}
	// Longitude wraps: tell pick the axis is one larger, then wrap index 720 → 0
	llng, err := pick(0, 0.5, 720+1, lng, "lng")
	if err != nil {
		return [8]lerp3{}, err
	}
	if llng[1].index == 720 {
		llng[1].index = 0
	}

	var out [8]lerp3
	i := 0
	for _, a := range lhour {
		for _, b := range llat {
			for _, c := range llng {
				out[i] = lerp3{
					hour: a.index,
					lat:  b.index,
					lng:  c.index,
					lerp: a.lerp * b.lerp * c.lerp,
				}
				i++
			}
		}
	}
	return out, nil
}

// interp3 performs 8-point weighted interpolation at a given variable and pressure level.
func interp3(ds *dataset.File, lerps [8]lerp3, variable, level int) float64 {
	var r float64
	for i := 0; i < 8; i++ {
		v := ds.Val(lerps[i].hour, level, variable, lerps[i].lat, lerps[i].lng)
		r += float64(v) * lerps[i].lerp
	}
	return r
}

// search finds the largest pressure level index where interpolated geopotential
// height is less than the target altitude. Searches levels 0..45 (excludes topmost).
func search(ds *dataset.File, lerps [8]lerp3, target float64) int {
	lower, upper := 0, 45

	for lower < upper {
		mid := (lower + upper + 1) / 2
		test := interp3(ds, lerps, dataset.VarHeight, mid)
		if target <= test {
			upper = mid - 1
		} else {
			lower = mid
		}
	}

	return lower
}

// interp4 performs altitude-interpolated wind lookup using two bracketing levels.
func interp4(ds *dataset.File, lerps [8]lerp3, altLerp lerp1, variable int) float64 {
	lower := interp3(ds, lerps, variable, altLerp.index)
	upper := interp3(ds, lerps, variable, altLerp.index+1)
	return lower*altLerp.lerp + upper*(1-altLerp.lerp)
}

// GetWind returns interpolated (u, v) wind components for the given position.
// hour: fractional hours since dataset start.
// lat: latitude in degrees (-90 to +90).
// lng: longitude in degrees (0 to 360).
// alt: altitude in metres above sea level.
func GetWind(ds *dataset.File, warnings *Warnings, hour, lat, lng, alt float64) (u, v float64, err error) {
	lerps, err := pick3(hour, lat, lng)
	if err != nil {
		return 0, 0, err
	}

	altidx := search(ds, lerps, alt)
	lower := interp3(ds, lerps, dataset.VarHeight, altidx)
	upper := interp3(ds, lerps, dataset.VarHeight, altidx+1)

	var altLerp float64
	if lower != upper {
		altLerp = (upper - alt) / (upper - lower)
	} else {
		altLerp = 0.5
	}

	if altLerp < 0 {
		warnings.AltitudeTooHigh.Add(1)
	}

	alt1 := lerp1{index: altidx, lerp: altLerp}
	u = interp4(ds, lerps, alt1, dataset.VarWindU)
	v = interp4(ds, lerps, alt1, dataset.VarWindV)

	return u, v, nil
}