import sys
from walk_direction_utils import *
import classes
import numpy as np
import comm_test6
import random
import pickle
import time
import math
from math import pi, cos, sin, sqrt
import TD3
import numpy as np
import matplotlib.pyplot as plt
from utils import plotLearning, ReplayBuffer
import sympy as sp

#joint_limits = [[[0, 2.36], [-1.57, 0.52], [-1.57, 1.57], [-2.618, 0.0], [-1.13, 0.44], [-1.57, 0.52]], [[0, 2.36], [-1.57, .52], [-1.57, 1.57], [-2.618, 0.0], [-1.13, 0.44], [-1.57, 0.52]], [[-0.52, 1.57]], [[-0.52, 1.57]], [[-0.52, 1.57], [-0.79, 0.79], [-1.57, 1.57]]]  # added soft limits in utils.py
pkl_folder = 'pkl_walk_direction'

score = 0
reward = 0.0
distance_new = 0.0
distance_old = 0.0
joint_states_init = [0.0 - hip_x_min_temp, 0.0 - hip_y_min_temp, 0.0 - knee_min_temp, 0.0 - ankle_y_min_temp,  0.0 - ankle_x_min_temp, 0.0 - hip_x_min_temp, 0.0 - hip_y_min_temp, 0.0 - knee_min_temp, 0.0 - ankle_y_min_temp,  0.0 - ankle_x_min_temp, 0.0 - hip_z_min_temp, 0.0 - hip_z_min_temp]#, [0.0, 0.0, 0.0, 0.0]]#, [0.0], [0.0], [0.0, 0.0, 0.0]]
joint_states = joint_states_init
true_joint_states_init = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
true_joint_states = true_joint_states_init
score_hist = []
distance_hist_long = []
game_avg = []

fallen_status = 0
vel_init = [0, 0, 0]
positions_init = [0.0, 0.0, 0.80]
rpy_init = [0, 0, -1.0, 0.0]
rpy_vel_init = [0, 0, 0]
y_pos_init = [0.0]

max_rad = 2*pi
min_rad = 0
#sign = 1 if random.random() < .5 else -1
target_yaw = random.random() * pi/2 + pi/4 # * sign
#while min_rad < target_yaw < max_rad:
#    target_yaw = random.random() * pi * sign

use_avg_foot_location = True
s = s1 = classes.State(joint_states, rpy_init, rpy_vel_init, positions_init, vel_init, true_joint_states, y_pos_init)
#if use_avg_foot_location == True:

if use_avg_foot_location == False:
    if pi/2 < target_yaw < 3*pi/2: # goal on left
        s.positions[0] = .1 # initial x_min
        s1.positions[0] = .1
    else:
        s.positions[0] = -.1
        s1.positions[0] = -.1
gamma = 0.99                # discount for future rewards
batch_size = 128            # num of transitions sampled from replay buffer
num_actions = 12
#action_init = [(hip_x_min_temp/(hip_x_min_temp-hip_x_max_temp)), (hip_y_min_temp/(hip_y_min_temp-hip_y_max_temp)), (knee_min_temp/(knee_min_temp-knee_max_temp)), (ankle_y_min_temp/(ankle_y_min_temp-ankle_y_max_temp)), (ankle_x_min_temp/(ankle_x_min_temp-ankle_x_max_temp)),
#                   (hip_x_min_temp/(hip_x_min_temp-hip_x_max_temp)), (hip_y_min_temp/(hip_y_min_temp-hip_y_max_temp)), (knee_min_temp/(knee_min_temp-knee_max_temp)), (ankle_y_min_temp/(ankle_y_min_temp-ankle_y_max_temp)), (ankle_x_min_temp/(ankle_x_min_temp-ankle_x_max_temp)),
#                   (hip_z_min_temp/(hip_z_min_temp-hip_z_max_temp)), (hip_z_min_temp/(hip_z_min_temp-hip_z_max_temp))]
#action_init = np.array([(x*2)-1 for x in action_init])

exploration_noise_init = 0.10 #.10, 0.05
exploration_noise = exploration_noise_init
polyak = 0.995              # target policy update parameter (1-tau)
policy_noise = 0.20 #.20, .10          # target policy smoothing noise
noise_clip = 0.5
policy_delay = 2            # delayed policy updates parameter
her = False
adv = True # advanced version of her
adv_obs = []
adv_obs_ = []
adv_reward = []
adv_action = []
adv_pos = [[0.0, 0.0]]
imagine = False
reward_noise = 0.0 #.20

last_action = []
time.sleep(2)
testing = False
if testing == True:
    policy_noise = 0.0
    exploration_noise_init = 0.0
    exploration_noise = 0.0
last_saved_index = 534000 # raised exp policy to from .05 .05 to .10 .20 @ 10000
                          # corrected distance equation and increased time reward from .1 to 1 @ 108000
                          # implemented her @ 157000; 0 distance_bonus when fallen @ 160000
                          # implemented imagine @ 251000 (obsolete)
                          # corrected imagine @ 259000 (obsolete)
                          # set her and adv True @ 272000
                          # hip_y min from -.22 to -.52, ankle_y max from .22 to .52 @ 274000 (ckpt) (ckpt)
                          # increased yaw_penalty from .2 to 2 @ 351000 (accidental override @ 279000) (undid increase & abs() change)
                          # yaw bonus always 0 @ 351000 (ckpt)
                          # fixed distance measurement @ 357000
                          # corrected her_yaw @ 372000
                          # fixed distance bonus persistently 0 @ 357000
                          # corrected distance again @ 351000
                          # added her again @ 359000
                          # increased exp from .10 to .20 @ 362000
                          # started from pretrained model and reverted joints to -.22 & .22 hip_y min ankle_y max resp
                          # fixed repetitive sympy import location @ 101000
                          # disabled her @ 132000
                          # tested device=cpu @211000-213000
                          # trained extra 5000 @ 240000
                          # lowered exp, policy to .02 @ 242000
                          # raised exp to .20, inc range from pi/2+pi/4 to pi+0
                          # lowered exp to .10 @ 328000
                          # decomposed yaw @ 358000
                          # exp 10 policy 10 to exp 15 policy 15 @ 361000
                          # target_yaw pi/2 @ 487000 (ckpt)
'''
                          # exp .10 policy .10 to exp .20 policy .20 @ 520000
                          # exp .20 policy .20 to exp .10 policy .10 @ 540000
                          # lr .0001 to .001 @ 594000 (before this was changed, it kept tripping after 2 steps all night...)
                          # exp .10 to .20, read replay false @ 598000
                          # exp, policy to .25 @ 600000
'''
                          # replaced delta_pos with pos for distance, widened range from just pi/2, to 0 to pi @ 509000
                          # shortened range to ip/2 + pi/4, raised lr from .001 to .01 @ 514000 (learned to go backwards due to high reward for some reason... lowered lr to .005)
                          # time bonus 0 @ 615000
                          # lowered learning rate from .005 to .001 @ 637000
                          # adv True, lr .0005 @ 650000
                          # adv False, min distance feet & torso_lower @ 664000
                          # lowered policy noise .15 to .20, lr .001 @ 679000
                          # exp .05 policy .10 @ 688000
                          # exp .20 @ 764000
                          # hip_z +-.78 to +-.4, exp .15 @ 766000
                          # q = self.l5(q)!!! previously was q = self.l3(q), removed one hidden layer, started agent fresh @ 0
                          # lr .0002 @ 3000
                          # exp .20 policy .20 @ 4000
                          # exp .40, read False @ 8000
                          # took replay from 766000, exp .20 @ 15000
                          # lr .001 @ 16000
                          # increasing exp (.10 -> .20) @ 33000
                          # pi/2 * random() + pi/4 @ 144000
                          # adv True @ 152000
                          # average foot location @ 171000
                          # exp .10, fixed adv distance not updating (never cleared adv_pos, so it kept grabbing positions from the start of training) @ 201000
                          # hip_z min/max from +-.4 to +-.6, lr .005 @ 292000 (before it kept sidestepping rather than turning its body) (ckpt)
'''
                          # lr .001 @ 394000 (kept taking very large steps and distance never increased, but maybe that's a result of avg foot distance)
'''
'''
                          # use_avg False @ 292000
                          # adv False @ 293000
                          # lr .001 @ 321000
'''
                          # iterations has a large effect on distance. set to 450 @ 292000 (read replay False)
                          # lr .002, exp .05 (<5) .10 (else) @ 330000
                          # lr .005, exp .10, hip_x/knee/ankle_x max/min/max +-.8 to +- .4, duration .3 to .2 @ 332000
                          # lr .002, exp .05 @ 375000
                          # Lowering duration rate requires fresh replay buffer
                          # duration rate .2 to .1, iter 500 @ 392000
                          # lr .0005, exp .02 @ 498000

distance_hist = []
game_hist = []
transfer_learning = False
if last_saved_index > 0 or transfer_learning == True:
    read_pkl = True
else:
    read_pkl = False

i = last_saved_index

v = True
j = 0
lr = .0005
num_states = 25
remove_states = []
load_weights = True
read_replay_buffer = True
add_num_states = 0
add_actions = 0
layer_height=250

if read_replay_buffer == True:
    print('reading replay buffer...')
    replay_buffer = pickle.load(open('./{}/replay_{}.pkl'.format(pkl_folder, i), 'rb'))
    replay_buffer.max_size = 5e4
    #TODO: max size may be too low for transfer learning?
    '''
    new_rb = ReplayBuffer()
    for num, item in enumerate(replay_buffer.buffer):
        if num % 10000 == 0:
            print('original:{}'.format(item))
            line0 = np.append(replay_buffer.buffer[num][0], [pi])
            line12 = replay_buffer.buffer[num][1:3]
            line3 = np.append(replay_buffer.buffer[num][3], [pi])
            line4 = replay_buffer.buffer[num][4]
            line = np.empty(0)
            line = np.array(np.append(line, line0))
            line = np.array(np.append(line, line12))
            line = np.array(np.append(line, line3))
            line = np.array(np.append(line, line4))
            #new_rb.add(line)
            print('line:{}'.format(np.array(line)))
    '''
else:
    replay_buffer = ReplayBuffer()

if read_pkl == True:

    #agent = NewAgent.load_model('./pkl/agent_{}.pkl'.format(i))

    if load_weights == True:
        print('reading weights...')
        agent = TD3.TD3(lr=lr, state_dim=num_states + add_num_states - len(remove_states), action_dim=num_actions + add_actions, max_action=1.0, layer_height=layer_height)

        agent.load('./{}'.format(pkl_folder), i, additional_dims=add_num_states, additional_actions=add_actions, remove_dimensions_=remove_states)
        num_actions = num_actions + add_actions
        #print('STATES:{}'.format(agent.))
        #agent.state_dim += add_num_states
        #if add_state > 0:

    else:
        print('WARNING: LOADING FULL AGENT')
        agent = TD3.TD3.load_model('./{}/agent_{}.pkl'.format(pkl_folder, i))
        agent.use_scheduler = False
else:
    print('creating agent')
    #agent = NewAgent(alpha=0.000005, beta=0.00001, input_dims=[3], gamma=1.01, layer1_size=30, layer2_size=30, n_outputs=1, n_actions=26) # 26=13*2
    #agent = Agent(alpha=0.000025, beta=0.00025, input_dims=[19], tau=0.001, env='dummy', sigma=.5,
    #          batch_size=100,  layer1_size=200, layer2_size=250, n_actions=12, max_size=100000)
    agent = TD3.TD3(lr=lr, state_dim=num_states, action_dim=num_actions, max_action=1.0, layer_height=layer_height)
    #replay_buffer = ReplayBuffer()
    # may need to increase layer2_size
extra_update = 0
if extra_update > 0:
    print('updating...')
    agent.update(replay_buffer, extra_update, batch_size, gamma, polyak, policy_noise, noise_clip, policy_delay)
    extra_update = 0
extra_update = 0
additional_her_update = 0
print('connecting...')
server = comm_test6.Server(55430)

