# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/05_labelme2coco.ipynb (unless otherwise specified).

__all__ = ['read_content_labelme', 'visualize_polygons', 'labelme_visual', 'random_labelme_visual', 'labelme2coco',
           'segmentation_mask', 'coco_random_annot_visual', 'get_coco_segmentation', 'get_concat_h', 'get_concat_v',
           'draw_grid', 'keypoints_labelme_clean', 'process_parts', 'labelme2coco_cli', 'visualize_keypoints']

# Cell
import os
import json
import random
import glob
import yaml
import torch
import cv2
import numpy as np
import skimage.io as io
import matplotlib.pyplot as plt
from tqdm import tqdm
from labelme import utils
from fastscript import call_parse, Param
from pycocotools import mask as CocoMask
from pycocotools import coco
from PIL import Image, ImageDraw, ImageFont
from fastcore.transform import store_attr
from pycocotools.coco import COCO

from .perspective_correction import order_points
from .utils import load_json, xywh_to_xyxy, xyxy_to_xywh, load_yaml

# Cell
def read_content_labelme(file):
    labels = [i["label"] for i in file["shapes"]]
    polygons = [i["points"] for i in file["shapes"]]
    return labels, polygons



def visualize_polygons(img, polygons: list, labels: list, text: bool=False, font_loc: str="arial.ttf"):
    """input pil image
    """
    font = ImageFont.truetype(font_loc, 10)
    draw = ImageDraw.Draw(img, "RGBA")
    for k, polygon in enumerate(polygons):
        draw.polygon([tuple(i) for i in polygon], fill=None, outline=(0, 255, 0, 255))
        if text:
            draw.text(polygon[0], labels[k], font=font, fill=(0, 0, 0, 0))



def labelme_visual(annot_file: str, font_loc: str="arial.ttf"):
    """ visualize an labelme json file
    """
    file = load_json(annot_file)
    labels, polygons = read_content_labelme(file)
    img = Image.open(annot_file.rsplit("/", 1)[0] + "/" + file["imagePath"])
    visualize_polygons(img, polygons, labels, text=True, font_loc=font_loc)
    img.thumbnail((512, 512))
    return img


# Cell
def random_labelme_visual(annoto_files: list, font_loc: str="assets/arial.ttf"):
    """ Randomly visualize labelme annotation file

    """
    file_name = annot_files[np.random.randint(len(annot_files))]  # Randomly take a file
    img = labelme_visual(file_name, font_loc=font_loc)
    return img


# Cell

