"""
Continually chase another aircraft, which is receiving random control input.
"""

from vpython import *
from math import pi
from ..global_variables import *
from ..fast_jet import dir_to_vec


Y_MID = (Y_MIN + Y_MAX) / 2
Z_MID = Z_MAX / 2

HDG_DEV, PITCH_DEV, ROLL_DEV = pi/2, pi/4, pi/4

INITIAL_OFFSET = 100
REFERENCE_ACTION_FREQ = 75


def reset(env):
    # Both jets start with identical heading, pitch, roll and velocity
    pos_ego = vector(0, env.np_random.uniform(Y_MIN, Y_MAX), env.np_random.uniform(0, Z_MAX))
    hdg = env.np_random.uniform(-HDG_DEV, HDG_DEV)
    pitch = env.np_random.uniform(-PITCH_DEV, +PITCH_DEV)
    roll = env.np_random.uniform(-ROLL_DEV, +ROLL_DEV)
    vel_norm, _ = dir_to_vec(hdg, pitch, roll)
    vel = vel_norm * env.np_random.uniform(START_SPEED_MIN, START_SPEED_MAX)
    env.jets[0].reset(pos=pos_ego, vel=vel, hdg=hdg, pitch=pitch, roll=roll)
    # The reference jet starts in front of the ego one, 
    env.jets[1].reset(pos=pos_ego+(vel_norm * INITIAL_OFFSET), vel=vel, hdg=hdg, pitch=pitch, roll=roll)

def dynamics(env, action):
    if env.t % (REFERENCE_ACTION_FREQ) == 0:
        # Resample a random action
        env._reference_jet_action = env.action_space.sample()
    env.jets[1].step(env._reference_jet_action)

task = {
    "obs": DEFAULT_OBS,
    "obs_lims": DEFAULT_OBS_LIMS,
    "reset":    reset,
    "dynamics": dynamics,
    "reference_jet": {"arrow_length": 5},
}