def solve_equation(xr, yr, mr, mt):
    #print('xr:{} yr:{} mr:{} mt:{}'.format(xr, yr, mr, mt))
    #sp.init_printing()
    x,y = sp.symbols('x,y')
    b = sp.symbols('b')
    h = sp.Eq(yr, mr * xr + b)
    br = sp.solve([h], (b))[b]

    f = sp.Eq(y, mt*x + 0)
    g = sp.Eq(y, mr*x + br)

    solution = sp.solve([f,g],(x,y))
    xd = solution[x]
    yd = solution[y]

    d = sqrt(xd**2 + yd**2) # distance along line from start when bt = 0
    print('br:{} xd:{} yd:{}'.format(br, xd, yd)) # doesn't care about positive or negitive!
    return(d)

#sym.solve([f,g,h],(x,y,z))
adjustment_points = [0]
exp_points = [.10]
policy_points = [.10]
adj_counter = 0
while True:
    if i in adjustment_points:
        exploration_noise_init = exp_points[adj_counter]
        exploration_noise = exp_points[adj_counter]
        adj_counter += 1

    print('\nround {}'.format(i))
    print('j:{}'.format(j))
    print('target_yaw:{}'.format(target_yaw))
    if pi/2 < target_yaw < 3*pi/2:
        print('LEFT')
    else:
        print('RIGHT')
    target_yaw_decomposed = [float(cos(target_yaw)), float(sin(target_yaw))]
    # update current state
    #print('last action:{}'.format(action))
    s = classes.State(s1.joint_states, s1.rpy, s1.rpy_vel, s1.positions, s1.position_vel, s1.true_joint_states, s1.y_pos)
    #print('s rpy original:{}'.format(s.rpy))
    #s.rpy.append(float(sin(s.rpy[2]))) # decompose yaw
    #s.rpy[2] = float(cos(s.rpy[2]))
    print('s positions:{}'.format(s.positions))
    #print('s velocities:{}'.format(s.position_vel))
    #print('s rpy:{}'.format(s.rpy))
    #print('s rpy_vel:{}'.format(s.rpy_vel))

    obs = np.array([s.positions[2]] + s.position_vel + s.rpy + s.rpy_vel + s.joint_states + target_yaw_decomposed)# + s.y_pos)#
    #print('obs:{}'.format(obs))
    action = agent.select_action(obs)
    #print('raw action:{}'.format(action))
    if testing == False:
        if j < 5:
            exploration_noise = .2
        #elif j < 7:
        #    exploration_noise = .10
        else:
            exploration_noise = .2

        noise = np.random.normal(0, exploration_noise, size=num_actions)

        action = action + noise
        override = False
        if override == True:
            action[10] = max(min(action[10], .2), -.2) # "Clip" hip_z action between +-20%
            action[11] = max(min(action[11], .2), -.2)
        action = action.clip(-1.0, 1.0)
        #print('clipped actions:{}'.format(action))
        print('exploration noise:{}'.format(exploration_noise))

    #print('ACTION:{}'.format(action))
    s1, fallen_status, _, roll, pitch, yaw, duration, sim_time = new_state_after_action(s, action, target_yaw_decomposed, server, i, last_saved_index)  # new state after taking the best action
    print('s1 positions:{}'.format(s1.positions))
    s1.rpy.append(sin(s1.rpy[2])) # decompose yaw
    s1.rpy[2] = cos(s1.rpy[2])
    x_pos, y_pos, z_pos = s1.positions
    x_pos_prev, y_pos_prev, z_pos_prev = s.positions
    print('x_pos:{} y_pos:{}'.format(x_pos, y_pos))
    print('x_pos_prev:{} y_pos_prev:{}'.format(x_pos_prev, y_pos_prev))
    #x_pos, y_pos = adj_x, adj_y # use adjusted min/max coordinates
    adv_pos.append([x_pos, y_pos])
    delta_x = x_pos - x_pos_prev
    delta_y = y_pos - y_pos_prev
    if her == True:

        her_yaw = math.atan(delta_y / delta_x) #formerly just used actual yaw
        if delta_x < 0:
            her_yaw += pi
        obs_her = np.array([s.positions[2]] + s.position_vel + s.rpy + s.rpy_vel + s.joint_states + [her_yaw])

    #distance = sqrt((x_pos * cos(target_yaw - pi/2)) ** 2 + (y_pos * sin(target_yaw - pi/2)) ** 2)
    #distance_bonus = 10.0 * sqrt(((x_pos - x_pos_prev) * cos(target_yaw)) ** 2 + ((y_pos - y_pos_prev) * sin(target_yaw)) ** 2)
    slope_target = sin(target_yaw)/cos(target_yaw)
    slope_robot = -1 / slope_target
    #yd = solpe_target * xd + 0
    distance = solve_equation(x_pos, y_pos, slope_robot, slope_target) # get SCALAR distance from (0, 0)
    #distance = distance if delta_x/target_yaw_decomposed[0] >= 0 and delta_y/target_yaw_decomposed[1] >= 0 else -1 * distance
    if abs(target_yaw_decomposed[0]) < abs(target_yaw_decomposed[1]):
        distance = distance if y_pos/target_yaw_decomposed[1] >= 0 else -1 * distance #only applied to change in distance!
    elif abs(target_yaw_decomposed[0]) >= abs(target_yaw_decomposed[1]):
        distance = distance if x_pos/target_yaw_decomposed[0] >= 0 else -1 * distance

    # theta=pi/2 + n*pi ==> delta_x <= 0 (This is not currently the case, but it should be! Direction should be considered in function BEFORE distance is calculated!)
    print('delta_x:{}'.format(delta_x))
    print('delta_y:{}'.format(delta_y))
    distance_old = distance_new
    distance_new = distance
    distance_bonus = 10 * (distance_new - distance_old)
    if her == True:
        #distance_bonus_her = 10.0 * sqrt(((x_pos - x_pos_prev) * cos(her_yaw - pi/2)) ** 2 + ((y_pos - y_pos_prev) * sin(her_yaw - pi/2)) ** 2)
        distance_bonus_her = 10 * sqrt(delta_x**2 + delta_y**2)

    if imagine == True:
        obs_imagine = []
        target_yaw_imagine = []
        distance_bonus_imagine = []
        reward_imagine = []
        for _ in range(10):
            sign = 1 if random.random() < .5 else -1
            target_yaw_imagine_ = random.random() * pi * sign
            while min_rad < target_yaw_imagine_ < max_rad:
                target_yaw_imagine_ = random.random() * pi * sign
            target_yaw_imagine.append(target_yaw_imagine_)
            obs_imagine.append(np.array([s.positions[2]] + s.position_vel + s.rpy + s.rpy_vel + s.joint_states + [target_yaw_imagine_]))
            distance_old_imagine = sqrt((x_pos_prev * cos(target_yaw_imagine_ - pi/2)) ** 2 + (y_pos_prev * sin(target_yaw_imagine_ - pi/2)) ** 2)
            distance_new_imagine = sqrt((x_pos * cos(target_yaw_imagine_ - pi/2)) ** 2 + (y_pos * sin(target_yaw_imagine_ - pi/2)) ** 2)
            distance_bonus_imagine = 10.0 * (distance_new_imagine - distance_old_imagine)
            if fallen_status == 1 and distance_bonus_imagine > 0:
                distance_bonus_imagine = 0
            yaw_bonus_imagine = -.2 * abs(np.tanh(abs(target_yaw_imagine_) - abs(yaw)))
            reward_imagine.append(distance_bonus_imagine + yaw_bonus_imagine)

    if adv == True:
        adv_obs.append(np.array([s.positions[2]] + s.position_vel + s.rpy + s.rpy_vel + s.joint_states + [0, 0])) # target_yaw will be determined once fallen == 1
        adv_action.append(action)
    #distance_old = distance_new
    #distance_new = distance
    distance_hist.append(distance)
    reward = 0.0

    #print('ROLL:{} PITCH:{} YAW:{}'.format(abs(roll), abs(pitch), abs(yaw)))

    flat_bonus = 0
    roll_bonus = 0 #3 * abs(np.tanh(.05/roll))
    try:
        pitch_bonus = .1 * abs(np.tanh(.04/pitch))
        if pitch_bonus > 1:
            pitch_bonus = 1.
    except:
        pitch_bonus = 1.
    '''
    if yaw/target_yaw >= 0:
        yaw_bonus = -.2 * abs(np.tanh(abs(target_yaw - yaw))) # if same L/R direction as target, find difference
    else:
        yaw_bonus = -.2 * abs(np.tanh(abs(target_yaw - yaw))) # iff opposite L/R direction,
    '''
    yaw_bonus = 0
    yaw_bonus_her = 0

    time_bonus = 0 #.01

    #distance_bonus = 0

    #distance_bonus = 10.0 * (distance_new - distance_old)
    #distance_bonus_her = 10.0 * (distance_new_her - distance_old_her)

    '''
    if distance_bonus > 0 and fallen_status == 1:
        distance_bonus = 0
    if her == True and distance_bonus_her > 0 and fallen_status == 1:
        distance_bonus_her = 0
    '''

    if adv == True:
        adv_reward.append(pitch_bonus + time_bonus) # we'll add the yaw and distance reward once we know the target angle
    # All rewards shuold be differentiable
    # Rewards should not be cumulative (ex. reward=1 for lasting 1 second, reward=5 for lasting 5 seconds, etc) or this will throw off the reward estimator
    reward = distance_bonus + pitch_bonus + yaw_bonus + time_bonus
    if her == True:
        reward_her = distance_bonus_her + pitch_bonus + yaw_bonus_her + time_bonus

    if imagine == True:
        obs_imagine_ = []
        for k in range(10):
            #print('reward_imagine_old:{}'.format(reward_imagine[k]))
            reward_imagine[k] += time_bonus + pitch_bonus
            #print('reward_imagine_new:{}'.format(reward_imagine[k]))
            obs_imagine_.append(np.array([s1.positions[2]] + s1.position_vel + s1.rpy + s1.rpy_vel + s1.joint_states + [target_yaw_imagine[k]]))
            replay_buffer.add((obs_imagine[k], action, reward_imagine[k], obs_imagine_[k], float(fallen_status)))

    #print(s1.y_pos)
    reward += random.random() * reward_noise - reward_noise/2 # add noise to reward
    obs_ = np.array([s1.positions[2]] + s1.position_vel + s1.rpy + s1.rpy_vel + s1.joint_states + target_yaw_decomposed)# + s1.y_pos)# + list(action))

    if adv == True:
        adv_obs_.append(np.array([s1.positions[2]] + s1.position_vel + s1.rpy + s1.rpy_vel + s1.joint_states + [0, 0]))

    replay_buffer.add((obs, action, reward, obs_, float(fallen_status))) #discount factor already takes future rewards into account!
    if her == True:
        obs_her_ = np.array([s1.positions[2]] + s1.position_vel + s1.rpy + s1.rpy_vel + s1.joint_states + [her_yaw])
        print('her dist_bonus:{}'.format(distance_bonus_her))
        replay_buffer.add((obs_her, action, reward_her, obs_her_, float(fallen_status)))
        #agent.update(replay_buffer, 1, batch_size, gamma, polyak, policy_noise, noise_clip, policy_delay) # update doubled when her active

    j += 1

    score += reward

    if v == True:
        print('distance_bonus:{}'.format(distance_bonus))
        print('distance_total:{}'.format(distance))
        print('pitch_bonus:{}'.format(pitch_bonus))
        #print('yaw_bonus:{}'.format(yaw_bonus))
        print('reward:{}'.format(reward))
        if her == True:
            print('reward_her:{}'.format(reward_her))
        print('score:{}'.format(score))
    if fallen_status == 1:
        min_iters = 2
        if adv == True and j >= min_iters:
            # calculate adv_target_yaw by taking the angle of the position right before when it landed
            adv_target_x = adv_pos[-min_iters][0]
            adv_target_y = adv_pos[-min_iters][1] # score may seem way off when j ends at a low number, since its first step usually puts the robot far to the left
            adv_target_yaw = math.atan(adv_target_y / adv_target_x)#math.atan(s.positions[1] / s.positions[0])
            #math.atan(delta_y / delta_x)
            adv_score = 0
            distance_old_adv = distance_new_adv = distance_adv = 0.0
            if adv_target_x < 0:
                adv_target_yaw += pi
            adv_target_yaw_decomposed = [cos(adv_target_yaw), sin(adv_target_yaw)]
            print('adv_target:{}'.format(adv_target_yaw))
            print('adv_target_decomposed:{}'.format(adv_target_yaw_decomposed))
            print('adv_pos:{}'.format(adv_pos))
            #slope_target = sin(adv_target_yaw)/cos(adv_target_yaw)
            slope_target = adv_target_yaw_decomposed[1]/adv_target_yaw_decomposed[0]
            slope_robot = -1 / slope_target
            if use_avg_foot_location == False:
                # cannot use adv with use_avg False because it doesn't previously know whether to take the max or min foot positions :(
                if pi/2 < adv_target_yaw < 3*pi/2: # goal on left
                    adv_pos[0] = [.1, 0.0]
                else:
                    adv_pos[0] = [-.1, 0.0]
            for k in range(j):
                adv_fallen = 0
                adv_obs[k][-2] = adv_target_yaw_decomposed[0]
                adv_obs[k][-1] = adv_target_yaw_decomposed[1]
                adv_obs_[k][-2] = adv_target_yaw_decomposed[0]
                adv_obs_[k][-1] = adv_target_yaw_decomposed[1]
                #print('obs[k][-2]:{} obs[k][-1]:{} obs_[k][-2]:{} obs[k][-1]:{}'.format(adv_obs[k][-2], adv_obs[k][-1], adv_obs_[k][-2], adv_obs_[k][-1]))
                adv_x_prev = adv_pos[k][0]
                adv_y_prev = adv_pos[k][1]
                adv_x = adv_pos[k+1][0]
                adv_y = adv_pos[k+1][1]
                print('adv_x_prev:{} adv_y_prev:{} adv_x:{} adv_y:{}'.format(adv_x_prev, adv_y_prev, adv_x, adv_y))
                adv_reward[k] += 0#-.2 * abs(np.tanh(abs(abs(adv_target_yaw) - abs(adv_obs_[k][6]))))

                time_start = time.time()
                distance_adv = solve_equation(adv_x, adv_y, slope_robot, slope_target)
                #print('distance_adv_orig:{}'.format(distance_adv))

                #print('time:{}'.format(time.time() - time_start))

                delta_x_adv = adv_x - adv_x_prev
                delta_y_adv = adv_y - adv_y_prev
                #distance_adv = distance_adv if delta_x_adv/adv_target_yaw_decomposed[0] >= 0 and delta_y_adv/adv_target_yaw_decomposed[1] >= 0 else -1 * distance_adv
                if abs(adv_target_yaw_decomposed[0]) < abs(adv_target_yaw_decomposed[1]):
                    distance_adv = distance_adv if adv_y/adv_target_yaw_decomposed[1] >= 0 else -1 * distance_adv
                elif abs(adv_target_yaw_decomposed[0]) >= abs(adv_target_yaw_decomposed[1]):
                    distance_adv = distance_adv if adv_x/adv_target_yaw_decomposed[0] >= 0 else -1 * distance_adv

                distance_old_adv = distance_new_adv
                distance_new_adv = distance_adv
                distance_adv_bonus = 10 * (distance_new_adv - distance_old_adv)
                print('dist_adv_bonus:{}'.format(distance_adv_bonus))
                if k == j-1:
                    adv_fallen = 1
                    #distance_adv_bonus = 0
                adv_reward[k] += distance_adv_bonus
                #print('adv_reward:{}'.format(adv_reward[k]))
                adv_score += adv_reward[k]
                print('{}: adv_obs:{} adv_action:{} adv_reward:{} adv_obs_:{} adv_fallen:{}'.format(k, adv_obs[k], adv_action[k], adv_reward[k], adv_obs_[k], float(adv_fallen)))
                replay_buffer.add((adv_obs[k], adv_action[k], adv_reward[k], adv_obs_[k], float(adv_fallen))) # add to rb iff not fallen
            print(adv_score)
            adv_obs = []
            adv_obs_ = []
            adv_reward = []
            adv_action = []

        if testing == False:

            '''
            if imagine == True and her == True:
                j = j * 11
            '''
            if her == True and adv == True:
                j = j * 3
            elif her == True:
                j = j * 2
            elif adv == True:
                j = j * 2

            if i > last_saved_index + extra_update:
                agent.update(replay_buffer, j + extra_update + additional_her_update, batch_size, gamma, polyak, policy_noise, noise_clip, policy_delay)
                extra_update = 0
                additional_her_update = 0

        distance_hist_long.append(max(distance_hist))
        game_hist.append(max(distance_hist))
        distance_new = distance_old = distance = 0
        #sign = 1 if random.random() < .5 else -1
        #target_yaw = random.random() * pi * sign
        target_yaw = random.random() * pi/2 + pi/4# + pi/2 + 0

        #action = action_init
        # if the robot falls, s1 is automatically back to neutral position, so we'll update all states to also be neutral

        distance_hist = []
        s.rpy = rpy_init
        s.rpy_vel = rpy_vel_init
        s.positions = positions_init
        if use_avg_foot_location == False:

            if pi/2 < target_yaw < 3*pi/2: # goal on left
                s.positions[0] = .1 # initial x_min
                adv_pos = [[50000, 50000]] # dummy, since we don't know which direction it lands yet
            else:
                s.positions[0] = -.1
                adv_pos = [[-50000, -50000]]
        else:
            adv_pos = [[0.0, 0.0]]
        s.position_vel = vel_init
        s.joint_states = joint_states_init#, [0.0, 0.0, 0.0, 0.0]]
        s.true_joint_states = true_joint_states_init
        s.y_pos = y_pos_init
        s1 = s
        j = 0
        score_hist.append(score)
        score = 0

    if i % 1000 == 0 and i != last_saved_index:
        game_avg.append(sum(game_hist)/len(game_hist))
        game_hist = []
        exploration_noise = exploration_noise_init
        print('sending pause')
        server.send_message('pause')
        if i % 1000 == 0:
            agent.save_model('./{}/agent_{}.pkl'.format(pkl_folder, i))
            agent.save('./{}'.format(pkl_folder), i)
            pickle.dump(replay_buffer, open('./{}/replay_{}.pkl'.format(pkl_folder, i), 'wb'))
        filename = './{}/TD3_{}.png'.format(pkl_folder, i)
        plotLearning(score_hist, filename=filename, window=5, erase=True)

        #filename = './pkl/reward_recent_{}.png'.format(i)
        #reward_hist_recent = reward_hist[-1000:]
        #plotLearning(reward_hist, filename=filename, window=5, erase=True)
        #reward_hist = []

        filename = './{}/distance_hist_{}.png'.format(pkl_folder, i)
        #reward_hist_recent = reward_hist[-1000:]
        plotLearning(distance_hist_long, filename=filename, window=5, erase=True, ylabel_='Distance (m)')
        filename = './{}/game_avg_{}.png'.format(pkl_folder, i)
        plotLearning(game_avg, filename=filename, window=1, erase=True, xlabel_='1000 iters', ylabel_='Avg max distance (m)')

        #TODO: Record screen with multithreading
        server.send_message('unpause')
        print('sending unpause')
        time.sleep(0.002)

    i += 1


