"""
Aircraft object containing state variables and parameters.
"""
Base.@kwdef mutable struct Aircraft
	r::Vector{R} = zeros(3) # position [ft]
	v::Vector{R} = zeros(3) # velocity [ft/s]

	# constants
	a_max::R     = G/2      # maximum acceleration along any axis [ft/s^2]
	v_min::R     = 50.0     # minimum velocity [ft/s]
	z_min::R     = 1.0      # minimum altitude [ft]
end

vec(ac::Aircraft) = [ac.r; ac.v]

"""
Command object specifying input to aircraft controls. Typically generated by transition
distribution.
"""
Base.@kwdef mutable struct Command
	h_d::R   = 0.0  # vertical rate [ft/s]
	vg_d::R  = 0.0  # ground acceleration [ft/s^2]
	psi_d::R = 0.0  # turn rate [deg/s]
end

Command(arr::Vector) = Command(arr...)
vec(cmd::Command) = [cmd.h_d; cmd.vg_d; cmd.psi_d]

function observation(ac::Aircraft; scale::Bool=true)
	r0 = scale ? 50.0 : 1.0
	v0 = scale ? 100.0 : 1.0
	obs = [ac.r / r0; ac.v / v0]
	return obs
end

function initialize(ac::Aircraft, ac_enc::Aircraft)
	ac.r = deepcopy(ac_enc.r)
	ac.v = deepcopy(ac_enc.v)
end

function update(ac::Aircraft, cmd::Command, dt)
    r, v = ac.r, ac.v
    vg = sqrt(v[1]^2 + v[2]^2)
	h_d, vg_d, psi_d = cmd.h_d, cmd.vg_d, deg2rad(cmd.psi_d)

	# minimum velocity
	vg_d = vg < ac.v_min ? max(vg_d, 0.0) : vg_d

	# compute accelerations
    ax = vg_d * v[1] / vg + v[2] * psi_d
    ay = vg_d * v[2] / vg - v[1] * psi_d
    az = (h_d - v[3]) / dt
    a = [ax; ay; az]
    clamp!(a, -ac.a_max, ac.a_max)

	# update state
    v_new = v + a * dt
    r_new = r + (v + v_new) * dt / 2

	# minimum altitude
	if r_new[3] < ac.z_min
		r_new[3] = ac.z_min
		v_new[3] = 0.0
	end

    ac.r, ac.v = r_new, v_new
end

distance(ac1::Aircraft, ac2::Aircraft) = norm(ac1.r - ac2.r)

"""
Return tuple of horizontal and vertial separation, optionally scaled by NMAC parameters.
"""
function separation(ac1::Aircraft, ac2::Aircraft; scale::Tuple{R,R}=(1.0, 1.0))
	r1, r2 = ac1.r, ac2.r
	return (norm(r1[1:2] - r2[1:2]), (scale[1] / scale[2]) * abs(r1[3] - r2[3]))
end

"""
Reduce symmetries of two-aircraft system, yielding R^12 -> R^8. Transformed coordinate space
is in a rotated frame of reference where the center of the system is at the origin. The
aircraft lie along the x-axis (except for displacement along the z-axis) with aircraft 1 at
position (d, 0, z) and aircraft 2 at (-d, 0, -z). Horizontal velocity is expressed as a
magnitude and angle is measured clockwise from the +y direction for both aircraft.

Relative positions are further transformed with an arcsinh function. This provides log-like
squashing at large magnitudes, preserves linear behavior at small magnitudes, and extends
smoothly and antisymmetrically to negative values. Since aircraft may travel large distances
before collision, the transformation encourages the learning of scale-dependent features.
"""
function kernel(ac1::Aircraft, ac2::Aircraft; scale::Bool=true)
	ac1, ac2 = observation.((ac1, ac2); scale=scale)
	r1, v1 = ac1[1:3], ac1[4:6]
	r2, v2 = ac2[1:3], ac2[4:6]

	rc = (r1 + r2) / 2        # center of two-aircraft system
	dr = r1 - rc              # ac1 position relative to center
	d = norm(dr[1:2])         # horizontal separation
	phi = atan(dr[2], dr[1])  # angle of two-aircraft system
	vg1 = norm(v1[1:2])       # ac1 ground speed
	vg2 = norm(v2[1:2])       # ac2 ground speed
	psi1 = atan(v1[1:2]...)   # ac1 heading
	psi2 = atan(v2[1:2]...)   # ac2 heading

	state = zeros(8)
	# 2D components
	state[1] = asinh(d)
	state[2] = wrap_to_pi(psi1 + phi) # ac1 relative heading
	state[3] = vg1
	state[4] = wrap_to_pi(psi2 + phi) # ac2 relative heading
	state[5] = vg2
	# 3D components
	state[6] = asinh(dr[3])
	state[7] = v1[3]
	state[8] = v2[3]
	return state
end
