"""
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()

[((118, 216), 'logging.basicConfig', 'logging.basicConfig', ([], {'format': '"""%(levelname)s :: %(asctime)s :: %(message)s"""', 'level': 'logging.DEBUG'}), "(format='%(levelname)s :: %(asctime)s :: %(message)s',\n    level=logging.DEBUG)\n", (137, 216), False, 'import logging\n'), ((2490, 2505), 'click.command', 'click.command', ([], {}), '()\n', (2503, 2505), False, 'import click\n'), ((2507, 2538), 'click.argument', 'click.argument', (['"""num_processes"""'], {}), "('num_processes')\n", (2521, 2538), False, 'import click\n'), ((2540, 2569), 'click.argument', 'click.argument', (['"""weights_pth"""'], {}), "('weights_pth')\n", (2554, 2569), False, 'import click\n'), ((2571, 2603), 'click.option', 'click.option', (['"""--skip/--no-skip"""'], {}), "('--skip/--no-skip')\n", (2583, 2603), False, 'import click\n'), ((946, 1008), 'logging.info', 'logging.info', (['"""Starting post-processing over detected objects"""'], {}), "('Starting post-processing over detected objects')\n", (958, 1008), False, 'import logging\n'), ((1026, 1037), 'time.time', 'time.time', ([], {}), '()\n', (1035, 1037), False, 'import time\n'), ((1051, 1087), 'pymongo.MongoClient', 'MongoClient', (["os.environ['DBCONNECT']"], {}), "(os.environ['DBCONNECT'])\n", (1062, 1087), False, 'from pymongo import MongoClient\n'), ((1092, 1139), 'logging.info', 'logging.info', (['f"""Connected to client: {client}."""'], {}), "(f'Connected to client: {client}.')\n", (1104, 1139), False, 'import logging\n'), ((1667, 1678), 'time.time', 'time.time', ([], {}), '()\n', (1676, 1678), False, 'import time\n'), ((1683, 1757), 'logging.info', 'logging.info', (['f"""Exiting post-processing. Time up: {end_time - start_time}"""'], {}), "(f'Exiting post-processing. Time up: {end_time - start_time}')\n", (1695, 1757), False, 'import logging\n'), ((1216, 1273), 'logging.info', 'logging.info', (['"""Loaded next batch. Running postprocessing"""'], {}), "('Loaded next batch. Running postprocessing')\n", (1228, 1273), False, 'import logging\n'), ((2316, 2357), 'logging.info', 'logging.info', (['f"""Updated result: {result}"""'], {}), "(f'Updated result: {result}')\n", (2328, 2357), False, 'import logging\n'), ((1307, 1337), 'joblib.Parallel', 'Parallel', ([], {'n_jobs': 'num_processes'}), '(n_jobs=num_processes)\n', (1315, 1337), False, 'from joblib import Parallel, delayed\n'), ((1453, 1520), 'logging.error', 'logging.error', (['f"""Postprocessing error in referenced page: {e.page}"""'], {}), "(f'Postprocessing error in referenced page: {e.page}')\n", (1466, 1520), False, 'import logging\n'), ((1533, 1593), 'logging.error', 'logging.error', (['f"""Original Exception: {e.original_exception}"""'], {}), "(f'Original Exception: {e.original_exception}')\n", (1546, 1593), False, 'import logging\n'), ((2418, 2490), 'logging.error', 'logging.error', (['f"""Document write error: {e}\n Document id: obj["_id"]"""'], {}), '(f"""Document write error: {e}\n Document id: obj["_id"]""")\n', (2431, 2490), False, 'import logging\n'), ((1338, 1360), 'joblib.delayed', 'delayed', (['run_inference'], {}), '(run_inference)\n', (1345, 1360), False, 'from joblib import Parallel, delayed\n')]