import numpy as np
import json
from pathlib import Path, PosixPath
from pycocotools.coco import COCO

def min_index(arr1, arr2):
    """
    Find a pair of indexes with the shortest distance.

    Args:
        arr1: (N, 2).
        arr2: (M, 2).
    Return:
        a pair of indexes (tuple)
    """
    dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
    return np.unravel_index(np.argmin(dis, axis=None), dis.shape)


def merge_multi_segment(segments):
    """
    Merge multi segments to one list.
    Find coordinates with min distance between each segment,
    then connect these coordinates with one thin line to merge all
    segments into one.

    Args:
        segments (List(List)): original segmentations in coco's json file
            like [segmentation1, segmentation2, ...], where
            each segmentation is a list of coordinates
    """
    s = []
    segments = [np.array(i).reshape(-1,2) for i in segments]
    idx_list = [[] for _ in range(len(segments))]

    # record the indexes with the min distance between each segment
    for i in range(1, len(segments)):
        idx1, idx2 = min_index(segments[i - 1, segments[i]])
        idx_list[i - 1].append(idx1)
        idx_list[i].append(idx2)

    # use two round to connect all the segments
    for k in range(2):
        # forward connection
        if k == 0:
            for i, idx in enumerate(idx_list):
                # middle segments have two indexes
                # reverse the index of middle segments
                if len(idx) == 2 and idx[0] > idx[1]:
                    idx = idx[::-1]
                    segments[i] = segments[i][::-1, :]

                segments[i] = np.roll(segments[i], -idx[0], axis=0)
                segments[i] = np.concatenate([segments[i], segments[i][:1]])

                # deal with the first segment and the last one
                if i in [0, len(idx_list) - 1]:
                    s.append(segments[i])
                else:
                    idx = [0, idx[1] - idx[0]]
                    s.append(segments[i][idx[0]:idx[1] + 1])

        else:
            for i in range(len(idx_list) - 1, -1, -1):
                if i not in [0, len(idx_list) - 1]:
                    idx = idx_list[i]
                    nidx = abs(idx[1] - idx[0])
                    s.append(segments[i][nidx:])
    return s


def get_yolo_labels(path2json, use_segment=False):
    if not isinstance(path2json, PosixPath):
        path2json = Path(path2json)
    path2labels = path2json.parents[0] / "labels"
    path2labels.mkdir(parents=True, exist_ok=True)

    coco = COCO(path2json)

    img2anns = {}
    for ann in coco.dataset['annotations']:
        img_id = ann['image_id']
        if img_id not in img2anns:
            img2anns[img_id] = [ann]
        else:
            img2anns[img_id].append(ann)

    id2img = {img["id"]: img for img in coco.dataset["images"]}

    for img_id, anns in img2anns.items():
        img = id2img[img_id]
        h, w, f = img['height'], img['width'], img['file_name']

        bboxes = []
        segments = []
        for ann in anns:
            if ann['iscrowd']:
                continue
            # coco box format: [top left x, top left y, width, height]
            box = np.array(ann['bbox'], dtype=np.float64)
            box[:2] += box[2:] / 2  # center coordinates
            box[[0, 2]] /= w  # normalize x
            box[[1, 3]] /= h  # normalize y
            if box[2] <= 0 or box[3] <= 0:
                continue

            cls = ann['category_id'] - 1
            box = [cls] + box.tolist()
            if box not in bboxes:
                bboxes.append(box)

            # segmentation?
            if use_segment:
                if len(ann['segmentation']) > 1:
                    s = merge_multi_segment(ann['segmentation'])
                    s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
                else:
                    s = [j for i in ann['segmentation'] for j in i]  # all segments concatenated
                    s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
                s = [cls] + s
                if s not in segments:
                    segments.append(s)

        # write
        with open((path2labels / f).with_suffix('.txt'), 'a') as file:
            for i in range(len(bboxes)):
                line = *(segments[i] if use_segment else bboxes[i]),
                file.write(('%g ' * len(line)).rstrip() % line + '\n')