import time
import pickle

from src.utils.loading import get_mazes
from src.utils.testing import Test, count_start_neighbors, count_start_neighbors_nonparallel, count_start_neighbors
from src.utils.plotting import plot_mazes, plot_extrapolation

# # Test start neighbors counting
# inputs, solutions = get_mazes(maze_size=19, num_mazes=100, percolation=0.0, deadend_start=False)

# start_time = time.time()
# count_1 = count_start_neighbors_nonparallel(inputs)
# print(count_1)
# print(f'Time: {time.time() - start_time}')
# print()

# start_time = time.time()
# count_2 = count_start_neighbors(inputs)
# print(count_2)
# print(f'Time: {time.time() - start_time}')

# print(f'Counts match:{(count_1 == count_2).all()}')

test = Test(
        model_name='dt_net',
        dataset='maze-dataset',
        num_mazes=1000,
        iters=[100],
        maze_sizes=[9, 11, 13, 15],
        percolations=[0.1],
        deadend_start=False)

test.get_accuracies()
test.print()
print(test.get_name())

# Save the Test object
file_name = f'outputs/extrapolation/{test.get_name()}.pkl'
with open(file_name, 'wb') as f:
    pickle.dump(test, f)