"""
Tests of the "convenience_Horizons" routines that are used for testing.

Some of these tests really function as demos/documentation to
remind myself/ourselves of how these Horizons functions are
intended to work
"""

# Import standard packages
# --------------------------------------------------------------
import numpy as np

# Import the module to be tested
# --------------------------------------------------------------
import convenience_Horizons as Horizons

def test_read_Horizons_state_from_text():
    """
    This is NOT testing a built-in Horizons function
    This is just testing a little convenience routine created by MJP
    This convenience routine is ONLY used as part of the testing code for Cheby Checker
    """

    # input text
    lines = """
X =-2.590350154796811E+00 Y =-7.949342693459856E-02 Z = 1.245107691757731E-01
VX=-1.454708370733871E-03 VY=-9.503445860627428E-03 VZ=-3.846514535533382E-03
""".split('\n')[1:-1]

    # use the target function to extract the coordinaets
    result = Horizons.read_Horizons_state_from_text( lines )

    # check that the results are as expected
    expected_array = np.array([ float('-2.590350154796811E+00'), float('-7.949342693459856E-02'), float('1.245107691757731E-01'),
                                float('-1.454708370733871E-03'), float('-9.503445860627428E-03'), float('-3.846514535533382E-03') ] )

    assert np.allclose(expected_array, result, rtol=1e-08, atol=1e-08)

def test_extract_first_state_from_text():
    """
    This is NOT testing a built-in Horizons function
    This is just testing a little convenience routine created by MJP
    This convenience routine is ONLY used as part of the testing code for Cheby Checker
    """
    # input text
    lines = """
*******************************************************************************
JPL/HORIZONS                  12345 (1993 FT8)             2022-Jan-28 14:39:42
Rec #:   12345 (+COV) Soln.date: 2021-Nov-10_08:38:58   # obs: 1959 (1993-2021)

IAU76/J2000 helio. ecliptic osc. elements (au, days, deg., period=Julian yrs):

  EPOCH=  2457108.5 ! 2015-Mar-27.00 (TDB)         Residual RMS= .2812
   EC= .1603033905689926   QR= 2.056207695854036   TP= 2457050.1973502915
   OM= 106.4549280993016   W=  314.1929318541605   IN= 3.350816780296945
   A= 2.448750742541829    MA= 14.99600220651154   ADIST= 2.841293789229623
   PER= 3.832              N= .25720961            ANGMOM= .02657056
   DAN= 2.14602            DDN= 2.68596            L= 60.6968709
   B= -2.401858            MOID= 1.06974006        TP= 2015-Jan-27.6973502915

Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= 1.506              ROTPER= n.a.
   H= 14.52                G= .150                 B-V= n.a.
                           ALBEDO= .407            STYP= n.a.

ASTEROID comments:
1: soln ref.= JPL#32, OCC=0
2: source=ORB
*******************************************************************************

*******************************************************************************
Ephemeris / WWW_USER Fri Jan 28 14:39:42 2022 Pasadena, USA      / Horizons
*******************************************************************************
Target body name: 12345 (1993 FT8)                {source: JPL#32}
Center body name: Sun (10)                        {source: DE441}
Center-site name: BODY CENTER
*******************************************************************************
Start time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Stop  time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Step-size       : DISCRETE TIME-LIST
*******************************************************************************
Center geodetic : 0.00000000,0.00000000,0.0000000 {E-lon(deg),Lat(deg),Alt(km)}
Center cylindric: 0.00000000,0.00000000,0.0000000 {E-lon(deg),Dxy(km),Dz(km)}
Center radii    : 696000.0 x 696000.0 x 696000.0 k{Equator, meridian, pole}
Small perturbers: Yes                             {source: SB441-N16}
Output units    : AU-D
Output type     : GEOMETRIC cartesian states
Output format   : 3 (position, velocity, LT, range, range-rate)
Reference frame : ICRF
*******************************************************************************
Initial IAU76/J2000 heliocentric ecliptic osculating elements (au, days, deg.):
  EPOCH=  2457108.5 ! 2015-Mar-27.00 (TDB)         Residual RMS= .2812
   EC= .1603033905689926   QR= 2.056207695854036   TP= 2457050.1973502915
   OM= 106.4549280993016   W=  314.1929318541605   IN= 3.350816780296945
  Equivalent ICRF heliocentric cartesian coordinates (au, au/d):
   X= 3.047919278950221E-01  Y= 1.902892265722551E+00  Z= 7.692605770652556E-01
  VX=-1.255238959074424E-02 VY= 2.052146789677108E-03 VZ= 1.612315394505861E-03
Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= 1.506              ROTPER= n.a.
   H= 14.52                G= .150                 B-V= n.a.
                           ALBEDO= .407            STYP= n.a.
*******************************************************************************
JDTDB
   X     Y     Z
   VX    VY    VZ
   LT    RG    RR
*******************************************************************************
$$SOE
2458850.000000000 = A.D. 2020-Jan-01 12:00:00.0000 TDB [del_T=     69.183915 s]
 X =-2.590350154796811E+00 Y =-7.949342693459856E-02 Z = 1.245107691757731E-01
 VX=-1.454708370733871E-03 VY=-9.503445860627428E-03 VZ=-3.846514535533382E-03
 LT= 1.498492268422344E-02 RG= 2.594558933811760E+00 RR= 1.558928955626413E-03
$$EOE
*******************************************************************************

TIME

  Barycentric Dynamical Time ("TDB" or T_eph) output was requested. This
continuous relativistic coordinate time is equivalent to the relativistic
proper time of a clock at rest in a reference frame comoving with the
solar system barycenter but outside the system's gravity well. It is the
independent variable in the solar system relativistic equations of motion.

  TDB runs at a uniform rate of one SI second per second and is independent
of irregularities in Earth's rotation.

  Calendar dates prior to 1582-Oct-15 are in the Julian calendar system.
Later calendar dates are in the Gregorian system.

REFERENCE FRAME AND COORDINATES

  International Celestial Reference Frame (ICRF)

    The ICRF is an adopted reference frame whose axes are defined relative to
    fixed extragalactic radio sources distributed across the sky.

    The ICRF was aligned with the prior FK5/J2000 dynamical system at the ~0.02
    arcsecond level but is not identical and has no associated standard epoch.

  Symbol meaning [1 au= 149597870.700 km, 1 day= 86400.0 s]:

    JDTDB    Julian Day Number, Barycentric Dynamical Time
    del_T    Time-scale conversion difference TDB - UT (s)
      X      X-component of position vector (au)
      Y      Y-component of position vector (au)
      Z      Z-component of position vector (au)
      VX     X-component of velocity vector (au/day)
      VY     Y-component of velocity vector (au/day)
      VZ     Z-component of velocity vector (au/day)
      LT     One-way down-leg Newtonian light-time (day)
      RG     Range; distance from coordinate center (au)
      RR     Range-rate; radial velocity wrt coord. center (au/day)

ABERRATIONS AND CORRECTIONS

 Geometric state vectors have NO corrections or aberrations applied.

Computations by ...

    Solar System Dynamics Group, Horizons On-Line Ephemeris System
    4800 Oak Grove Drive, Jet Propulsion Laboratory
    Pasadena, CA  91109   USA

    General site: https://ssd.jpl.nasa.gov/
    Mailing list: https://ssd.jpl.nasa.gov/email_list.html
    System news : https://ssd.jpl.nasa.gov/horizons/news.html
    User Guide  : https://ssd.jpl.nasa.gov/horizons/manual.html
    Connect     : browser        https://ssd.jpl.nasa.gov/horizons/app.html#/x
                  API            https://ssd-api.jpl.nasa.gov/doc/horizons.html
                  command-line   telnet ssd.jpl.nasa.gov 6775
                  e-mail/batch   https://ssd.jpl.nasa.gov/ftp/ssd/hrzn_batch.txt
                  scripts        https://ssd.jpl.nasa.gov/ftp/ssd/SCRIPTS
    Author      : Jon.D.Giorgini@jpl.nasa.gov
*******************************************************************************
""".split('\n')[1:-1]

    #print(lines)
    # use the target function to extract the coordinaets
    result = Horizons.extract_first_state_from_text( lines )

    # check that the results are as expected
    expected_array = np.array([ float('-2.590350154796811E+00'), float('-7.949342693459856E-02'), float('1.245107691757731E-01'),
                                float('-1.454708370733871E-03'), float('-9.503445860627428E-03'), float('-3.846514535533382E-03') ] )

    assert np.allclose(expected_array, result, rtol=1e-08, atol=1e-08)

