import pycocotools.coco
import argparse
import logging
import os
import pprint
import sqlite3

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("coco")

argParser = argparse.ArgumentParser(description='Create shape SVG\'s')
argParser.add_argument(
        '--annotations',
        type=str,
        default='../../datasets/COCO/annotations/instances_val2017.json'
    )
argParser.add_argument(
        '--categories',
        action='store_true',
        help='Show categories'
    )
argParser.add_argument(
        '--propagate',
        type=str,
        metavar='DATABASE',
        help='Store data in sqlite db'
    )
args = argParser.parse_args()


logger.info(f"Load {args.annotations}")
coco = pycocotools.coco.COCO(args.annotations)


if args.categories:
    cats = {}
    for id, cat in coco.cats.items():
        if cat['supercategory'] not in cats:
            cats[cat['supercategory']] = []
        cats[cat['supercategory']].append(cat)
#     pp = pprint.PrettyPrinter(indent=4)
    pprint.pprint(cats, sort_dicts=False)
    
if args.propagate:
    if not os.path.exists(args.propagate):
        con = sqlite3.connect(args.propagate)
        cur = con.cursor()
        with open('coco.sql', 'r') as fp:
            cur.executescript(fp.read())
        con.close()
        
    con = sqlite3.connect(args.propagate)
    logger.info("Create categories")
    cur = con.cursor()
    cur.executemany('INSERT OR IGNORE INTO categories(id, supercategory, name) VALUES (:id, :supercategory, :name)', coco.cats.values())
    con.commit()
    
    logger.info("Images...")
    cur.executemany('''
        INSERT OR IGNORE INTO images(id, flickr_url, coco_url, width, height, date_captured)
        VALUES (:id, :flickr_url, :coco_url, :width, :height, :date_captured)
        ''', coco.imgs.values())
    con.commit()
    
    logger.info("Annotations...")
    
    
    def annotation_generator():
        for c in coco.anns.values():
            ann = c.copy()
            ann['bbox_top'] = ann['bbox'][1]
            ann['bbox_left'] = ann['bbox'][0]
            ann['bbox_width'] = ann['bbox'][2]
            ann['bbox_height'] = ann['bbox'][3]
            yield ann
            
    cur.executemany('''
        INSERT OR IGNORE INTO annotations(id, image_id, category_id, iscrowd, area, bbox_top, bbox_left, bbox_width, bbox_height)
        VALUES (:id, :image_id, :category_id, :iscrowd, :area, :bbox_top, :bbox_left, :bbox_width, :bbox_height)
        ''', annotation_generator())
    con.commit()
    
    
    logger.info("Segments...")
    
    def segment_generator():
        for ann in coco.anns.values():
            for i, seg in enumerate(ann['segmentation']):
                yield {
                    'id': ann['id']*10 + i, # create a uniqe segment id, supports max 10 segments per annotation
                    'annotation_id': ann['id'],
                    'points': str(seg)[1:-1],
                }
            
    cur.executemany('''
        INSERT OR IGNORE INTO segments(id, annotation_id, points)
        VALUES (:id, :annotation_id, :points)
        ''', segment_generator())
    con.commit()
    
    
    logger.info("Done...")
    
#     for id, cat in coco.cats.items():
#         cur = con.cursor()
#         cur.execute