"""
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, 'logging.basicConfig', 'logging.basicConfig', ([], {'format': '"""%(levelname)s :: %(asctime)s :: %(message)s"""', 'level': 'logging.DEBUG'}), False, 'import logging\n'), (65, 'click.command', 'click.command', ([], {}), False, 'import click\n'), (66, 'click.argument', 'click.argument', (['"""num_processes"""'], {}), False, 'import click\n'), (67, 'click.argument', 'click.argument', (['"""weights_pth"""'], {}), False, 'import click\n'), (68, 'click.option', 'click.option', (['"""--skip/--no-skip"""'], {}), False, 'import click\n'), (32, 'logging.info', 'logging.info', (['"""Starting post-processing over detected objects"""'], {}), False, 'import logging\n'), (33, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (34, 'pymongo.MongoClient', 'MongoClient', (["os.environ['DBCONNECT']"], {}), False, 'from pymongo import MongoClient\n'), (35, 'logging.info', 'logging.info', (['f"""Connected to client: {client}."""'], {}), False, 'import logging\n'), (47, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (48, 'logging.info', 'logging.info', (['f"""Exiting post-processing. Time up: {end_time - start_time}"""'], {}), False, 'import logging\n'), (38, 'logging.info', 'logging.info', (['"""Loaded next batch. Running postprocessing"""'], {}), False, 'import logging\n'), (61, 'logging.info', 'logging.info', (['f"""Updated result: {result}"""'], {}), False, 'import logging\n'), (40, 'joblib.Parallel', 'Parallel', ([], {'n_jobs': 'num_processes'}), False, 'from joblib import Parallel, delayed\n'), (42, 'logging.error', 'logging.error', (['f"""Postprocessing error in referenced page: {e.page}"""'], {}), False, 'import logging\n'), (43, 'logging.error', 'logging.error', (['f"""Original Exception: {e.original_exception}"""'], {}), False, 'import logging\n'), (63, 'logging.error', 'logging.error', (['f"""Document write error: {e}\n Document id: obj["_id"]"""'], {}), False, 'import logging\n'), (40, 'joblib.delayed', 'delayed', (['run_inference'], {}), False, 'from joblib import Parallel, delayed\n')]