def test_nice_Horizons_A():
    """
    Testing Mike A's convenience wrapper around Horizon query functionality
     - Much of this test is being done to provide some reminder
       to myself/ourselves as to how to use the Horizons tool

    Deliberately *not* using all of the functionalities of pytest here.
    Just want to keep it simple and keep it obvious what everything is supposed to be doing.

    Here we extract the
        HELIOCENTRIC state
    for
        Asteroid number 12345 (== 1993 FT8)
    in an
        EQUATORIAL FRAME (refplane='earth')

    """

    # Define the variables that will be used in the query
    target  = '12345'   # <<-- Asteroid number 12345 == 1993    FT8
    centre  = '500@10'
    epochs  = '2458850.0'
    id_type = 'smallbody'
    refplane= 'earth'

    # Hardpaste the expected results from a by-hand query of horizons
    hardpasted_results = """
*******************************************************************************
JPL/HORIZONS                  12345 (1993 FT8)             2022-Jan-28 14:39:42
Rec #:   12345 (+COV) Soln.date: 2021-Nov-10_08:38:58   # obs: 1959 (1993-2021)

IAU76/J2000 helio. ecliptic osc. elements (au, days, deg., period=Julian yrs):

  EPOCH=  2457108.5 ! 2015-Mar-27.00 (TDB)         Residual RMS= .2812
   EC= .1603033905689926   QR= 2.056207695854036   TP= 2457050.1973502915
   OM= 106.4549280993016   W=  314.1929318541605   IN= 3.350816780296945
   A= 2.448750742541829    MA= 14.99600220651154   ADIST= 2.841293789229623
   PER= 3.832              N= .25720961            ANGMOM= .02657056
   DAN= 2.14602            DDN= 2.68596            L= 60.6968709
   B= -2.401858            MOID= 1.06974006        TP= 2015-Jan-27.6973502915

Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= 1.506              ROTPER= n.a.
   H= 14.52                G= .150                 B-V= n.a.
                           ALBEDO= .407            STYP= n.a.

ASTEROID comments:
1: soln ref.= JPL#32, OCC=0
2: source=ORB
*******************************************************************************

*******************************************************************************
Ephemeris / WWW_USER Fri Jan 28 14:39:42 2022 Pasadena, USA      / Horizons
*******************************************************************************
Target body name: 12345 (1993 FT8)                {source: JPL#32}
Center body name: Sun (10)                        {source: DE441}
Center-site name: BODY CENTER
*******************************************************************************
Start time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Stop  time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Step-size       : DISCRETE TIME-LIST
*******************************************************************************
Center geodetic : 0.00000000,0.00000000,0.0000000 {E-lon(deg),Lat(deg),Alt(km)}
Center cylindric: 0.00000000,0.00000000,0.0000000 {E-lon(deg),Dxy(km),Dz(km)}
Center radii    : 696000.0 x 696000.0 x 696000.0 k{Equator, meridian, pole}
Small perturbers: Yes                             {source: SB441-N16}
Output units    : AU-D
Output type     : GEOMETRIC cartesian states
Output format   : 3 (position, velocity, LT, range, range-rate)
Reference frame : ICRF
*******************************************************************************
Initial IAU76/J2000 heliocentric ecliptic osculating elements (au, days, deg.):
  EPOCH=  2457108.5 ! 2015-Mar-27.00 (TDB)         Residual RMS= .2812
   EC= .1603033905689926   QR= 2.056207695854036   TP= 2457050.1973502915
   OM= 106.4549280993016   W=  314.1929318541605   IN= 3.350816780296945
  Equivalent ICRF heliocentric cartesian coordinates (au, au/d):
   X= 3.047919278950221E-01  Y= 1.902892265722551E+00  Z= 7.692605770652556E-01
  VX=-1.255238959074424E-02 VY= 2.052146789677108E-03 VZ= 1.612315394505861E-03
Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= 1.506              ROTPER= n.a.
   H= 14.52                G= .150                 B-V= n.a.
                           ALBEDO= .407            STYP= n.a.
*******************************************************************************
JDTDB
   X     Y     Z
   VX    VY    VZ
   LT    RG    RR
*******************************************************************************
$$SOE
2458850.000000000 = A.D. 2020-Jan-01 12:00:00.0000 TDB [del_T=     69.183915 s]
 X =-2.590350154796811E+00 Y =-7.949342693459856E-02 Z = 1.245107691757731E-01
 VX=-1.454708370733871E-03 VY=-9.503445860627428E-03 VZ=-3.846514535533382E-03
 LT= 1.498492268422344E-02 RG= 2.594558933811760E+00 RR= 1.558928955626413E-03
$$EOE
*******************************************************************************

TIME

  Barycentric Dynamical Time ("TDB" or T_eph) output was requested. This
continuous relativistic coordinate time is equivalent to the relativistic
proper time of a clock at rest in a reference frame comoving with the
solar system barycenter but outside the system's gravity well. It is the
independent variable in the solar system relativistic equations of motion.

  TDB runs at a uniform rate of one SI second per second and is independent
of irregularities in Earth's rotation.

  Calendar dates prior to 1582-Oct-15 are in the Julian calendar system.
Later calendar dates are in the Gregorian system.

REFERENCE FRAME AND COORDINATES

  International Celestial Reference Frame (ICRF)

    The ICRF is an adopted reference frame whose axes are defined relative to
    fixed extragalactic radio sources distributed across the sky.

    The ICRF was aligned with the prior FK5/J2000 dynamical system at the ~0.02
    arcsecond level but is not identical and has no associated standard epoch.

  Symbol meaning [1 au= 149597870.700 km, 1 day= 86400.0 s]:

    JDTDB    Julian Day Number, Barycentric Dynamical Time
    del_T    Time-scale conversion difference TDB - UT (s)
      X      X-component of position vector (au)
      Y      Y-component of position vector (au)
      Z      Z-component of position vector (au)
      VX     X-component of velocity vector (au/day)
      VY     Y-component of velocity vector (au/day)
      VZ     Z-component of velocity vector (au/day)
      LT     One-way down-leg Newtonian light-time (day)
      RG     Range; distance from coordinate center (au)
      RR     Range-rate; radial velocity wrt coord. center (au/day)

ABERRATIONS AND CORRECTIONS

 Geometric state vectors have NO corrections or aberrations applied.

Computations by ...

    Solar System Dynamics Group, Horizons On-Line Ephemeris System
    4800 Oak Grove Drive, Jet Propulsion Laboratory
    Pasadena, CA  91109   USA

    General site: https://ssd.jpl.nasa.gov/
    Mailing list: https://ssd.jpl.nasa.gov/email_list.html
    System news : https://ssd.jpl.nasa.gov/horizons/news.html
    User Guide  : https://ssd.jpl.nasa.gov/horizons/manual.html
    Connect     : browser        https://ssd.jpl.nasa.gov/horizons/app.html#/x
                  API            https://ssd-api.jpl.nasa.gov/doc/horizons.html
                  command-line   telnet ssd.jpl.nasa.gov 6775
                  e-mail/batch   https://ssd.jpl.nasa.gov/ftp/ssd/hrzn_batch.txt
                  scripts        https://ssd.jpl.nasa.gov/ftp/ssd/SCRIPTS
    Author      : Jon.D.Giorgini@jpl.nasa.gov
*******************************************************************************
""".split('\n')[1:-1]

    # Extract the hardpasted results into an array
    # (using convenience func, "extract_first_state_from_text", tested above)
    expected_array = Horizons.extract_first_state_from_text( hardpasted_results )

    # Call the nice_Horizons function (i.e. the focus of the test)
    result = Horizons.nice_Horizons(target, centre, epochs, id_type, refplane=refplane )

    # Check that the results are as expected
    # - Lowered the accuracy from 1e-11 to 1e-8 as the discrepany grows when JPL re-fits the orbit,
    # and I don't want to keep hand-pasting different sets of results
    assert np.allclose(expected_array, result, rtol=1e-8, atol=1e-8)

