import numpy as np

import collections
import pickle
from sklearn.cluster import DBSCAN, k_means
from sklearn.neighbors import KNeighborsClassifier
import sobol_seq
import numpy as np
import logging
logging.basicConfig(level=logging.DEBUG,  # Set the minimum logging level
                    format='%(asctime)s %(levelname)s: %(message)s',  # Define the log message format
                    datefmt='%Y-%m-%d %H:%M:%S',  # Define the date format for log messages
                    filename='Trajectory-Transformer/trajectory/datasets/print_out.txt',  # Specify the file to save log messages
                    filemode='w')  # Set the file mode to write
with open('Trajectory-Transformer/trajectory/datasets/random_data/train_f2_view100_kmeans.pkl', 'rb') as f:
    dataset = pickle.load(f)
#with open('Trajectory-Transformer/trajectory/datasets/random_data/train_f2_view100.pkl', 'rb') as f:
#    dataset2 = pickle.load(f)
a = 0
b = 0
for i in range(10000):
    #logging.info('{}'.format(dataset[0]['observations'][i]))
    #if -dataset2[0]['rewards'][i] < 0.1:
    #    a+=1
    if -dataset[0]['rewards'][i] < 0.1:
        print(dataset[0]['rewards'][i])
        print(dataset[0]['actions'][i])
        a+=1
        
    
    #if i % 99 == 0:
        #logging.info('1: {}'.format(dataset[0]['rewards'][i]))
        #logging.info('2: {}'.format(dataset2[0]['rewards'][i]))
print(a)
print(b)
# print(dataset[0]['observations'].shape)

#print(dataset[0]['next_observations'])
# print(dataset[0]['actions'].shape)
# print(dataset[0]['rewards'].shape)
# print(dataset[0]['terminals'].shape)
#X = sobol_seq.i4_sobol_generate(2, 1000, np.random.randint(0, 100))
#print(X)
#view = k_means(X, n_clusters= 100)[0]
#view_point = []
#for i in range(len(view)):
#    view_point.append(np.argmin(np.linalg.norm(view[i]-X, axis=1)))
#print(view_point)
#with open('/home/kjlin0508/multi_objective_bo/trajectory-transformer/env/Functions/domain/domain_DRZ_1000.npy', 'rb') as f:
#    dataset = np.load(f)
#print(dataset)#