import numpy as np
import pickle
import pycuda.driver as cuda
import sys
import os
from datetime import datetime

current_dir = os.path.dirname(os.path.abspath("__file__"))
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
output_dir = os.path.join(parent_dir, "output")
os.makedirs(output_dir, exist_ok=True)
sys.path.append(parent_dir)

cuda.init()

context = cuda.Device(3).make_context()

DATA_TYPE = np.float32

from env_manifold import CustomEnv

context.push()
env = CustomEnv(device_id=3)
context.pop()

size = [10, 10, 10]

def generate_random_group(size, low, high):
    intervals = np.linspace(low, high, size + 1)
    group = np.concatenate([np.random.uniform(intervals[i], intervals[i+1], 1) for i in range(size)])
    return np.sort(group)

group1 = generate_random_group(size[0], -1, 1)
group2 = generate_random_group(size[1], -1, 1)
group3 = generate_random_group(size[2], -1, 1)

data = np.empty(size, dtype=object)
for i, a1 in enumerate(sorted(group1)):
    for j, a2 in enumerate(sorted(group2)):
        for k, a3 in enumerate(sorted(group3)):
            context.push()
            action = np.array([a1, a2, a3], dtype=np.float32)
            env.reset()
            for _ in range(400):
                _, _, _, _, _ = env.step(action)
            context.pop()
            fifo = np.array(env.fifo_states.copy())
            data[i, j, k] = {'action': action, 'fifo': fifo}
            current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            print(f"{current_time} - ({i}, {j}, {k})")

with open(os.path.join(output_dir, "manifold_1k.pkl"), 'wb') as f:
    pickle.dump(data, f)