class labelme2coco(object):
    def __init__(self, labelme_json: list, labels_loc: str, save_json_path="./coco.json"):
        """
        Source: https://github.com/Tony607/labelme2coco/blob/master/labelme2coco.py
        :param labelme_json: the list of all labelme json file paths
        :param save_json_path: the path to save new json
        """
        store_attr(self, "labelme_json, save_json_path, labels_loc")
        self.images = []
        self.categories = []
        self.annotations = []
        self.label = []
        self.annID = 1
        self.height = 0
        self.width = 0

        self.label_data = load_yaml(self.labels_loc)
        self.labels = self.label_data["labels"]
        print(self.labels)

    def data_transfer(self):
        for num, json_file in tqdm(enumerate(self.labelme_json)):
            with open(json_file, "r") as fp:
                data = json.load(fp)
                self.images.append(self.image(data, num))
                for shapes in data["shapes"]:
                    label = shapes["label"]
                    if self.labels is not None:
                        if label not in self.labels:
                            continue
                    if label not in self.label:
                        self.label.append(label)
                    points = shapes["points"]
                    points = order_points(np.asarray(points)).tolist()
                    self.annotations.append(self.annotation(points, label, num))
                    self.annID += 1

        # Sort all text labels so they are in the same order across data splits.
        self.label.sort()
        for label in self.label:
            self.categories.append(self.category(label))
        for annotation in self.annotations:
            annotation["category_id"] = self.getcatid(annotation["category_id"])

    def image(self, data, num):
        image = {}
        img = utils.img_b64_to_arr(data["imageData"])
        height, width = img.shape[:2]
        img = None
        image["height"] = height
        image["width"] = width
        image["id"] = num
        image["file_name"] = data["imagePath"].split("/")[-1]

        self.height = height
        self.width = width

        return image

    def category(self, label):
        category = {}
        category["supercategory"] = label
        category["id"] = len(self.categories)
        category["name"] = label

        if 'keypoints' in self.label_data.keys() and len(self.label_data["keypoints"])>0:
            category["keypoints"] = self.label_data["keypoints"]

        if 'skeleton' in self.label_data.keys() and len(self.label_data["skeleton"])>0:
            category["skeleton"] = self.label_data["skeleton"]
        return category

    def annotation(self, points, label, num):
        annotation = {}
        contour = np.array(points)
        x = contour[:, 0]
        y = contour[:, 1]
        area = 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
        annotation["segmentation"] = [list(np.asarray(points).flatten())]
        annotation["iscrowd"] = 0
        annotation["area"] = area
        annotation["image_id"] = num

        annotation["bbox"] = list(map(float, self.getbbox(points)))

        annotation["category_id"] = label  # self.getcatid(label)

        annotation["id"] = self.annID

        if 'keypoints' in self.label_data.keys() and len(self.label_data["keypoints"])>0:
            annotation["num_keypoints"] = len(points)
            annotation["keypoints"] = sum([[int(x),int(y),2] for [x,y] in points],[])


        return annotation

    def getcatid(self, label):
        for category in self.categories:
            if label == category["name"]:
                return category["id"]
        print("label: {} not in categories: {}.".format(label, self.categories))
        exit()
        return -1

    def getbbox(self, points):
        polygons = points
        mask = self.polygons_to_mask([self.height, self.width], polygons)
        return self.mask2box(mask)

    def mask2box(self, mask):

        index = np.argwhere(mask == 1)
        rows = index[:, 0]
        clos = index[:, 1]

        left_top_r = np.min(rows)  # y
        left_top_c = np.min(clos)  # x

        right_bottom_r = np.max(rows)
        right_bottom_c = np.max(clos)

        return [
            left_top_c,
            left_top_r,
            right_bottom_c - left_top_c,
            right_bottom_r - left_top_r,
        ]

    def polygons_to_mask(self, img_shape, polygons):
        mask = np.zeros(img_shape, dtype=np.uint8)
        mask = Image.fromarray(mask)
        xy = list(map(tuple, polygons))
        ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
        mask = np.array(mask, dtype=bool)
        return mask

    def polygon2rle(self, polygons):
        mask = self.polygons_to_mask([self.height, self.width], polygons)
        segm = segmentation_mask(mask)
        assert len([segm]) == 1
        return segm

    def data2coco(self):
        data_coco = {}
        data_coco["images"] = self.images
        data_coco["categories"] = self.categories
        data_coco["annotations"] = self.annotations
        return data_coco

    def process_json(self):
        print("save coco json")
        self.data_transfer()
        self.data_coco = self.data2coco()

    def save_json(self):

        print(self.save_json_path)
        os.makedirs(
            os.path.dirname(os.path.abspath(self.save_json_path)), exist_ok=True
        )
        json.dump(self.data_coco, open(self.save_json_path, "w"), indent=4)

def segmentation_mask(mask):
    contours, hierarchy, _ = cv2.findContours(mask.astype(np.uint8),  cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    segmentation = []

    for contour in contours:
        contour = contour.flatten().tolist()
        segmentation.append(contour)
        if len(contour) > 4:
            segmentation.append(contour)
    return segmentation

# Cell
def coco_random_annot_visual(data, img_loc: str):
    ids = list(data.imgs.keys())
    img_id = ids[np.random.randint(len(ids))]
    ann_ids = data.getAnnIds(imgIds=img_id)
    print(ann_ids)
    target = data.loadAnns(ann_ids)
    path = data.loadImgs(img_id)[0]["file_name"]
    img_pil = Image.open(os.path.join(img_loc,path))
    label = [i["category_id"] for i in target]
    boxes = np.asarray([i["bbox"] for i in target])
    print(boxes)
    boxes_xyxy = xywh_to_xyxy(boxes)
    # colors = compute_colors_for_labels(label)
    img_seg = get_coco_segmentation(data, target)
    draw_grid(img_seg, boxes_xyxy, label, outline="red")
    img = get_concat_h(img_pil, img_seg)
    img.thumbnail((1024, 512))
    return img

def get_coco_segmentation(data, anns):
    # seg = [i["segmentation"] for i in anns]
    all_masks = sum([data.annToMask(i) for i in anns])
    all_masks[all_masks > 1] = 1
    return Image.fromarray(all_masks * 255)


def get_concat_h(im1, im2):
    """concats two images horizontally

    https://note.nkmk.me/en/python-pillow-concat-images/
    """
    dst = Image.new("RGB", (im1.width + im2.width, im1.height))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (im1.width, 0))
    return dst


