"""
Turn to follow another aicraft moving along a constant path
in the opposite direction to initial velocity.
"""

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/4, pi/4, pi/4


def reset(env):
    pos_ego = vector(0, env.np_random.uniform(Y_MIN, Y_MAX), env.np_random.uniform(0, Z_MAX))
    hdg_ego = env.np_random.uniform(-HDG_DEV, HDG_DEV)
    pitch_ego = env.np_random.uniform(-PITCH_DEV, +PITCH_DEV)
    roll_ego = env.np_random.uniform(-ROLL_DEV, +ROLL_DEV)
    vel_norm, _ = dir_to_vec(hdg_ego, pitch_ego, roll_ego)
    vel_ego = vel_norm * env.np_random.uniform(START_SPEED_MIN, START_SPEED_MAX)
    env.jets[0].reset(pos=pos_ego, vel=vel_ego, hdg=hdg_ego, pitch=pitch_ego, roll=roll_ego)
    pos_reference = vector(2*X_MAX,(Y_MIN + Y_MAX) / 2,Z_MAX / 2)
    vel_reference = vector(-1, 0, 0) * env.np_random.uniform(START_SPEED_MIN, START_SPEED_MAX)
    env.jets[1].reset(pos=pos_reference, vel=vel_reference, hdg=pi)

def dynamics(env, action):
    env.jets[1].pos += env.jets[1].vel / HZ * RENDER_SCALE

task = {
    "obs": DEFAULT_OBS,
    "obs_lims": DEFAULT_OBS_LIMS,
    "reset":    reset,
    "dynamics": dynamics,
    "reference_jet": {"arrow_length": 5},
    "render": lambda env: (
        curve(canvas=env.scene, pos=[vector(-2*X_MAX,Y_MID,Z_MID), vector(3*X_MAX,Y_MID,Z_MID)]))
}