def test_nice_Horizons_B():
    """
    Testing Mike A's convenience wrapper around Horizon query functionality
     - Much of this test is being done to provide some reminder
       to myself/ourselves as to how to use the Horizons tool

    Deliberately *not* using all of the functionalities of pytest here.
    Just want to keep it simple and keep it obvious what everything is supposed to be doing.

    Here we extract the
        HELIOCENTRIC state
    for
        GEOCENTER
    in an
        EQUATORIAL FRAME (refplane='earth')

    """

    # Define the variables that will be used in the query
    target  = '399'       # Earth
    centre  = '500@10'
    epochs  = '2458850.0'
    id_type = 'majorbody'
    refplane= 'earth'

    # Hardpaste the expected results from a by-hand query of horizons
    hardpasted_results = """
*******************************************************************************
 Revised: April 12, 2021                 Earth                              399

 GEOPHYSICAL PROPERTIES (revised Aug 15, 2018):
  Vol. Mean Radius (km)    = 6371.01+-0.02   Mass x10^24 (kg)= 5.97219+-0.0006
  Equ. radius, km          = 6378.137        Mass layers:
  Polar axis, km           = 6356.752          Atmos         = 5.1   x 10^18 kg
  Flattening               = 1/298.257223563   oceans        = 1.4   x 10^21 kg
  Density, g/cm^3          = 5.51              crust         = 2.6   x 10^22 kg
  J2 (IERS 2010)           = 0.00108262545     mantle        = 4.043 x 10^24 kg
  g_p, m/s^2  (polar)      = 9.8321863685      outer core    = 1.835 x 10^24 kg
  g_e, m/s^2  (equatorial) = 9.7803267715      inner core    = 9.675 x 10^22 kg
  g_o, m/s^2               = 9.82022         Fluid core rad  = 3480 km
  GM, km^3/s^2             = 398600.435436   Inner core rad  = 1215 km
  GM 1-sigma, km^3/s^2     =      0.0014     Escape velocity = 11.186 km/s
  Rot. Rate (rad/s)        = 0.00007292115   Surface area:
  Mean sidereal day, hr    = 23.9344695944     land          = 1.48 x 10^8 km
  Mean solar day 2000.0, s = 86400.002         sea           = 3.62 x 10^8 km
  Mean solar day 1820.0, s = 86400.0         Love no., k2    = 0.299
  Moment of inertia        = 0.3308          Atm. pressure   = 1.0 bar
  Mean temperature, K      = 270             Volume, km^3    = 1.08321 x 10^12
  Mean effect. IR temp, K  = 255             Magnetic moment = 0.61 gauss Rp^3
  Geometric albedo         = 0.367           Vis. mag. V(1,0)= -3.86
  Solar Constant (W/m^2)   = 1367.6 (mean), 1414 (perihelion), 1322 (aphelion)
 HELIOCENTRIC ORBIT CHARACTERISTICS:
  Obliquity to orbit, deg  = 23.4392911  Sidereal orb period  = 1.0000174 y
  Orbital speed, km/s      = 29.79       Sidereal orb period  = 365.25636 d
  Mean daily motion, deg/d = 0.9856474   Hill's sphere radius = 234.9
*******************************************************************************

*******************************************************************************
Ephemeris / WWW_USER Fri Jan 28 16:02:02 2022 Pasadena, USA      / Horizons
*******************************************************************************
Target body name: Earth (399)                     {source: DE441}
Center body name: Sun (10)                        {source: DE441}
Center-site name: BODY CENTER
*******************************************************************************
Start time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Stop  time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Step-size       : DISCRETE TIME-LIST
*******************************************************************************
Center geodetic : 0.00000000,0.00000000,0.0000000 {E-lon(deg),Lat(deg),Alt(km)}
Center cylindric: 0.00000000,0.00000000,0.0000000 {E-lon(deg),Dxy(km),Dz(km)}
Center radii    : 696000.0 x 696000.0 x 696000.0 k{Equator, meridian, pole}
Output units    : AU-D
Output type     : GEOMETRIC cartesian states
Output format   : 3 (position, velocity, LT, range, range-rate)
Reference frame : ICRF
*******************************************************************************
JDTDB
   X     Y     Z
   VX    VY    VZ
   LT    RG    RR
*******************************************************************************
$$SOE
2458850.000000000 = A.D. 2020-Jan-01 12:00:00.0000 TDB [del_T=     69.183915 s]
 X =-1.749585912701602E-01 Y = 8.877645495087018E-01 Z = 3.848482875671789E-01
 VX=-1.721190438300784E-02 VY=-2.874039035670773E-03 VZ=-1.245648654352060E-03
 LT= 5.678966496273616E-03 RG= 9.832825679666131E-01 RR=-1.981645766688001E-05
$$EOE
*******************************************************************************

TIME

  Barycentric Dynamical Time ("TDB" or T_eph) output was requested. This
continuous relativistic coordinate time is equivalent to the relativistic
proper time of a clock at rest in a reference frame comoving with the
solar system barycenter but outside the system's gravity well. It is the
independent variable in the solar system relativistic equations of motion.

  TDB runs at a uniform rate of one SI second per second and is independent
of irregularities in Earth's rotation.

  Calendar dates prior to 1582-Oct-15 are in the Julian calendar system.
Later calendar dates are in the Gregorian system.

REFERENCE FRAME AND COORDINATES

  International Celestial Reference Frame (ICRF)

    The ICRF is an adopted reference frame whose axes are defined relative to
    fixed extragalactic radio sources distributed across the sky.

    The ICRF was aligned with the prior FK5/J2000 dynamical system at the ~0.02
    arcsecond level but is not identical and has no associated standard epoch.

  Symbol meaning [1 au= 149597870.700 km, 1 day= 86400.0 s]:

    JDTDB    Julian Day Number, Barycentric Dynamical Time
    del_T    Time-scale conversion difference TDB - UT (s)
      X      X-component of position vector (au)
      Y      Y-component of position vector (au)
      Z      Z-component of position vector (au)
      VX     X-component of velocity vector (au/day)
      VY     Y-component of velocity vector (au/day)
      VZ     Z-component of velocity vector (au/day)
      LT     One-way down-leg Newtonian light-time (day)
      RG     Range; distance from coordinate center (au)
      RR     Range-rate; radial velocity wrt coord. center (au/day)

ABERRATIONS AND CORRECTIONS

 Geometric state vectors have NO corrections or aberrations applied.

Computations by ...

    Solar System Dynamics Group, Horizons On-Line Ephemeris System
    4800 Oak Grove Drive, Jet Propulsion Laboratory
    Pasadena, CA  91109   USA

    General site: https://ssd.jpl.nasa.gov/
    Mailing list: https://ssd.jpl.nasa.gov/email_list.html
    System news : https://ssd.jpl.nasa.gov/horizons/news.html
    User Guide  : https://ssd.jpl.nasa.gov/horizons/manual.html
    Connect     : browser        https://ssd.jpl.nasa.gov/horizons/app.html#/x
                  API            https://ssd-api.jpl.nasa.gov/doc/horizons.html
                  command-line   telnet ssd.jpl.nasa.gov 6775
                  e-mail/batch   https://ssd.jpl.nasa.gov/ftp/ssd/hrzn_batch.txt
                  scripts        https://ssd.jpl.nasa.gov/ftp/ssd/SCRIPTS
    Author      : Jon.D.Giorgini@jpl.nasa.gov
*******************************************************************************
""".split('\n')[1:-1]

    # Extract the hardpasted results into an array
    # (using convenience func, "extract_first_state_from_text", tested above)
    expected_array = Horizons.extract_first_state_from_text( hardpasted_results )

    # Call the nice_Horizons function (i.e. the focus of the test)
    result = Horizons.nice_Horizons(target, centre, epochs, id_type, refplane=refplane)
    print('result=\n' , result)
    # Check that the results are as expected
    assert np.allclose(expected_array, result, rtol=1e-10, atol=1e-10)

def test_nice_Horizons_C():
    """
    Testing Mike A's convenience wrapper around Horizon query functionality
     - Much of this test is being done to provide some reminder
       to myself/ourselves as to how to use the Horizons tool

    Deliberately *not* using all of the functionalities of pytest here.
    Just want to keep it simple and keep it obvious what everything is supposed to be doing.

    Here we extract the
        TOPOCENTRIC (F51) state
    for
        Asteroid number 54321 (== 2000 JA81)
    in an
        EQUATORIAL FRAME (refplane='earth')

    """

    # Define the variables that will be used in the query
    target  = '54321'       # <<-- Asteroid number 54321 == 2000 JA81
    centre  = 'F51'
    epochs  = '2458850.0'
    id_type = 'smallbody'
    refplane= 'earth'

    # Hardpaste the expected results from a by-hand query of horizons
    hardpasted_results = """
*******************************************************************************
JPL/HORIZONS                  54321 (2000 JA81)            2022-Jan-28 16:08:57
Rec #:   54321 (+COV) Soln.date: 2021-Oct-08_04:39:24   # obs: 1315 (1979-2021)

IAU76/J2000 helio. ecliptic osc. elements (au, days, deg., period=Julian yrs):

  EPOCH=  2456698.5 ! 2014-Feb-10.00 (TDB)         Residual RMS= .27282
   EC= .2508846058943067   QR= 1.938411174247326   TP= 2456093.1011463138
   OM= 91.32740861093403   W=  91.37096816741918   IN= 6.76912753748867
   A= 2.58760024089671     MA= 143.3507250314229   ADIST= 3.236789307546094
   PER= 4.1625             N= .236787236           ANGMOM= .026786318
   DAN= 2.43937            DDN= 2.41026            L= 182.707997
   B= 6.7671807            MOID= .95572901         TP= 2012-Jun-14.6011463138

Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= n.a.               ROTPER= n.a.
   H= 14.45                G= .150                 B-V= n.a.
                           ALBEDO= n.a.            STYP= n.a.

ASTEROID comments:
1: soln ref.= JPL#33, OCC=0
2: source=ORB
*******************************************************************************

*******************************************************************************
Ephemeris / WWW_USER Fri Jan 28 16:08:57 2022 Pasadena, USA      / Horizons
*******************************************************************************
Target body name: 54321 (2000 JA81)               {source: JPL#33}
Center body name: Earth (399)                     {source: DE441}
Center-site name: Pan-STARRS 1, Haleakala
*******************************************************************************
Start time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Stop  time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Step-size       : DISCRETE TIME-LIST
*******************************************************************************
Center geodetic : 203.744100,20.7071888,3.0763821 {E-lon(deg),Lat(deg),Alt(km)}
Center cylindric: 203.744100,5971.48324,2242.1878 {E-lon(deg),Dxy(km),Dz(km)}
Center pole/equ : ITRF93                          {East-longitude positive}
Center radii    : 6378.1 x 6378.1 x 6356.8 km     {Equator, meridian, pole}
Small perturbers: Yes                             {source: SB441-N16}
Output units    : AU-D
Output type     : GEOMETRIC cartesian states
Output format   : 3 (position, velocity, LT, range, range-rate)
EOP file        : eop.220127.p220422
EOP coverage    : DATA-BASED 1962-JAN-20 TO 2022-JAN-27. PREDICTS-> 2022-APR-21
Reference frame : ICRF
*******************************************************************************
Initial IAU76/J2000 heliocentric ecliptic osculating elements (au, days, deg.):
  EPOCH=  2456698.5 ! 2014-Feb-10.00 (TDB)         Residual RMS= .27282
   EC= .2508846058943067   QR= 1.938411174247326   TP= 2456093.1011463138
   OM= 91.32740861093403   W=  91.37096816741918   IN= 6.76912753748867
  Equivalent ICRF heliocentric cartesian coordinates (au, au/d):
   X= 2.934573285149345E+00  Y=-8.702901499041770E-01  Z=-7.535748078855007E-01
  VX= 3.948600090813408E-03 VY= 7.155151609877323E-03 VZ= 2.568700850735469E-03
Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= n.a.               ROTPER= n.a.
   H= 14.45                G= .150                 B-V= n.a.
                           ALBEDO= n.a.            STYP= n.a.
*******************************************************************************
JDTDB
   X     Y     Z
   VX    VY    VZ
   LT    RG    RR
*******************************************************************************
$$SOE
2458850.000000000 = A.D. 2020-Jan-01 12:00:00.0000 TDB [del_T=     69.183916 s]
 X = 3.140272938432556E-01 Y = 1.401450872643150E+00 Z = 5.824305212783573E-01
 VX= 6.595793009138184E-03 VY= 5.366428257622971E-04 VZ= 1.577496239642071E-03
 LT= 8.950941094980980E-03 RG= 1.549807407979245E+00 RR= 2.414570690380698E-03
$$EOE
*******************************************************************************

TIME

  Barycentric Dynamical Time ("TDB" or T_eph) output was requested. This
continuous relativistic coordinate time is equivalent to the relativistic
proper time of a clock at rest in a reference frame comoving with the
solar system barycenter but outside the system's gravity well. It is the
independent variable in the solar system relativistic equations of motion.

  TDB runs at a uniform rate of one SI second per second and is independent
of irregularities in Earth's rotation.

  Calendar dates prior to 1582-Oct-15 are in the Julian calendar system.
Later calendar dates are in the Gregorian system.

REFERENCE FRAME AND COORDINATES

  International Celestial Reference Frame (ICRF)

    The ICRF is an adopted reference frame whose axes are defined relative to
    fixed extragalactic radio sources distributed across the sky.

    The ICRF was aligned with the prior FK5/J2000 dynamical system at the ~0.02
    arcsecond level but is not identical and has no associated standard epoch.

  Symbol meaning [1 au= 149597870.700 km, 1 day= 86400.0 s]:

    JDTDB    Julian Day Number, Barycentric Dynamical Time
    del_T    Time-scale conversion difference TDB - UT (s)
      X      X-component of position vector (au)
      Y      Y-component of position vector (au)
      Z      Z-component of position vector (au)
      VX     X-component of velocity vector (au/day)
      VY     Y-component of velocity vector (au/day)
      VZ     Z-component of velocity vector (au/day)
      LT     One-way down-leg Newtonian light-time (day)
      RG     Range; distance from coordinate center (au)
      RR     Range-rate; radial velocity wrt coord. center (au/day)

ABERRATIONS AND CORRECTIONS

 Geometric state vectors have NO corrections or aberrations applied.

Computations by ...

    Solar System Dynamics Group, Horizons On-Line Ephemeris System
    4800 Oak Grove Drive, Jet Propulsion Laboratory
    Pasadena, CA  91109   USA

    General site: https://ssd.jpl.nasa.gov/
    Mailing list: https://ssd.jpl.nasa.gov/email_list.html
    System news : https://ssd.jpl.nasa.gov/horizons/news.html
    User Guide  : https://ssd.jpl.nasa.gov/horizons/manual.html
    Connect     : browser        https://ssd.jpl.nasa.gov/horizons/app.html#/x
                  API            https://ssd-api.jpl.nasa.gov/doc/horizons.html
                  command-line   telnet ssd.jpl.nasa.gov 6775
                  e-mail/batch   https://ssd.jpl.nasa.gov/ftp/ssd/hrzn_batch.txt
                  scripts        https://ssd.jpl.nasa.gov/ftp/ssd/SCRIPTS
    Author      : Jon.D.Giorgini@jpl.nasa.gov
*******************************************************************************
""".split('\n')[1:-1]

    # Extract the hardpasted results into an array
    # (using convenience func, "extract_first_state_from_text", tested above)
    expected_array = Horizons.extract_first_state_from_text( hardpasted_results )

    # Call the nice_Horizons function (i.e. the focus of the test)
    result = Horizons.nice_Horizons(target, centre, epochs, id_type, refplane=refplane )
    print('result=\n' , result)
    # Check that the results are as expected
    assert np.allclose(expected_array, result, rtol=1e-11, atol=1e-11)