def get_concat_v(im1, im2):
    """ concats two images vertically

    https://note.nkmk.me/en/python-pillow-concat-images/
    """
    dst = Image.new("RGB", (im1.width, im1.height + im2.height))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (0, im1.height))
    return dst

def draw_grid(image, bbox, labels, outline="white", input_format="xyxy"):
    """draws rectangles on the image given by bbox

    image: PIL image
    bbox: numpy with each box representing the format defined by "format"
    outline: color of the bbox
    input_format: "xxyy" or "xyhw": use one of this

    """
    draw = ImageDraw.Draw(image)

    if isinstance(bbox, torch.Tensor):
        bbox = bbox.tolist()

    if input_format == "xywh":
        bbox = xywh_to_xyxy(bbox)

    for num, i in enumerate(bbox):
        x0, y0, x1, y1 = i
        label = labels[num]
        draw.rectangle([x0, y0, x1, y1], outline=outline, width=5)
        draw.text((x0, y0), str(float(label)))
    return image


# Cell
def keypoints_labelme_clean(files: list, labels: list, points: int=4):
    cleaned_json_files = []
    for file in tqdm(files):
        js = load_json(file)
        if len(js["shapes"]) ==0:
            continue
        for annot in js["shapes"]:
            if annot["label"] in labels and len(annot["points"]) ==points:
                cleaned_json_files.append(file)
    return cleaned_json_files

# Cell
def process_parts(files, labels,  save_json_path):
    data = labelme2coco(files, labels, save_json_path)
    data.process_json()
    data.save_json()

@call_parse
def labelme2coco_cli(labelme_json_location: Param("json files location of labelme", type=str)=None,
                 labels_loc: Param("labels yaml file location",type=str)=None,
                 save_json: Param("path to save coco json",type=str)='data/',
                 train_ratio: Param("fraction of images to be used for train split",type=float)=0.9,
                 seed: Param("random seed for train val split",type=int)=50):
    json_files = glob.glob(os.path.join(labelme_json_location,'*.json'))
    yaml_data = load_yaml(labels_loc)

    if "keypoints" in yaml_data.keys():
        print("only valid keypoint files are considered")
        json_files = keypoints_labelme_clean(json_files, yaml_data["labels"], len(yaml_data["keypoints"]))

    random.seed(50)
    dim = int(0.9*len(json_files))
    random.shuffle(json_files)
    train_json_files = json_files[:dim]
    val_json_files = json_files[dim:]

    if not os.path.exists(save_json): os.makedirs(save_json)

    process_parts(train_json_files, labels_loc, save_json_path=os.path.join(save_json,'train_coco.json'))
    process_parts(val_json_files, labels_loc, save_json_path=os.path.join(save_json,'val_coco.json'))

# Cell
## Visualization keypoints
def visualize_keypoints(annFile: str, img_dir:str, cat_names: list):
    coco_kps=COCO(annFile)
    catIds = coco_kps.getCatIds(catNms=cat_names);
    imgIds = coco_kps.getImgIds(catIds=catIds );
    img = coco_kps.loadImgs(imgIds[np.random.randint(0,len(imgIds))])[0]
    I = io.imread(os.path.join(img_dir, img["file_name"]))
    # load and display keypoints annotations
    plt.imshow(I); plt.axis('off')
    ax = plt.gca()
    annIds = coco_kps.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
    anns = coco_kps.loadAnns(annIds)
    coco_kps.showAnns(anns)