# Extract ego-vehicle information
# Written by Jiageng Mao

def extract_ego_inputs(data_dict):
    ego_prompts = "*"*5 + "Ego States:" + "*"*5 + "\n"

    """
    Ego-States:
        ego_states: [vx, vy, ?, ?, v_yaw (rad/s), ego_length, ego_width, v0 (vy from canbus), Kappa (steering)]
    """
    vx = data_dict['ego_states'][0]*0.5
    vy = data_dict['ego_states'][1]*0.5
    v_yaw = data_dict['ego_states'][4]
    ax = data_dict['ego_hist_traj_diff'][-1, 0] - data_dict['ego_hist_traj_diff'][-2, 0]
    ay = data_dict['ego_hist_traj_diff'][-1, 1] - data_dict['ego_hist_traj_diff'][-2, 1]
    cx = data_dict['ego_states'][2]
    cy = data_dict['ego_states'][3]
    vhead = data_dict['ego_states'][7]*0.5
    steeling = data_dict['ego_states'][8]
    ego_prompts += f"Current State:\n"
    ego_prompts += f" - Velocity (vx,vy): ({vx:.2f},{vy:.2f})\n"
    ego_prompts += f" - Heading Angular Velocity (v_yaw): ({v_yaw:.2f})\n"
    ego_prompts += f" - Acceleration (ax,ay): ({ax:.2f},{ay:.2f})\n"
    ego_prompts += f" - Can Bus: ({cx:.2f},{cy:.2f})\n"
    ego_prompts += f" - Heading Speed: ({vhead:.2f})\n"
    ego_prompts += f" - Steering: ({steeling:.2f})\n"

    """
    Historical Trjectory:
        ego_hist_traj: [5, 2] last 2 seconds 
        ego_hist_traj_diff: [4, 2] last 2 seconds, differential format, viewed as velocity 
    """
    xh1 = data_dict['ego_hist_traj'][0][0]
    yh1 = data_dict['ego_hist_traj'][0][1]
    xh2 = data_dict['ego_hist_traj'][1][0]
    yh2 = data_dict['ego_hist_traj'][1][1]
    xh3 = data_dict['ego_hist_traj'][2][0]
    yh3 = data_dict['ego_hist_traj'][2][1]
    xh4 = data_dict['ego_hist_traj'][3][0]
    yh4 = data_dict['ego_hist_traj'][3][1]
    ego_prompts += f"Historical Trajectory (last 2 seconds):"
    ego_prompts += f" [({xh1:.2f},{yh1:.2f}), ({xh2:.2f},{yh2:.2f}), ({xh3:.2f},{yh3:.2f}), ({xh4:.2f},{yh4:.2f})]\n"
    
    """
    Mission goal:
        goal
    """
    cmd_vec = data_dict['goal']
    right, left, forward = cmd_vec
    if right > 0:
        mission_goal = "RIGHT"
    elif left > 0:
        mission_goal = "LEFT"
    else:
        assert forward > 0
        mission_goal = "FORWARD"
    ego_prompts += f"Mission Goal: "
    ego_prompts += f"{mission_goal}\n"
    
    ego_data = {}
    keys_to_extract = ['token', 'ego_states', 'ego_hist_traj_diff', 'ego_hist_traj', 'goal']
    ego_data = {key: data_dict[key] for key in keys_to_extract if key in data_dict}
    return ego_prompts, ego_data

def get_ego_prompts(data_dict):
    """Get ego prompt inputs from driving data (used as inital inputs)"""
    ego_prompts = extract_ego_inputs(data_dict)
    return ego_prompts