import gevent
from gevent import monkey
monkey.patch_all()

import os
import tqdm
import logging
import subprocess


from gevent.pool import Pool


logger = logging.getLogger()
logger.setLevel(logging.INFO)

file_handler = logging.FileHandler('layout_log.txt')

formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)

logger.addHandler(file_handler)

def _generate_single_data(file, data_path, out_path, bar):
    cmd = ['python', '-m', 'data_driven.layout_single_gen', '-j', file, '-d', data_path, '-o', out_path]
    # for 3D-FRONT data cleaning
    # cmd = ['python', '-m', 'data_driven.layout_single_gen', '-j', file, '-d', data_path, '-o', out_path, '-c']
    # for 3D-FRONT data room cleanning and plot
    # cmd = ['python', '-m', 'data_driven.room_single_gen', '-j', file, '-d', data_path, '-o', out_path, '-c']
    cmd_str = " ".join(cmd)
    logger.info(f'start {cmd_str}')
    try:
        ret = subprocess.run(cmd, shell=False, timeout=30, capture_output=True, check=False)
        if ret.returncode != 0:
            logger.info(f"cmd error, {cmd_str}, {ret.stderr.decode('utf-8')}")
        else:
            logger.info(f"cmd done, {cmd_str}, {ret.stdout.decode('utf-8')}")
    except subprocess.TimeoutExpired:
        logger.info(f'cmd timeout {cmd_str}')
    bar.update(1)
    return 1


if __name__ == '__main__':
    data_path = '/anonymous/data/fp_json_out'
    out_path = '/anonymous/data/fp_json_plot'
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    for x in ['frame', 'layout', 'frame_labels', 'layout_labels']:
        if not os.path.exists(os.path.join(out_path, x)):
            os.makedirs(os.path.join(out_path, x))
    file_li = [f for f in os.listdir(data_path) if f.endswith('.json')]
    print('len(file_li)', len(file_li))
    done_file_li = [f.replace('.png', '.json') for f in os.listdir(os.path.join(out_path, 'layout')) if f.endswith('.png')]
    print('done_file_li', len(done_file_li))
    file_li = set(file_li) - set(done_file_li)
    print('len(file_li)', len(file_li))
    num_worker = min(30, len(file_li))
    pool = Pool(num_worker)
    bar = tqdm.tqdm(total=len(file_li))
    jobs = []
    for file in file_li:
        jobs.append(pool.spawn(_generate_single_data, file, data_path, out_path, bar))#callback=lambda _: bar.update()
    gevent.joinall(jobs)
    bar.close()
