# 将 MDP 数据裁剪到指定尺寸，仅用于构造小数据集加快debug时加快数据加载速度
import os
import sys
import glob
import pickle
base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(base_path)

import numpy as np
from utils.utils import load_data_and_check_quantity

# 保留的数据量
ENV_NAME = 'TSP_V2'
num_problem = 1000
num_episode = 1000

# 获取指定环境的所有数据集   
file_list = glob.glob(os.path.join(f'{base_path}/data/used/{ENV_NAME}', '*.pkl'))
file_names = [os.path.basename(file) for file in file_list]
file_names = [name for name in file_names if 'tsp20' in name]

# 加载数据 & 打印样本量
file_data = load_data_and_check_quantity(file_names, ENV_NAME, num_problem, num_episode, check_ave_obj=True)

for file_name in file_names:
    data_name = file_name[:-4]
    # 可视化评估问题
    if data_name.endswith('problem'):
        problems = file_data[data_name]
        with open(f'{base_path}/data/used/{ENV_NAME}/{data_name}.pkl', 'wb') as f:
            pickle.dump(problems, f)    
            print(f'retained_problems saved')
        
    # 可视化训练序列
    else:
        episodes = file_data[data_name]
        with open(f'{base_path}/data/used/{ENV_NAME}/{data_name}.pkl', 'wb') as f:
            pickle.dump(episodes, f)    
            print(f'retained_episodes saved')
        