"""
Post processing on detected objects
"""
import pymongo
from pymongo import MongoClient
import time
import logging
logging.basicConfig(format='%(levelname)s :: %(asctime)s :: %(message)s', level=logging.DEBUG)
from joblib import Parallel, delayed
import click
from xgboost_model.inference import run_inference, PostprocessException
import os

def load_detected_pages(db, buffer_size):
    """
    """
    current_docs = []
    for doc in db.propose_pages.find({'postprocess': None, 'ocr': True}, no_cursor_timeout=True):
        current_docs.append(doc)
        if len(current_docs) == buffer_size:
            yield current_docs
            current_docs = []
    yield current_docs

def do_skip(page, client):
    db = client.pdfs
    coll = db.postprocess_pages
    return coll.count_documents({'pdf_name': page['pdf_name'], 'page_num': page['page_num']}, limit=1) != 0


def postprocess(db_insert_fn, num_processes, weights_pth, skip):
    logging.info('Starting post-processing over detected objects')
    start_time = time.time()
    client = MongoClient(os.environ["DBCONNECT"])
    logging.info(f'Connected to client: {client}.')
    db = client.pdfs
    for batch in load_detected_pages(db, 100):
        logging.info('Loaded next batch. Running postprocessing')
        try:
            pages = Parallel(n_jobs=num_processes)(delayed(run_inference)(page, weights_pth) for page in batch)
        except PostprocessException as e:
            logging.error(f'Postprocessing error in referenced page: {e.page}')
            logging.error(f'Original Exception: {e.original_exception}')
            continue

        db_insert_fn(pages, client)
    end_time = time.time()
    logging.info(f'Exiting post-processing. Time up: {end_time - start_time}')

def mongo_insert_fn(objs, client):
    db = client.pdfs
    for obj in objs:
        try:
            result = db.propose_pages.update_one({'_id': obj['_id']},
                                             {'$set':
                                                {
                                                    'pp_detected_objs': obj['pp_detected_objs'],
                                                    'postprocess': True
                                                }
                                             }, upsert=False)
            logging.info(f'Updated result: {result}')
        except pymongo.errors.WriterError as e:
            logging.error(f'Document write error: {e}\n Document id: obj["_id"]')

@click.command()
@click.argument("num_processes")
@click.argument("weights_pth")
@click.option('--skip/--no-skip')
def click_wrapper(num_processes, weights_pth, skip):
	postprocess(mongo_insert_fn, int(num_processes), weights_pth, skip)

if __name__ == '__main__':
	click_wrapper()

[((8, 0, 8, 94), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((65, 1, 65, 16), 'click.command', 'click.command', ({}, {}), '()', False, 'import click\n'), ((66, 1, 66, 32), 'click.argument', 'click.argument', ({(66, 16, 66, 31): '"""num_processes"""'}, {}), "('num_processes')", False, 'import click\n'), ((67, 1, 67, 30), 'click.argument', 'click.argument', ({(67, 16, 67, 29): '"""weights_pth"""'}, {}), "('weights_pth')", False, 'import click\n'), ((68, 1, 68, 33), 'click.option', 'click.option', ({(68, 14, 68, 32): '"""--skip/--no-skip"""'}, {}), "('--skip/--no-skip')", False, 'import click\n'), ((32, 4, 32, 66), 'logging.info', 'logging.info', ({(32, 17, 32, 65): '"""Starting post-processing over detected objects"""'}, {}), "('Starting post-processing over detected objects')", False, 'import logging\n'), ((33, 17, 33, 28), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((34, 13, 34, 49), 'pymongo.MongoClient', 'MongoClient', ({(34, 25, 34, 48): "os.environ['DBCONNECT']"}, {}), "(os.environ['DBCONNECT'])", False, 'from pymongo import MongoClient\n'), ((35, 4, 35, 51), 'logging.info', 'logging.info', ({(35, 17, 35, 50): 'f"""Connected to client: {client}."""'}, {}), "(f'Connected to client: {client}.')", False, 'import logging\n'), ((47, 15, 47, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((48, 4, 48, 78), 'logging.info', 'logging.info', ({(48, 17, 48, 77): 'f"""Exiting post-processing. Time up: {end_time - start_time}"""'}, {}), "(f'Exiting post-processing. Time up: {end_time - start_time}')", False, 'import logging\n'), ((38, 8, 38, 65), 'logging.info', 'logging.info', ({(38, 21, 38, 64): '"""Loaded next batch. Running postprocessing"""'}, {}), "('Loaded next batch. Running postprocessing')", False, 'import logging\n'), ((61, 12, 61, 53), 'logging.info', 'logging.info', ({(61, 25, 61, 52): 'f"""Updated result: {result}"""'}, {}), "(f'Updated result: {result}')", False, 'import logging\n'), ((40, 20, 40, 50), 'joblib.Parallel', 'Parallel', (), '', False, 'from joblib import Parallel, delayed\n'), ((42, 12, 42, 79), 'logging.error', 'logging.error', ({(42, 26, 42, 78): 'f"""Postprocessing error in referenced page: {e.page}"""'}, {}), "(f'Postprocessing error in referenced page: {e.page}')", False, 'import logging\n'), ((43, 12, 43, 72), 'logging.error', 'logging.error', ({(43, 26, 43, 71): 'f"""Original Exception: {e.original_exception}"""'}, {}), "(f'Original Exception: {e.original_exception}')", False, 'import logging\n'), ((63, 12, 63, 81), 'logging.error', 'logging.error', ({(63, 26, 63, 80): 'f"""Document write error: {e}\n Document id: obj["_id"]"""'}, {}), '(f"""Document write error: {e}\n Document id: obj["_id"]""")', False, 'import logging\n'), ((40, 51, 40, 73), 'joblib.delayed', 'delayed', ({(40, 59, 40, 72): 'run_inference'}, {}), '(run_inference)', False, 'from joblib import Parallel, delayed\n')]