import os
import sys

#if "/opt/ros/kinetic/lib/python2.7/dist-packages" in sys.path:
 #   sys.path.remove("/opt/ros/kinetic/lib/python2.7/dist-packages")
import cv2
import numpy as np
from skimage import io
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import patches, lines
from matplotlib.patches import Polygon, Rectangle
from matplotlib.collections import PatchCollection


from pycocotools.coco import COCO
#matplotlib.use('TkAgg')


def showAnns(cat_names,anns):
    if len(anns) == 0:
        return 0
    ax = plt.gca()
    ax.set_autoscale_on(False)
    captions = []
    polygons = []
    rectangles = []
    color = []
    for ann in anns:
        c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0]
        # if 'segmentation' in ann:
        #     if type(ann['segmentation']) == list:
        #         # polygon
        #         for seg in ann['segmentation']:
        #             print(132131,ann['category_id'])
        #             print(cat_names[0])
        captions.append(cat_names[ann['category_id']-1])
        #poly = np.array(seg).reshape((int(len(seg) / 2), 2))
        l_corner, w, h = (ann['bbox'][0], ann['bbox'][1]), ann['bbox'][2], ann['bbox'][3]
        rectangles.append(Rectangle(l_corner, w, h))
        #polygons.append(Polygon(poly))
        color.append(c)

    p = PatchCollection(rectangles, facecolor='none', edgecolors=color, alpha=1, linestyle='--', linewidths=2)
    ax.add_collection(p)

    for i in range(len(captions)):
        x = rectangles[i].xy[0]
        y = rectangles[i].xy[1]
        ax.text(x, y, captions[i], size=10, verticalalignment='top', color='w', backgroundcolor="none")

    # p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.6)
    # ax.add_collection(p)
    # #p = PatchCollection(polygons, facecolor='none', edgecolors=color, linewidths=2)
    # p = PatchCollection(polygons, facecolor='none', edgecolors='b', linewidths=0.5)
    # ax.add_collection(p)
    print('Ok!')



import random

#annfile = "/data/DataManagement/Project/CAS_ProjectBird/CDB_ProjectBird/coco/annotations/val.json"  #####coco json
#imgroot = "/data/DataManagement/Project/CAS_ProjectBird/CDB_ProjectBird/images/test"   ####image_path
def draw(coco,imgids,cat_names,imgroot,m, n, i):
    img = coco.loadImgs([str(imgids[np.random.randint(0, len(imgids))])])[0]
    I = io.imread(os.path.join(imgroot, img['file_name']))
    plt.subplot(m, n, i)
    plt.axis('off')
    plt.title(img['file_name'], fontsize=8, color='blue')
    plt.imshow(I, aspect='equal')
    annids = coco.getAnnIds(imgIds=[img['id']])
    anns = coco.loadAnns(annids)
    showAnns(cat_names,anns)
def checkCOCO(annfile,imgroot):
    coco = COCO(annfile)
    cats = coco.loadCats(coco.getCatIds())
    cat_names = [cat['name'] for cat in cats]
    print(cat_names)
    catids = coco.getCatIds(catNms=random.randint(0,len(cat_names)-1))
    imgids = coco.getImgIds(catIds=catids)
    if 1:
        m=4
        n=4
        plt.figure(figsize=(m*6,n*4))
        plt.subplots_adjust(top=1,bottom=0,left=0,right=1,hspace=0,wspace=0)
        plt.margins(0,0)
        #fig = plt.figure(figsize=(18*m,12*n))
        for i in range(1,m*n+1):
            draw(coco,imgids,cat_names,imgroot,m,n,i)
        plt.savefig('static_map.png')
        plt.show()
if __name__ == "__main__":
    annfile = "/data3/Project/CAS_ProjectBird/CDB_ProjectBird_BJCH_YNDC_HHK/COCO/annotations_str/train.json"  #####coco json
    imgroot = "/data3/Project/CAS_ProjectBird/CDB_ProjectBird_BJCH_YNDC_HHK/COCO/train"   ####image_path
    checkCOCO(annfile,imgroot)