package gfs

import (
	"time"

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

// Wind is a WindField backed by a GFS dataset file.
type Wind struct {
	file *File

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

// NewWind returns a Wind backed by file. The axes are constructed from the
// file's variant geometry.
func NewWind(file *File) *Wind {
	v := file.variant
	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:    v.NumLatitudes(),
			Name: "lat",
		},
		lngAxis: numerics.Axis{
			Left: LonStart,
			Step: v.Resolution,
			N:    v.NumLongitudes(),
			Wrap: true,
			Name: "lng",
		},
	}
}

// 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
	}
	bs := [3]numerics.Bracket{bh, bla, bln}

	height := func(level int) func(i, j, k int) float64 {
		return func(i, j, k int) float64 {
			return float64(w.file.Val(i, level, VarHeight, j, k))
		}
	}

	levelIdx := numerics.Bisect(0, w.file.variant.NumLevels()-2, alt, func(level int) float64 {
		return numerics.EvalTrilinear(bs, height(level))
	})

	lowerHGT := numerics.EvalTrilinear(bs, height(levelIdx))
	upperHGT := numerics.EvalTrilinear(bs, height(levelIdx+1))

	var altFrac float64
	if lowerHGT != upperHGT {
		altFrac = (upperHGT - alt) / (upperHGT - lowerHGT)
	} else {
		altFrac = 0.5
	}

	component := func(level, variable int) float64 {
		return numerics.EvalTrilinear(bs, func(i, j, k int) float64 {
			return float64(w.file.Val(i, level, variable, j, k))
		})
	}

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

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