def test_nice_Horizons_D():
    """
    Similar to test_nice_Horizons_C, but ECLIPTIC insted of equatorial

    Here we extract the
        TOPOCENTRIC (F51) state
    for
        Asteroid number 54321 (== 2000 JA81)
    in an
        ECLIPTIC FRAME (refplane='ecliptic')

    """

    # Define the variables that will be used in the query
    target  = '54321'       # <<-- Asteroid number 54321 == 2000 JA81
    centre  = 'F51'
    epochs  = '2458850.0'
    id_type = 'smallbody'
    refplane= 'ecliptic'

    # Hardpaste the expected results from a by-hand query of horizons
    hardpasted_results = """
*******************************************************************************
JPL/HORIZONS                  54321 (2000 JA81)            2022-Jan-28 16:19:21
Rec #:   54321 (+COV) Soln.date: 2021-Oct-08_04:39:24   # obs: 1315 (1979-2021)

IAU76/J2000 helio. ecliptic osc. elements (au, days, deg., period=Julian yrs):

  EPOCH=  2456698.5 ! 2014-Feb-10.00 (TDB)         Residual RMS= .27282
   EC= .2508846058943067   QR= 1.938411174247326   TP= 2456093.1011463138
   OM= 91.32740861093403   W=  91.37096816741918   IN= 6.76912753748867
   A= 2.58760024089671     MA= 143.3507250314229   ADIST= 3.236789307546094
   PER= 4.1625             N= .236787236           ANGMOM= .026786318
   DAN= 2.43937            DDN= 2.41026            L= 182.707997
   B= 6.7671807            MOID= .95572901         TP= 2012-Jun-14.6011463138

Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= n.a.               ROTPER= n.a.
   H= 14.45                G= .150                 B-V= n.a.
                           ALBEDO= n.a.            STYP= n.a.

ASTEROID comments:
1: soln ref.= JPL#33, OCC=0
2: source=ORB
*******************************************************************************

*******************************************************************************
Ephemeris / WWW_USER Fri Jan 28 16:19:22 2022 Pasadena, USA      / Horizons
*******************************************************************************
Target body name: 54321 (2000 JA81)               {source: JPL#33}
Center body name: Earth (399)                     {source: DE441}
Center-site name: Pan-STARRS 1, Haleakala
*******************************************************************************
Start time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Stop  time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Step-size       : DISCRETE TIME-LIST
*******************************************************************************
Center geodetic : 203.744100,20.7071888,3.0763821 {E-lon(deg),Lat(deg),Alt(km)}
Center cylindric: 203.744100,5971.48324,2242.1878 {E-lon(deg),Dxy(km),Dz(km)}
Center pole/equ : ITRF93                          {East-longitude positive}
Center radii    : 6378.1 x 6378.1 x 6356.8 km     {Equator, meridian, pole}
Small perturbers: Yes                             {source: SB441-N16}
Output units    : AU-D
Output type     : GEOMETRIC cartesian states
Output format   : 3 (position, velocity, LT, range, range-rate)
EOP file        : eop.220127.p220422
EOP coverage    : DATA-BASED 1962-JAN-20 TO 2022-JAN-27. PREDICTS-> 2022-APR-21
Reference frame : Ecliptic of J2000.0
*******************************************************************************
Initial IAU76/J2000 heliocentric ecliptic osculating elements (au, days, deg.):
  EPOCH=  2456698.5 ! 2014-Feb-10.00 (TDB)         Residual RMS= .27282
   EC= .2508846058943067   QR= 1.938411174247326   TP= 2456093.1011463138
   OM= 91.32740861093403   W=  91.37096816741918   IN= 6.76912753748867
  Equivalent ICRF heliocentric cartesian coordinates (au, au/d):
   X= 2.934573285149345E+00  Y=-8.702901499041770E-01  Z=-7.535748078855007E-01
  VX= 3.948600090813408E-03 VY= 7.155151609877323E-03 VZ= 2.568700850735469E-03
Asteroid physical parameters (km, seconds, rotational period in hours):
   GM= n.a.                RAD= n.a.               ROTPER= n.a.
   H= 14.45                G= .150                 B-V= n.a.
                           ALBEDO= n.a.            STYP= n.a.
*******************************************************************************
JDTDB
   X     Y     Z
   VX    VY    VZ
   LT    RG    RR
*******************************************************************************
$$SOE
2458850.000000000 = A.D. 2020-Jan-01 12:00:00.0000 TDB [del_T=     69.183916 s]
 X = 3.140272938432556E-01 Y = 1.517483592803339E+00 Z =-2.309558662379520E-02
 VX= 6.595793009138184E-03 VY= 1.119852134073137E-03 VZ= 1.233860245870196E-03
 LT= 8.950941094980980E-03 RG= 1.549807407979244E+00 RR= 2.414570690380698E-03
$$EOE
*******************************************************************************

TIME

  Barycentric Dynamical Time ("TDB" or T_eph) output was requested. This
continuous relativistic coordinate time is equivalent to the relativistic
proper time of a clock at rest in a reference frame comoving with the
solar system barycenter but outside the system's gravity well. It is the
independent variable in the solar system relativistic equations of motion.

  TDB runs at a uniform rate of one SI second per second and is independent
of irregularities in Earth's rotation.

  Calendar dates prior to 1582-Oct-15 are in the Julian calendar system.
Later calendar dates are in the Gregorian system.

REFERENCE FRAME AND COORDINATES

  Ecliptic at the standard reference epoch

    Reference epoch: J2000.0
    X-Y plane: adopted Earth orbital plane at the reference epoch
               Note: IAU76 obliquity of 84381.448 arcseconds wrt ICRF X-Y plane
    X-axis   : ICRF
    Z-axis   : perpendicular to the X-Y plane in the directional (+ or -) sense
               of Earth's north pole at the reference epoch.

  Symbol meaning [1 au= 149597870.700 km, 1 day= 86400.0 s]:

    JDTDB    Julian Day Number, Barycentric Dynamical Time
    del_T    Time-scale conversion difference TDB - UT (s)
      X      X-component of position vector (au)
      Y      Y-component of position vector (au)
      Z      Z-component of position vector (au)
      VX     X-component of velocity vector (au/day)
      VY     Y-component of velocity vector (au/day)
      VZ     Z-component of velocity vector (au/day)
      LT     One-way down-leg Newtonian light-time (day)
      RG     Range; distance from coordinate center (au)
      RR     Range-rate; radial velocity wrt coord. center (au/day)

ABERRATIONS AND CORRECTIONS

 Geometric state vectors have NO corrections or aberrations applied.

Computations by ...

    Solar System Dynamics Group, Horizons On-Line Ephemeris System
    4800 Oak Grove Drive, Jet Propulsion Laboratory
    Pasadena, CA  91109   USA

    General site: https://ssd.jpl.nasa.gov/
    Mailing list: https://ssd.jpl.nasa.gov/email_list.html
    System news : https://ssd.jpl.nasa.gov/horizons/news.html
    User Guide  : https://ssd.jpl.nasa.gov/horizons/manual.html
    Connect     : browser        https://ssd.jpl.nasa.gov/horizons/app.html#/x
                  API            https://ssd-api.jpl.nasa.gov/doc/horizons.html
                  command-line   telnet ssd.jpl.nasa.gov 6775
                  e-mail/batch   https://ssd.jpl.nasa.gov/ftp/ssd/hrzn_batch.txt
                  scripts        https://ssd.jpl.nasa.gov/ftp/ssd/SCRIPTS
    Author      : Jon.D.Giorgini@jpl.nasa.gov
*******************************************************************************
""".split('\n')[1:-1]

    # Extract the hardpasted results into an array
    # (using convenience func, "extract_first_state_from_text", tested above)
    expected_array = Horizons.extract_first_state_from_text( hardpasted_results )

    # Call the nice_Horizons function (i.e. the focus of the test)
    result = Horizons.nice_Horizons(target, centre, epochs, id_type, refplane=refplane)
    print('result=\n' , result)
    # Check that the results are as expected
    assert np.allclose(expected_array, result, rtol=1e-11, atol=1e-11)

