# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Convert COCO annotations to KITTI format"""

import os
import numpy as np

from pycocotools.coco import COCO
from pycocotools import mask


def find_corners_2d(np_array):
    """Find corners in a 2D binary mask."""
    result = np.where(np_array == np.amax(np_array))
    x1 = np.min(result[0])
    x2 = np.max(result[0])
    y1 = np.min(result[1])
    y2 = np.max(result[1])
    return y1, x1, y2, x2


def convert_coco_to_kitti(annotations_file, output_dir, refine_box=False):
    """Function to convert COCO annotations to KITTI format.

    Args:
        annotations_file (string): Path to the COCO annotation file.
        output_dir (string): Directory to output the KITTI files
    """
    if not os.path.isfile(annotations_file):
        print("Annotation file doenot exsit. Please check the path")

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    coco = COCO(annotations_file)
    category_ids = coco.getCatIds()
    categories = coco.loadCats(category_ids)

    category_id_dict = {}
    for category in categories:
        category_id_dict[category['id']] = category['name']
    for img, ann in coco.imgs.items():
        annotation_ids = coco.getAnnIds(imgIds=[img], catIds=category_ids)
        if len(annotation_ids) > 0:
            img_fname = ann['file_name']
            label_fname = img_fname.split('.')[0]
            with open(os.path.join(output_dir, f'{label_fname}.txt'), 'w', encoding='utf-8') as label_file:
                annotations = coco.loadAnns(annotation_ids)
                for annotation in annotations:
                    bbox = annotation['bbox']
                    bbox = [bbox[0], bbox[1], bbox[2] + bbox[0], bbox[3] + bbox[1]]
                    if refine_box and annotation.get('segmentation', None):
                        rle = annotation['segmentation']
                        binary_mask = mask.decode(rle)
                        bbox = find_corners_2d(binary_mask)
                    bbox = [str(b) for b in bbox]
                    catname = category_id_dict[annotation['category_id']]
                    out_str = [catname.replace(" ", "") + ' ' + ' '.join(['0'] * 3) + ' ' + ' '.join(list(bbox)) + ' ' + ' '.join(['0'] * 7) + '\n']
                    label_file.write(out_str[0])
