predictor/internal/prediction/models.go

package prediction

import (
	"math"
	"time"

	"predictor-refactored/internal/dataset"
	"predictor-refactored/internal/elevation"
)

// Exact port of the reference flight models (models.py).

const (
	pi180 = math.Pi / 180.0
	_180pi = 180.0 / math.Pi
)

// --- Up/Down Models ---

// ConstantAscent returns a model with constant vertical velocity (m/s).
func ConstantAscent(ascentRate float64) Model {
	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
		return 0, 0, ascentRate
	}
}

// DragDescent returns a descent-under-parachute model.
// seaLevelDescentRate is the descent rate at sea level (m/s, positive value).
// Uses the NASA atmosphere model for density at altitude.
func DragDescent(seaLevelDescentRate float64) Model {
	dragCoefficient := seaLevelDescentRate * 1.1045

	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
		return 0, 0, -dragCoefficient / math.Sqrt(nasaDensity(alt))
	}
}

// nasaDensity computes air density using the NASA atmosphere model.
// Reference: http://www.grc.nasa.gov/WWW/K-12/airplane/atmosmet.html
func nasaDensity(alt float64) float64 {
	var temp, pressure float64

	switch {
	case alt > 25000:
		temp = -131.21 + 0.00299*alt
		pressure = 2.488 * math.Pow((temp+273.1)/216.6, -11.388)
	case alt > 11000:
		temp = -56.46
		pressure = 22.65 * math.Exp(1.73-0.000157*alt)
	default:
		temp = 15.04 - 0.00649*alt
		pressure = 101.29 * math.Pow((temp+273.1)/288.08, 5.256)
	}

	return pressure / (0.2869 * (temp + 273.1))
}

// --- Sideways Models ---

// WindVelocity returns a model that gives lateral movement at the wind velocity.
// ds is the wind dataset, dsEpoch is the dataset start time as UNIX timestamp.
func WindVelocity(ds *dataset.File, dsEpoch float64, warnings *Warnings) Model {
	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
		tHours := (t - dsEpoch) / 3600.0
		u, v, err := GetWind(ds, warnings, tHours, lat, lng, alt)
		if err != nil {
			return 0, 0, 0
		}

		R := 6371009.0 + alt
		dlat = _180pi * v / R
		dlng = _180pi * u / (R * math.Cos(lat*pi180))
		return dlat, dlng, 0
	}
}

// --- Model Combinations ---

// LinearModel returns a model that sums all component models.
func LinearModel(models ...Model) Model {
	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
		for _, m := range models {
			d1, d2, d3 := m(t, lat, lng, alt)
			dlat += d1
			dlng += d2
			dalt += d3
		}
		return
	}
}

// --- Termination Criteria ---

// BurstTermination returns a terminator that fires when altitude >= burstAltitude.
func BurstTermination(burstAltitude float64) Terminator {
	return func(t, lat, lng, alt float64) bool {
		return alt >= burstAltitude
	}
}

// SeaLevelTermination fires when altitude <= 0.
func SeaLevelTermination(t, lat, lng, alt float64) bool {
	return alt <= 0
}

// TimeTermination returns a terminator that fires when t > maxTime.
func TimeTermination(maxTime float64) Terminator {
	return func(t, lat, lng, alt float64) bool {
		return t > maxTime
	}
}

// ElevationTermination returns a terminator that fires when alt < ground level.
// Uses ruaumoko-compatible elevation data. Longitude is normalised internally.
func ElevationTermination(elev *elevation.Dataset) Terminator {
	return func(t, lat, lng, alt float64) bool {
		return elev.Get(lat, lng) > alt
	}
}

// --- Pre-Defined Profiles ---

// Stage pairs a model with its termination criterion.
type Stage struct {
	Model      Model
	Terminator Terminator
}

// StandardProfile creates the chain for a standard high-altitude balloon flight:
// ascent at constant rate → burst → descent under parachute.
// If elev is non-nil, descent terminates at ground level; otherwise at sea level.
func StandardProfile(ascentRate, burstAltitude, descentRate float64,
	ds *dataset.File, dsEpoch float64, warnings *Warnings,
	elev *elevation.Dataset) []Stage {

	wind := WindVelocity(ds, dsEpoch, warnings)

	modelUp := LinearModel(ConstantAscent(ascentRate), wind)
	termUp := BurstTermination(burstAltitude)

	modelDown := LinearModel(DragDescent(descentRate), wind)
	var termDown Terminator
	if elev != nil {
		termDown = ElevationTermination(elev)
	} else {
		termDown = Terminator(SeaLevelTermination)
	}

	return []Stage{
		{Model: modelUp, Terminator: termUp},
		{Model: modelDown, Terminator: termDown},
	}
}

// FloatProfile creates the chain for a floating balloon flight:
// ascent to float altitude → float until stop time.
func FloatProfile(ascentRate, floatAltitude float64, stopTime time.Time,
	ds *dataset.File, dsEpoch float64, warnings *Warnings) []Stage {

	wind := WindVelocity(ds, dsEpoch, warnings)

	modelUp := LinearModel(ConstantAscent(ascentRate), wind)
	termUp := BurstTermination(floatAltitude)

	modelFloat := wind
	termFloat := TimeTermination(float64(stopTime.Unix()))

	return []Stage{
		{Model: modelUp, Terminator: termUp},
		{Model: modelFloat, Terminator: termFloat},
	}
}

// RunPrediction runs a prediction with the given profile stages.
// launchTime is a UNIX timestamp.
func RunPrediction(launchTime float64, lat, lng, alt float64, stages []Stage) []StageResult {
	chain := make([]struct {
		Model      Model
		Terminator Terminator
	}, len(stages))

	for i, s := range stages {
		chain[i].Model = s.Model
		chain[i].Terminator = s.Terminator
	}

	return Solve(launchTime, lat, lng, alt, chain)
}