def test_nice_Horizons_E():
    """
    Here we use Horizons to get the Heliocentric EQUATORIAL position of the Observatory
    (NB we use a hack, setting the target as the Sun, and the center as the observatory)

    Here we extract the
        TOPOCENTRIC (F51) state
    for
        The Sun
    in an
        EQUATORIAL FRAME (refplane='earth')

    """

    # Define the variables that will be used in the query
    target  = '10'
    centre  = 'F51'
    epochs  = '2458850.0'
    id_type = 'majorbody'
    refplane='earth'

    # Hardpaste the expected results from a by-hand query of horizons
    hardpasted_results = """
*******************************************************************************
 Revised: July 31, 2013                  Sun                                 10

 PHYSICAL PROPERTIES (updated 2018-Aug-15):
  GM, km^3/s^2          = 132712440041.93938  Mass, 10^24 kg        = ~1988500
  Vol. mean radius, km  = 695700              Volume, 10^12 km^3    = 1412000
  Solar radius (IAU)    = 696000 km           Mean density, g/cm^3  = 1.408
  Radius (photosphere)  = 696500 km           Angular diam at 1 AU  = 1919.3"
  Photosphere temp., K  = 6600 (bottom)       Photosphere temp., K  = 4400(top)
  Photospheric depth    = ~500 km             Chromospheric depth   = ~2500 km
  Flatness, f           = 0.00005             Adopted sid. rot. per.= 25.38 d
  Surface gravity       =  274.0 m/s^2        Escape speed, km/s    =  617.7
  Pole (RA,DEC), deg.   = (286.13, 63.87)     Obliquity to ecliptic = 7.25 deg.
  Solar constant (1 AU) = 1367.6 W/m^2        Luminosity, 10^24 J/s = 382.8
  Mass-energy conv rate = 4.260 x 10^9 kg/s   Effective temp, K     = 5772
  Sunspot cycle         = 11.4 yr             Cycle 24 sunspot min. = 2008 A.D.

  Motion relative to nearby stars = apex : R.A.= 271 deg.; DEC.= +30 deg.
                                    speed: 19.4 km/s (0.0112 au/day)
  Motion relative to 2.73K BB/CBR = apex : l= 264.7 +- 0.8; b= 48.2 +- 0.5 deg.
                                    speed: 369 +-11 km/s
*******************************************************************************

*******************************************************************************
Ephemeris / WWW_USER Fri Jan 28 16:31:17 2022 Pasadena, USA      / Horizons
*******************************************************************************
Target body name: Sun (10)                        {source: DE441}
Center body name: Earth (399)                     {source: DE441}
Center-site name: Pan-STARRS 1, Haleakala
*******************************************************************************
Start time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Stop  time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Step-size       : DISCRETE TIME-LIST
*******************************************************************************
Center geodetic : 203.744100,20.7071888,3.0763821 {E-lon(deg),Lat(deg),Alt(km)}
Center cylindric: 203.744100,5971.48324,2242.1878 {E-lon(deg),Dxy(km),Dz(km)}
Center pole/equ : ITRF93                          {East-longitude positive}
Center radii    : 6378.1 x 6378.1 x 6356.8 km     {Equator, meridian, pole}
Output units    : AU-D
Output type     : GEOMETRIC cartesian states
Output format   : 3 (position, velocity, LT, range, range-rate)
EOP file        : eop.220127.p220422
EOP coverage    : DATA-BASED 1962-JAN-20 TO 2022-JAN-27. PREDICTS-> 2022-APR-21
Reference frame : ICRF
*******************************************************************************
JDTDB
   X     Y     Z
   VX    VY    VZ
   LT    RG    RR
*******************************************************************************
$$SOE
2458850.000000000 = A.D. 2020-Jan-01 12:00:00.0000 TDB [del_T=     69.183916 s]
 X = 1.749807755042866E-01 Y =-8.877977147373203E-01 Z =-3.848633185157228E-01
 VX= 1.742085896638510E-02 VY= 3.013989047237847E-03 VZ= 1.245251026128347E-03
 LT= 5.679196211202660E-03 RG= 9.833223418736220E-01 RR=-1.085591308641175E-04
$$EOE
*******************************************************************************

TIME

  Barycentric Dynamical Time ("TDB" or T_eph) output was requested. This
continuous relativistic coordinate time is equivalent to the relativistic
proper time of a clock at rest in a reference frame comoving with the
solar system barycenter but outside the system's gravity well. It is the
independent variable in the solar system relativistic equations of motion.

  TDB runs at a uniform rate of one SI second per second and is independent
of irregularities in Earth's rotation.

  Calendar dates prior to 1582-Oct-15 are in the Julian calendar system.
Later calendar dates are in the Gregorian system.

REFERENCE FRAME AND COORDINATES

  International Celestial Reference Frame (ICRF)

    The ICRF is an adopted reference frame whose axes are defined relative to
    fixed extragalactic radio sources distributed across the sky.

    The ICRF was aligned with the prior FK5/J2000 dynamical system at the ~0.02
    arcsecond level but is not identical and has no associated standard epoch.

  Symbol meaning [1 au= 149597870.700 km, 1 day= 86400.0 s]:

    JDTDB    Julian Day Number, Barycentric Dynamical Time
    del_T    Time-scale conversion difference TDB - UT (s)
      X      X-component of position vector (au)
      Y      Y-component of position vector (au)
      Z      Z-component of position vector (au)
      VX     X-component of velocity vector (au/day)
      VY     Y-component of velocity vector (au/day)
      VZ     Z-component of velocity vector (au/day)
      LT     One-way down-leg Newtonian light-time (day)
      RG     Range; distance from coordinate center (au)
      RR     Range-rate; radial velocity wrt coord. center (au/day)

ABERRATIONS AND CORRECTIONS

 Geometric state vectors have NO corrections or aberrations applied.

Computations by ...

    Solar System Dynamics Group, Horizons On-Line Ephemeris System
    4800 Oak Grove Drive, Jet Propulsion Laboratory
    Pasadena, CA  91109   USA

    General site: https://ssd.jpl.nasa.gov/
    Mailing list: https://ssd.jpl.nasa.gov/email_list.html
    System news : https://ssd.jpl.nasa.gov/horizons/news.html
    User Guide  : https://ssd.jpl.nasa.gov/horizons/manual.html
    Connect     : browser        https://ssd.jpl.nasa.gov/horizons/app.html#/x
                  API            https://ssd-api.jpl.nasa.gov/doc/horizons.html
                  command-line   telnet ssd.jpl.nasa.gov 6775
                  e-mail/batch   https://ssd.jpl.nasa.gov/ftp/ssd/hrzn_batch.txt
                  scripts        https://ssd.jpl.nasa.gov/ftp/ssd/SCRIPTS
    Author      : Jon.D.Giorgini@jpl.nasa.gov
*******************************************************************************
""".split('\n')[1:-1]

    # Extract the hardpasted results into an array
    # (using convenience func, "extract_first_state_from_text", tested above)
    expected_array = Horizons.extract_first_state_from_text( hardpasted_results )

    # Call the nice_Horizons function (i.e. the focus of the test)
    result = Horizons.nice_Horizons(target, centre, epochs, id_type, refplane=refplane)
    print('result=\n' , result)
    # Check that the results are as expected
    assert np.allclose(expected_array, result, rtol=1e-11, atol=1e-11)

def test_nice_Horizons_F():
    """
    Similar to test_nice_Horizons_E, but ECLIPTIC instead of equatorial

    Here we use Horizons to get the Heliocentric ECLIPTIC position of the Observatory
        (NB we use a hack, setting the target as the Sun, and the center as the observatory)

    Here we extract the
        TOPOCENTRIC (F51) state
    for
        The Sun
    in an
        ECLIPTIC FRAME (refplane='ecliptic')

    """

    # Define the variables that will be used in the query
    target  = '10'
    centre  = 'F51'
    epochs  = '2458850.0'
    id_type = 'majorbody'
    refplane='ecliptic'

    # Hardpaste the expected results from a by-hand query of horizons
    hardpasted_results = """
*******************************************************************************
 Revised: July 31, 2013                  Sun                                 10

 PHYSICAL PROPERTIES (updated 2018-Aug-15):
  GM, km^3/s^2          = 132712440041.93938  Mass, 10^24 kg        = ~1988500
  Vol. mean radius, km  = 695700              Volume, 10^12 km^3    = 1412000
  Solar radius (IAU)    = 696000 km           Mean density, g/cm^3  = 1.408
  Radius (photosphere)  = 696500 km           Angular diam at 1 AU  = 1919.3"
  Photosphere temp., K  = 6600 (bottom)       Photosphere temp., K  = 4400(top)
  Photospheric depth    = ~500 km             Chromospheric depth   = ~2500 km
  Flatness, f           = 0.00005             Adopted sid. rot. per.= 25.38 d
  Surface gravity       =  274.0 m/s^2        Escape speed, km/s    =  617.7
  Pole (RA,DEC), deg.   = (286.13, 63.87)     Obliquity to ecliptic = 7.25 deg.
  Solar constant (1 AU) = 1367.6 W/m^2        Luminosity, 10^24 J/s = 382.8
  Mass-energy conv rate = 4.260 x 10^9 kg/s   Effective temp, K     = 5772
  Sunspot cycle         = 11.4 yr             Cycle 24 sunspot min. = 2008 A.D.

  Motion relative to nearby stars = apex : R.A.= 271 deg.; DEC.= +30 deg.
                                    speed: 19.4 km/s (0.0112 au/day)
  Motion relative to 2.73K BB/CBR = apex : l= 264.7 +- 0.8; b= 48.2 +- 0.5 deg.
                                    speed: 369 +-11 km/s
*******************************************************************************

*******************************************************************************
Ephemeris / WWW_USER Fri Jan 28 16:43:25 2022 Pasadena, USA      / Horizons
*******************************************************************************
Target body name: Sun (10)                        {source: DE441}
Center body name: Earth (399)                     {source: DE441}
Center-site name: Pan-STARRS 1, Haleakala
*******************************************************************************
Start time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Stop  time      : A.D. 2020-Jan-01 12:00:00.0000 TDB
Step-size       : DISCRETE TIME-LIST
*******************************************************************************
Center geodetic : 203.744100,20.7071888,3.0763821 {E-lon(deg),Lat(deg),Alt(km)}
Center cylindric: 203.744100,5971.48324,2242.1878 {E-lon(deg),Dxy(km),Dz(km)}
Center pole/equ : ITRF93                          {East-longitude positive}
Center radii    : 6378.1 x 6378.1 x 6356.8 km     {Equator, meridian, pole}
Output units    : AU-D
Output type     : GEOMETRIC cartesian states
Output format   : 3 (position, velocity, LT, range, range-rate)
EOP file        : eop.220127.p220422
EOP coverage    : DATA-BASED 1962-JAN-20 TO 2022-JAN-27. PREDICTS-> 2022-APR-21
Reference frame : Ecliptic of J2000.0
*******************************************************************************
JDTDB
   X     Y     Z
   VX    VY    VZ
   LT    RG    RR
*******************************************************************************
$$SOE
2458850.000000000 = A.D. 2020-Jan-01 12:00:00.0000 TDB [del_T=     69.183916 s]
 X = 1.749807755042866E-01 Y =-9.676283142792063E-01 Z = 4.045892447000447E-05
 VX= 1.742085896638510E-02 VY= 3.260613297708340E-03 VZ=-5.640051197420877E-05
 LT= 5.679196211202660E-03 RG= 9.833223418736221E-01 RR=-1.085591308641179E-04
$$EOE
*******************************************************************************

TIME

  Barycentric Dynamical Time ("TDB" or T_eph) output was requested. This
continuous relativistic coordinate time is equivalent to the relativistic
proper time of a clock at rest in a reference frame comoving with the
solar system barycenter but outside the system's gravity well. It is the
independent variable in the solar system relativistic equations of motion.

  TDB runs at a uniform rate of one SI second per second and is independent
of irregularities in Earth's rotation.

  Calendar dates prior to 1582-Oct-15 are in the Julian calendar system.
Later calendar dates are in the Gregorian system.

REFERENCE FRAME AND COORDINATES

  Ecliptic at the standard reference epoch

    Reference epoch: J2000.0
    X-Y plane: adopted Earth orbital plane at the reference epoch
               Note: IAU76 obliquity of 84381.448 arcseconds wrt ICRF X-Y plane
    X-axis   : ICRF
    Z-axis   : perpendicular to the X-Y plane in the directional (+ or -) sense
               of Earth's north pole at the reference epoch.

  Symbol meaning [1 au= 149597870.700 km, 1 day= 86400.0 s]:

    JDTDB    Julian Day Number, Barycentric Dynamical Time
    del_T    Time-scale conversion difference TDB - UT (s)
      X      X-component of position vector (au)
      Y      Y-component of position vector (au)
      Z      Z-component of position vector (au)
      VX     X-component of velocity vector (au/day)
      VY     Y-component of velocity vector (au/day)
      VZ     Z-component of velocity vector (au/day)
      LT     One-way down-leg Newtonian light-time (day)
      RG     Range; distance from coordinate center (au)
      RR     Range-rate; radial velocity wrt coord. center (au/day)

ABERRATIONS AND CORRECTIONS

 Geometric state vectors have NO corrections or aberrations applied.

Computations by ...

    Solar System Dynamics Group, Horizons On-Line Ephemeris System
    4800 Oak Grove Drive, Jet Propulsion Laboratory
    Pasadena, CA  91109   USA

    General site: https://ssd.jpl.nasa.gov/
    Mailing list: https://ssd.jpl.nasa.gov/email_list.html
    System news : https://ssd.jpl.nasa.gov/horizons/news.html
    User Guide  : https://ssd.jpl.nasa.gov/horizons/manual.html
    Connect     : browser        https://ssd.jpl.nasa.gov/horizons/app.html#/x
                  API            https://ssd-api.jpl.nasa.gov/doc/horizons.html
                  command-line   telnet ssd.jpl.nasa.gov 6775
                  e-mail/batch   https://ssd.jpl.nasa.gov/ftp/ssd/hrzn_batch.txt
                  scripts        https://ssd.jpl.nasa.gov/ftp/ssd/SCRIPTS
    Author      : Jon.D.Giorgini@jpl.nasa.gov
*******************************************************************************
""".split('\n')[1:-1]

    # Extract the hardpasted results into an array
    # (using convenience func, "extract_first_state_from_text", tested above)
    expected_array = Horizons.extract_first_state_from_text( hardpasted_results )

    # Call the nice_Horizons function (i.e. the focus of the test)
    result = Horizons.nice_Horizons(target, centre, epochs, id_type, refplane=refplane)
    print('result=\n' , result)
    # Check that the results are as expected
    assert np.allclose(expected_array, result, rtol=1e-11, atol=1e-11)


