package engine

import (
	"sort"

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

// Sum composes models by summing their derivatives at each evaluation point.
//
// Useful for combining a vertical-rate model with a horizontal wind model
// into a single propagator. Equivalent to Tawhiri's LinearModel.
func Sum(models ...Model) Model {
	if len(models) == 1 {
		return models[0]
	}
	return func(t float64, s State) numerics.Rate {
		var sum numerics.Rate
		for _, m := range models {
			sum = numerics.AddRate(sum, m(t, s))
		}
		return sum
	}
}

// ConstantRate returns a model with a constant vertical velocity (m/s).
// Positive rates are upward.
func ConstantRate(rate float64) Model {
	return func(_ float64, _ State) numerics.Rate { return numerics.Rate{Vertical: rate} }
}

// ParachuteDescent returns a model where vertical velocity grows with
// altitude because thinner air provides less drag. seaLevelRate is the
// descent speed at sea level (m/s, positive).
//
// Terminal velocity at altitude is computed as
//
//	v = -k / sqrt(rho(alt)),  k = seaLevelRate * 1.1045,
//
// using the NASA atmosphere model for rho. Equivalent to Tawhiri's drag_descent.
func ParachuteDescent(seaLevelRate float64) Model {
	return func(_ float64, s State) numerics.Rate {
		return numerics.Rate{Vertical: numerics.DragTerminalVelocity(seaLevelRate, s.Altitude)}
	}
}

// RateSegment is one entry in a Piecewise rate schedule. Until is the UNIX
// timestamp at which this segment ends — the model emits the segment's
// Rate for all t < Until. The final segment's Rate is held indefinitely.
type RateSegment struct {
	Until float64
	Rate  float64
}

// Piecewise returns a model that produces a piecewise-constant vertical
// rate over a sequence of intervals. The input is sorted ascending by
// Until on construction; later segments shadow earlier ones.
func Piecewise(segments []RateSegment) Model {
	if len(segments) == 0 {
		return ConstantRate(0)
	}
	sorted := append([]RateSegment(nil), segments...)
	sort.Slice(sorted, func(i, j int) bool { return sorted[i].Until < sorted[j].Until })
	finalRate := sorted[len(sorted)-1].Rate

	return func(t float64, _ State) numerics.Rate {
		idx := sort.Search(len(sorted), func(i int) bool { return sorted[i].Until > t })
		if idx == len(sorted) {
			return numerics.Rate{Vertical: finalRate}
		}
		return numerics.Rate{Vertical: sorted[idx].Rate}
	}
}

// WindTransport returns a model that moves laterally at the wind velocity
// sampled from field. The vertical component is zero. Sampling and the
// non-fatal "above_model" event live here (orchestration); the m/s → deg/s
// wind is handed straight to the integrator as a horizontal velocity.
//
// If events is non-nil, an "above_model" event is emitted whenever the
// wind field reports altitude above the highest pressure level.
func WindTransport(field weather.WindField, events *EventSink) Model {
	return func(t float64, s State) numerics.Rate {
		sample, err := field.Wind(t, s.Lat, s.Lng, s.Altitude)
		if err != nil {
			return numerics.Rate{}
		}
		if sample.AboveModel && events != nil {
			events.Emit("above_model", t, s,
				"altitude exceeded the highest pressure level of the wind dataset; samples extrapolated")
		}
		// The wind is already a horizontal velocity; it is handed over as-is.
		// Converting it to deg/s here is what used to blow up near the poles.
		return numerics.Rate{East: sample.U, North: sample.V}
	}
}