import numpy as np
import cv2
import os

path_prefix = os.environ['UDG_DATA_PATH']
data_path = path_prefix + "mopo-local/data/antcustom/2023-04-29_12-37-47/"
angles = [0, 60, 120, 180, 240, 300]
rewards = []
trajectories = []
for a in angles:
    r = np.load(data_path+"rewards_{}.np".format(a), allow_pickle=True)
    r = np.sum(np.squeeze(r), axis=-1)
    rewards.append(r)
    trajectories.append(np.load(data_path+"trajectories_{}.np".format(a), allow_pickle=True))

# rewards = np.stack(rewards, axis=0)
# trajectories = np.stack(trajectories, axis=0)
img = np.ones((1024, 1024, 3), np.uint8)*255
color_theme = [(152, 225, 204), (152, 213, 172), (151, 201, 137), (201, 204, 132), (244, 206, 126), 
    (218, 171, 136), (203, 151, 140), (188, 130, 143), (189, 137, 170), (191, 144, 196),
    (194, 158, 241), (159, 156, 242)]
scale_f = 32
boundaries = 16
for l in range(len(angles)):
    k = 3
    ind = np.argpartition(rewards[l], -k)[-k:]
    tr = trajectories[l][ind]
    for i in range(k):
        for j in range(len(tr[i])-1):
            s = tr[i][j]
            t = tr[i][j+1]
            s = ((s + boundaries) * scale_f).astype(np.int16).tolist()
            t = ((t + boundaries) * scale_f).astype(np.int16).tolist()
            cv2.line(img, tuple(s), tuple(t), color=color_theme[l%12], thickness=2)   

img = cv2.flip(img, 0)
img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
fig_path = "mopo-local/fig/unsu.png"
cv2.imwrite(fig_path, img)      