import numpy as np
import argparse
# from os.path import join, exists
from utils import compute_return
import glob

parser = argparse.ArgumentParser()
parser.add_argument('--dir', type=str, default='log_pendulum_ct_data',
                    help='directory.')
args = parser.parse_args()

paths = [path for path in glob.glob(f'{args.dir}/array/Rewards_*.npy')]
print(paths)
print(f"loading {paths[0]}")
rewards = np.load(paths[0])
N, M = rewards.shape
print(f'N is {N}, M is {M}')
print(f'rewards for the first 5 steps and episodes: {rewards[:5, :5]}')

paths = [path for path in glob.glob(f'{args.dir}/array/Dones_*.npy')]
print(f"loading {paths[0]}")
dones = np.load(paths[0])
N, M = dones.shape
print(f'N is {N}, M is {M}')

R = compute_return(rewards, dones)
dt = 0.001
time_limit = 10
print(f"infer> return: {R}")
print(f"infer> return scaled to physical time {time_limit}: {R*dt}")
