# Script for creating json with annotations in the COCO format out of the Cell Tracking Challenge format of annotations
# Author: Kristyna Janku

import json
import cv2
import glob
import numpy as np
from os import listdir, path
from numba import njit


def get_image_info(videos_folder, video, removed_frames):
    images = []

    video_path = path.join(videos_folder, video)

    # Filter out frames that were removed due to incomplete SEG annotations
    frame_paths = sorted(frame for frame in glob.glob(path.join(video_path, '*.png'))
                         if ''.join(filter(str.isdigit, path.basename(frame))) not in removed_frames)

    for filepath in frame_paths:
        img = cv2.imread(filepath)
        height, width, _ = img.shape
        images.append(
            {
                'file_name': path.basename(filepath),
                'height': height,
                'id': int(''.join(filter(str.isdigit, path.basename(filepath)))),
                'width': width
            }
        )

    return images


def get_bounding_box(mask, padding=5, min_area=6):
    # Get bounding box of object from maximal contour
    contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
    areas = [cv2.contourArea(c) for c in contours]
    max_index = np.argmax(areas)
    max_contour = contours[max_index]
    x, y, width, height = cv2.boundingRect(max_contour)

    if (width * height) < min_area:
        return 0, 0, 0, 0

    im_height, im_width = mask.shape

    # Add padding
    x1 = max(x - padding, 0)
    y1 = max(y - padding, 0)
    x2 = min(x + width + padding, im_width)
    y2 = min(y + height + padding, im_height)
    width = x2 - x1
    height = y2 - y1

    return x1, y1, width, height


@njit
def tracking_to_segmentation_labels(tracked_ids, track_image, seg_image):
    track_to_seg_labels = {}

    # Look for tracking markers in tracking image
    for i in range(track_image.shape[0]):
        for j in range(track_image.shape[1]):
            track_id = track_image[i, j]
            if track_id in tracked_ids:
                # Get segmentation label of tracked object
                seg_label = seg_image[i, j]
                if seg_label > 0 and track_id not in track_to_seg_labels:
                    track_to_seg_labels[track_id] = seg_label
        if len(tracked_ids) == len(track_to_seg_labels):
            break

    return track_to_seg_labels


def visualize_annotations(seg_image, frame_annotations):
    im_labeled = cv2.equalizeHist(cv2.convertScaleAbs(seg_image))
    im_labeled = cv2.merge([im_labeled, im_labeled, im_labeled])
    for ann in frame_annotations:
        bbox = ann['bbox']
        im_labeled = cv2.rectangle(im_labeled, (bbox[0], bbox[1]), (bbox[0] + bbox[2], bbox[1] + bbox[3]),
                                   (0, 255, 255), 3)
    cv2.imshow('Labeled', im_labeled)
    cv2.waitKey(0)


def get_annotations(videos_folder, video, mot, visualize=False):
    annotations = []

    gt_annotation_path = path.join(videos_folder, f'{video}_GT')
    gt_segmentation_folder_path = path.join(gt_annotation_path, 'SEG')
    gt_tracking_folder_path = path.join(gt_annotation_path, 'TRA')

    st_segmentation_folder_path = path.join(videos_folder, f'{video}_ST', 'SEG')

    # Load numbers of all frames from filenames (only PNG files)
    frames = sorted([''.join(filter(str.isdigit, frame)) for frame
                     in listdir(path.join(videos_folder, video)) if frame.endswith('.png')])

    removed_frames = []

    for frame in frames:
        frame_annotations = []

        # Load segmented image, gold truth if available, silver truth otherwise
        seg = glob.glob(path.join(gt_segmentation_folder_path, f'*{frame}*'))
        if len(seg) == 0:
            seg = glob.glob(path.join(st_segmentation_folder_path, f'*{frame}*'))
        seg_image = cv2.imread(seg[0], cv2.IMREAD_ANYDEPTH)

        # Load tracking image with object markers
        track = glob.glob(path.join(gt_tracking_folder_path, f'*{frame}*'))
        track_image = cv2.imread(track[0], cv2.IMREAD_ANYDEPTH)

        # Find tracked ids and map them to segmentation labels
        tracked_ids = np.unique(track_image)[1:]
        track_to_seg_labels = tracking_to_segmentation_labels(tracked_ids, track_image, seg_image)

        # Check if segmentation contains all tracked objects
        len_diff = len(tracked_ids) - len(track_to_seg_labels)
        if len_diff > 0:
            print(f'Video: {video} Frame: {frame} Missing {len_diff} tracked objects in seg')

            # Remove frame for MOT annotations (unannotated objects would confuse detector)
            if mot:
                removed_frames.append(frame)
                continue

        for obj_id in track_to_seg_labels:
            # Mask only pixels with current object, multiply by 1 to change True/False to 1/0
            mask = (seg_image == track_to_seg_labels[obj_id]) * 1

            # Convert to 8-bit image to enable further transformations
            mask = cv2.convertScaleAbs(mask)

            x, y, width, height = get_bounding_box(mask, padding=5, min_area=6)
            area = width * height

            # Ignore annotations with area less than min_area (width, height set to 0 from get_bounding_box)
            if area > 0:
                frame_annotations.append(
                    {
                        'area': area,
                        'bbox': [x, y, width, height],
                        'category_id': 1,
                        'id': int(obj_id) - 1,
                        'ignore': 0,
                        'image_id': int(frame),
                        'iscrowd': 0
                    }
                )

        if visualize:
            visualize_annotations(seg_image, frame_annotations)

        annotations += frame_annotations
    return annotations, removed_frames


def create_json(dataset_name, subset, videos_folder, videos, mot, visualize=False):
    for video in videos:
        # Add category info
        categories = [
                {
                    'id': 1,
                    'name': 'cell',
                    'supercategory': 'cell'
                }
            ]

        annotations, removed_frames = get_annotations(videos_folder, video, mot, visualize)

        coco_dict = {
            'categories': categories,
            'annotations': annotations,
            'images': get_image_info(videos_folder, video, removed_frames),
        }

        ann_folder = path.join(videos_folder, '../../annotation')

        with open(path.join(ann_folder, f'{"mot" if mot else "sot"}_{dataset_name}_{subset}_image_{video}.json'), 'w') as ann_json_file:
            json.dump(coco_dict, ann_json_file, indent=4, sort_keys=True)


if __name__ == '__main__':
    def_dataset_name = 'BF-C2DL-HSC'
    def_subset = 'test'
    def_videos_folder = path.join(path.expanduser('~'), f'data/{def_dataset_name}/raw_data/{def_subset}')
    def_videos = ['02']
    create_json(def_dataset_name, def_subset, def_videos_folder, def_videos, mot=True, visualize=False)

