import os
import uuid
from dataclasses import asdict, dataclass
from typing import Any, DefaultDict, Dict, List, Optional, Tuple
import sys
sys.path.append("/home/fn/MyRL/Simple-work2-CDT/OSRL-CDT")
import bullet_safety_gym  # noqa
import dsrl
import gymnasium as gym  # noqa
import numpy as np
import pyrallis
import torch
from dsrl.infos import DENSITY_CFG
from dsrl.offline_env import OfflineEnvWrapper, wrap_env  # noqa
from fsrl.utils import WandbLogger
from torch.utils.data import DataLoader
from tqdm.auto import trange  # noqa

from examples.configs.cdt_configs import CDT_DEFAULT_CONFIG, CDTTrainConfig
from osrl.algorithms import CDT, CDTTrainer
from osrl.common import SequenceDataset
from osrl.common.exp_util import auto_name, seed_all
import pickle

dataset_path_val = f'/home/fn/OSRL-CDT/examples/train/my_cdt_data_train_noauto.pkl'
dataset_path_train = f'/home/fn/MyRL/Simple-work1/DT/dt_idve_train_data.pkl'
dataset_path_test = f'/home/fn/MyRL/Simple-work1/DT/dt_idve_test_data.pkl'


with open(dataset_path_val, 'rb') as f:
    data_val = pickle.load(f)

with open(dataset_path_test, 'rb') as f:
    data_test = pickle.load(f)

with open(dataset_path_train,'rb') as f:
    data_train = pickle.load(f)

maxcost = max(data_train['costs'])
mincost = min(data_train['costs'])

data_test['costs'] = (data_test['costs']-mincost)/(maxcost-mincost)
data_train['costs'] = (data_train['costs']-mincost)/(maxcost-mincost)

save_train = f'/home/fn/OSRL-CDT/examples/train/my_cdt_data_train_noauto_new.pkl'
save_test = f'/home/fn/OSRL-CDT/examples/train/my_cdt_data_val_noauto_new.pkl'

with open(save_test, 'wb') as f:
    pickle.dump(data_test,f)

with open(save_train,'wb') as f:
    pickle.dump(data_train,f)


a = 1