from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


import argparse

# from config import get_parse_args
# from utils.logger import Logger
# from utils.random_seed import set_seed
# from utils.sat_utils import solve_sat_iteratively
from satb.data.dataset_factory import _dataset_factory
# from detectors.detector_factory import detector_factory

def get_parse_args():
    parser = argparse.ArgumentParser(description='Parse Dataset')
    parser.add_argument('--type', default='gate', type=str, choices=['gate', 'cnf', 'ckt'])
    parser.add_argument('--path', default='data/random_sr3_10_100')
    parser.add_argument('--use_aig', action='store_true', default=False)
    parser.add_argument('--aig_opt', action='store_true', default=False)
    args = parser.parse_args()
    return args

def test():
    # os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus_str

    # print(args)
    # args.num_rounds = args.test_num_rounds
    args = get_parse_args()
    print(args)

    if args.type == 'gate' or args.type == 'ckt':
        dataset = _dataset_factory[args.type](args.path, use_aig=args.use_aig, aig_opt=args.aig_opt)
    elif args.type == 'cnf':
        dataset = _dataset_factory[args.type](args.path)
        
    # dataset = _dataset_factory[args.dataset](args.dataset_dir, args)
    # Do the shuffle
    # perm = torch.randperm(len(dataset))
    # dataset = dataset[perm]
    # split = args.test_split
    # dataset = dataset[:100]
    data_len = len(dataset)
    print('Total # Test SAT problems: ', data_len)

    print(dataset[0])
    print(dataset[0].y)

    print('forward level:', dataset[0].forward_level)
    print('forward index:', dataset[0].forward_index)

    print('\nbackward level:', dataset[0].backward_level)
    print('backward index:', dataset[0].backward_index)


if __name__ == '__main__':
    # args = get_parse_args()
    test()
