package numerics

import "math"

// Spherical kinematics for trajectory integration.
//
// Positions are carried as GeoVec (degrees) because constraints, path recording
// and the API all speak latitude and longitude. Motion, however, is expressed
// as a velocity in metres per second and applied by rotating the position
// vector along a great circle. Nothing in this file forms a longitude
// derivative, which is what removes the polar singularity: dLng/dt carries a
// 1/cos(lat) factor that reached 1.46e12 deg/s at 90 degrees and made any
// fixed-step integrator meaningless there.
//
// Consequences worth knowing:
//   - A step that reaches a pole continues down the far side, and longitude
//     picks up 180 degrees on its own. No special case, no threshold latitude.
//   - Latitude cannot leave [-90, 90] by construction, because it is read back
//     from a unit vector rather than accumulated.

const (
	degToRad = math.Pi / 180
	radToDeg = 180 / math.Pi
)

// Rate is a velocity in the local horizontal frame at a point: metres per
// second toward east and north, plus metres per second upward.
type Rate struct {
	East     float64
	North    float64
	Vertical float64
}

// AddRate sums two rates componentwise. Rates compose linearly; positions do
// not, which is why they are advanced with GeoStep instead.
func AddRate(a, b Rate) Rate {
	return Rate{East: a.East + b.East, North: a.North + b.North, Vertical: a.Vertical + b.Vertical}
}

// basis returns the earth-centred unit vectors at (lat, lng): outward radial,
// east and north.
//
// All three are unit length at every latitude, the poles included. What happens
// at a pole is not a degeneracy but a genuine ambiguity: east and north there
// depend on which meridian the longitude names. That matches the data — NCEP
// resolves the GFS pole row per longitude for exactly this reason, so any
// choice yields the same physical vector.
func basis(latDeg, lngDeg float64) (radial, east, north [3]float64) {
	sinLat, cosLat := math.Sincos(latDeg * degToRad)
	sinLng, cosLng := math.Sincos(lngDeg * degToRad)
	radial = [3]float64{cosLat * cosLng, cosLat * sinLng, sinLat}
	east = [3]float64{-sinLng, cosLng, 0}
	north = [3]float64{-sinLat * cosLng, -sinLat * sinLng, cosLat}
	return radial, east, north
}

// toGeo reads a position back off a unit vector.
func toGeo(v [3]float64, altitude float64) GeoVec {
	return GeoVec{
		Lat:      math.Asin(math.Max(-1, math.Min(1, v[2]))) * radToDeg,
		Lng:      PyMod(math.Atan2(v[1], v[0])*radToDeg, 360),
		Altitude: altitude,
	}
}

// GeoStep advances a position by rate over dt seconds along a great circle.
// dt may be negative, which travels the same arc in the opposite direction.
//
// It is exact for a constant rate: the position vector is rotated in the plane
// it spans with the direction of travel. Since that direction is orthogonal to
// the radial by construction, the result stays on the unit sphere without
// renormalisation.
func GeoStep(y GeoVec, rate Rate, dt float64) GeoVec {
	altitude := y.Altitude + rate.Vertical*dt
	speed := math.Hypot(rate.East, rate.North)
	if speed == 0 {
		return GeoVec{Lat: y.Lat, Lng: y.Lng, Altitude: altitude}
	}

	radial, east, north := basis(y.Lat, y.Lng)
	var tangent [3]float64
	for i := range tangent {
		tangent[i] = (rate.East*east[i] + rate.North*north[i]) / speed
	}

	// Angle subtended at the earth's centre by the arc travelled.
	angle := speed * dt / (EarthRadius + y.Altitude)
	sinA, cosA := math.Sincos(angle)
	var out [3]float64
	for i := range out {
		out[i] = radial[i]*cosA + tangent[i]*sinA
	}
	return toGeo(out, altitude)
}

// GreatCircleMetres is the surface distance between two positions, ignoring
// altitude.
//
// Uses the chord rather than acos(dot): for nearby points acos loses most of
// its significant digits, and these distances are checked to sub-millimetre
// tolerances in tests.
func GreatCircleMetres(a, b GeoVec) float64 {
	ra, _, _ := basis(a.Lat, a.Lng)
	rb, _, _ := basis(b.Lat, b.Lng)
	var chordSq float64
	for i := range ra {
		d := ra[i] - rb[i]
		chordSq += d * d
	}
	return 2 * EarthRadius * math.Asin(math.Min(1, math.Sqrt(chordSq)/2))
}

// RateField returns the rate of change of state at (t, y). The rate is
// direction-independent; the integrator applies the sign of dt for reverse-time
// integration.
type RateField func(t float64, y GeoVec) Rate

// RK4Step performs one classical Runge-Kutta-4 step along the sphere.
//
// The four stage rates are combined as earth-centred vectors rather than as
// local east/north pairs. That distinction matters: the local frame rotates
// between stages, and near a pole it rotates fast enough that averaging
// components directly would reintroduce the very error this formulation exists
// to remove. The combined velocity is then read back in the frame at the
// starting point — which also discards the small radial component averaging
// introduces — and applied as a single great-circle step.
func RK4Step(t float64, y GeoVec, dt float64, f RateField) GeoVec {
	half := dt / 2

	k1 := f(t, y)
	y2 := GeoStep(y, k1, half)
	k2 := f(t+half, y2)
	y3 := GeoStep(y, k2, half)
	k3 := f(t+half, y3)
	y4 := GeoStep(y, k3, dt)
	k4 := f(t+dt, y4)

	points := [4]GeoVec{y, y2, y3, y4}
	rates := [4]Rate{k1, k2, k3, k4}
	weights := [4]float64{1.0 / 6, 1.0 / 3, 1.0 / 3, 1.0 / 6}

	var vx, vy, vz, vertical float64
	for i := range points {
		_, east, north := basis(points[i].Lat, points[i].Lng)
		w := weights[i]
		vx += w * (rates[i].East*east[0] + rates[i].North*north[0])
		vy += w * (rates[i].East*east[1] + rates[i].North*north[1])
		vz += w * (rates[i].East*east[2] + rates[i].North*north[2])
		vertical += w * rates[i].Vertical
	}

	_, east0, north0 := basis(y.Lat, y.Lng)
	mean := Rate{
		East:     vx*east0[0] + vy*east0[1] + vz*east0[2],
		North:    vx*north0[0] + vy*north0[1] + vz*north0[2],
		Vertical: vertical,
	}
	return GeoStep(y, mean, dt)
}