import os
import random
import numpy as np
from itertools import chain, combinations

import PIL.Image as Image
from PIL import ImageFont
import torch
from torchvision import transforms
import torch.optim as optim
from sklearn.metrics import average_precision_score

from modalities.CelebaImg import Img
from modalities.CelebaText import Text
from celeba.CelebADataset import CelebaDataset
from celeba.networks.VAEbimodalCelebA import VAEbimodalCelebA
from celeba.networks.ConvNetworkImgClfCelebA import ClfImg as ClfImg
from celeba.networks.ConvNetworkTextClfCelebA import ClfText as ClfText

from celeba.networks.ConvNetworksImgCelebA import EncoderImg, DecoderImg
from celeba.networks.ConvNetworksTextCelebA import EncoderText, DecoderText

from utils.BaseExperiment import BaseExperiment

LABELS = ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive',
          'Bags_Under_Eyes', 'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair',
          'Blond_Hair', 'Blurry', 'Brown_Hair', 'Bushy_Eyebrows',
          'Chubby', 'Double_Chin', 'Eyeglasses', 'Goatee',
          'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones', 'Male',
          'Mouth_Slightly_Open', 'Mustache',
          'Narrow_Eyes', 'No_Beard', 'Oval_Face', 'Pale_Skin',
          'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks',
          'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair',
          'Wearing_Earrings', 'Wearing_Hat',
          'Wearing_Lipstick', 'Wearing_Necklace',
          'Wearing_Necktie', 'Young'];

class CelebaExperiment(BaseExperiment):
    def __init__(self, flags, alphabet):
        self.labels = LABELS;
        self.flags = flags;
        self.dataset = flags.dataset;
        self.plot_img_size = torch.Size((3, 64, 64))
        self.font = ImageFont.truetype('FreeSerif.ttf', 38)

        self.alphabet = alphabet;
        self.flags.num_features = len(alphabet);

        self.modalities = self.set_modalities();
        self.num_modalities = len(self.modalities.keys());
        self.subsets = self.set_subsets();
        self.dataset_train = None;
        self.dataset_test = None;
        self.set_dataset();

        self.mm_vae = self.set_model();
        self.clfs = self.set_clfs();
        self.optimizer = None;
        self.rec_weights = self.set_rec_weights();
        self.style_weights = self.set_style_weights();

        self.test_samples = self.get_test_samples();
        self.eval_metric = average_precision_score;
        self.paths_fid = self.set_paths_fid();

    def set_model(self):
        model = VAEbimodalCelebA(self.flags, self.modalities, self.subsets)
        model = model.to(self.flags.device);
        return model;

    def set_modalities(self):
        mod1 = Img(EncoderImg(self.flags),
                   DecoderImg(self.flags),
                   self.plot_img_size);
        mod2 = Text(EncoderText(self.flags),
                    DecoderText(self.flags),
                    self.flags.len_sequence,
                    self.alphabet,
                    self.plot_img_size,
                    self.font);
        mods = {mod1.name: mod1, mod2.name: mod2};
        return mods;

    def get_transform_celeba(self):
        offset_height = (218 - self.flags.crop_size_img) // 2
        offset_width = (178 - self.flags.crop_size_img) // 2
        crop = lambda x: x[:, offset_height:offset_height + self.flags.crop_size_img,
                         offset_width:offset_width + self.flags.crop_size_img]
        transform = transforms.Compose([transforms.ToTensor(),
                                        transforms.Lambda(crop),
                                        transforms.ToPILImage(),
                                        transforms.Resize(size=(self.flags.img_size,
                                                                self.flags.img_size),
                                                          interpolation=Image.BICUBIC),
                                        transforms.ToTensor()])

        return transform;

    def set_dataset(self):
        transform = self.get_transform_celeba();
        d_train = CelebaDataset(self.flags, self.alphabet, partition=0, transform=transform)
        d_eval = CelebaDataset(self.flags, self.alphabet, partition=1, transform=transform)
        self.dataset_train = d_train;
        self.dataset_test = d_eval;

    def set_clfs(self):
        model_clf_m1 = None;
        model_clf_m2 = None;
        if self.flags.use_clf:
            model_clf_m1 = ClfImg(self.flags);
            model_clf_m1.load_state_dict(torch.load(os.path.join(self.flags.dir_clf,
                                                                 self.flags.clf_save_m1)))
            model_clf_m1 = model_clf_m1.to(self.flags.device);

            model_clf_m2 = ClfText(self.flags);
            model_clf_m2.load_state_dict(torch.load(os.path.join(self.flags.dir_clf,
                                                                 self.flags.clf_save_m2)))
            model_clf_m2 = model_clf_m2.to(self.flags.device);

        clfs = {'img': model_clf_m1,
                'text': model_clf_m2}
        return clfs;

    def set_optimizer(self):
        # optimizer definition
        optimizer = optim.Adam(
            list(self.mm_vae.parameters()),
            lr=self.flags.initial_learning_rate,
            betas=(self.flags.beta_1, self.flags.beta_2))
        self.optimizer = optimizer;

    def set_rec_weights(self):
        rec_weights = dict();
        ref_mod_d_size = self.modalities['img'].data_size.numel()/3;
        for k, m_key in enumerate(self.modalities.keys()):
            mod = self.modalities[m_key];
            numel_mod = mod.data_size.numel()
            rec_weights[mod.name] = float(ref_mod_d_size/numel_mod)
        return rec_weights;

    def set_style_weights(self):
        weights = dict();
        weights['img'] = self.flags.beta_m1_style;
        weights['text'] = self.flags.beta_m2_style;
        return weights;

    def get_prediction_from_attr(self, values):
        return values.ravel();

    def get_prediction_from_attr_random(self, values, index=None):
        return values[:,index] > 0.5;

    def eval_label(self, values, labels, index=None):
        pred = values[:,index];
        gt = labels[:,index];
        try:
            ap = self.eval_metric(gt, pred)
        except ValueError:
            ap = 0.0;
        return ap;

    def get_test_samples(self, num_images=10):
        n_test = self.dataset_test.__len__();
        samples = []
        for i in range(10):
            ix = np.random.randint(0, len(self.dataset_test.img_names))
            sample, target = self.dataset_test.__getitem__(random.randint(0, n_test))
            for k, key in enumerate(sample):
                sample[key] = sample[key].to(self.flags.device);
            samples.append(sample)
        return samples

    def mean_eval_metric(self, values):
        return np.mean(np.array(values));


import omnifig as fig
import random
import networkx as nx
import torch
from torch import nn
from torch import distributions as distrib
import numpy as np
import zlib
from omnibelt import unspecified_argument

from omnilearn import util
from omnilearn import Function

from .methods import Autoencoder

@fig.AutoModifier('freezable')
class Freezable(Autoencoder):
	def __init__(self, A, freeze_encoder=None, freeze_decoder=None, **kwargs):

		if freeze_encoder is None:
			freeze_encoder = A.pull('freeze-encoder', False)

		if freeze_decoder is None:
			freeze_decoder = A.pull('freeze-decoder', False)

		if freeze_decoder and freeze_encoder:
			print('WARNING: Freezing both encoder and decoder, this is probably not intended')

		super().__init__(A, **kwargs)

		if freeze_encoder:
			print('Freezing encoder')
			for param in self.encoder.parameters():
				param.requires_grad = False

		if freeze_decoder:
			print('Freezing decoder')
			for param in self.decoder.parameters():
				param.requires_grad = False

		self.freeze_encoder = freeze_encoder
		self.freeze_decoder = freeze_decoder


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
import os
import sys
import tensorflow as tf
import importlib
import itertools
import argparse
from tensorflow.python.framework import graph_util
import numpy as np
import tensorflow.contrib.slim as slim

FLAGS = None
sys.path.append("./ToBeConvertModels/net")
net_name=os.listdir("./ToBeConvertModels/net")[0].split(".")[0]

def get_checkpoint_name(ckpt_dir):
    f=open("%s/checkpoint"%ckpt_dir,"r")
    for line  in f.readlines():
        line=line.strip()
        line=os.path.split( line.split("\"")[1] )[1]
        return line

def main():

    network = importlib.import_module(net_name)
    with tf.Graph().as_default():
        with tf.Session() as sess:
            input_tensor = tf.placeholder(dtype=tf.float32, shape=(1, 112, 112, 3), name='input')
            # Load the model metagraph and checkpoint
            cpkt_name=get_checkpoint_name("./ToBeConvertModels/checkpoint")
            cpkt_file_path=os.path.join("./ToBeConvertModels/checkpoint",cpkt_name)
			# Build the inference graph
            prelogits, _ = network.inference(input_tensor, 1.0, phase_train=False, bottleneck_layer_size=512,weight_decay=0.0)
            embeddings = tf.nn.l2_normalize(prelogits, 1, 1e-10, name='embeddings')
            saver = tf.train.Saver()
            tf.get_default_session().run(tf.global_variables_initializer())
            tf.get_default_session().run(tf.local_variables_initializer())
            saver.restore(tf.get_default_session(), cpkt_file_path)
            # Retrieve the protobuf graph definition and fix the batch norm nodes
            input_graph_def = sess.graph.as_graph_def()
            # Freeze the graph def
            output_graph_def = freeze_graph_def(sess, input_graph_def, 'embeddings')
        # Serialize and dump the output graph to the filesystem
        with tf.gfile.GFile("./ConvertedModels/%s.pb"%cpkt_name, 'wb') as f:
            f.write(output_graph_def.SerializeToString())

def freeze_graph_def(sess, input_graph_def, output_node_names):
    for node in input_graph_def.node:
        if node.op == 'RefSwitch':
            node.op = 'Switch'
            for index in xrange(len(node.input)):
                if 'moving_' in node.input[index]:
                    node.input[index] = node.input[index] + '/read'
        elif node.op == 'AssignSub':
            node.op = 'Sub'
            if 'use_locking' in node.attr: del node.attr['use_locking']
        elif node.op == 'AssignAdd':
            node.op = 'Add'
            if 'use_locking' in node.attr: del node.attr['use_locking']
        elif node.op == 'ArgMax':
            print (" network conain ArgMax ops, which is not support by snpe" )

    # Get the list of important nodes
    whitelist_names = []
    for node in input_graph_def.node:
        if (node.name.startswith('InceptionResnetV1') or node.name.startswith('embeddings') or node.name.startswith('phase_train') or node.name.startswith('Bottleneck') or node.name.startswith('Logits')):
            whitelist_names.append(node.name)

    # Replace all the variables in the graph with constants of the same values
    output_graph_def = graph_util.convert_variables_to_constants(
        sess, input_graph_def, output_node_names.split(","),
        variable_names_whitelist=whitelist_names)
    return output_graph_def

if __name__ == '__main__':
    main( )