#!/usr/bin/env python3

import os

from echo1_coco_builder.annotations_builder import CocoAnnotationsBuilder
from deps import getdeps # PYTHONPATH=$DATA_REPOSITORY/utils/python
from nabirds import NABirds
from cub import CUB
import parts

PBOXSZ = 20
part_id_base = 2
CAPTION_SEP = '.'

def get_caption(image_key,birds):
    """Generate caption string for the image_key and the map (keyed to
    category_id) of token (beg,end) position tuples for each section of the
    caption.

    """
    
    caption = 'bird'
    indices = { 1: (0,len(caption)) }
               
    for part_num in range(len(birds.image_parts[image_key])):
        nab_id = birds.image_class_labels[image_key] # bird species class id (str)
        [px,py,pv] = birds.image_parts[image_key][part_num]

        if pv==0: # Part not visible
            continue
        part_descr = parts.NABParts.parts[part_num].description(int(nab_id))
        caption += CAPTION_SEP + ' ' # Add separators ...
        # cf. https://github.com/microsoft/GLIP/blob/91f20921e691922167d8fa94fb55fa7a25ffe0ff/maskrcnn_benchmark/data/datasets/od_to_grounding.py#L46
        start_index = len(caption)
        caption += part_descr
        end_index = len(caption)
        indices[part_id_base+part_num] = (start_index,end_index)

    caption += '.'

    return caption, indices


def process_parts(part_id_base,birds,coco_data):
    # Process part categories    
    for class_id,class_name in birds.part_names.items():
        coco_data.add_category({
            "id"  : int(class_id) + part_id_base,
            "name": class_name # Species-generic class name
        })

def process_images(image_keys, birds, coco_data):
    image_id = 0 # COCO wants ints; no non-hacky way to associate NABird's hash :(

    image_ids = {}
    image_tokens = {}
    
    for image_key in image_keys:
        image_ids[image_key] = image_id
        image_width,image_height = birds.image_sizes[image_key]
        caption,indices = get_caption(image_key,birds)

        # Captions are generated in part_num sorted order, so let's make sure the
        # tokens positions are in the same order
        tokens = [[indices[part]] for part in sorted(indices.keys())]
        image_tokens[image_key] = indices
        coco_data.add_image({
            "id"       : image_id,
            "file_name": birds.image_paths[image_key],
            "width"    : image_width,
            "height"   : image_height,
            "caption"  : caption,
            "tokens"   : tokens
        })
        image_id = image_id + 1

    return image_ids,image_tokens

def process_bboxes(image_keys, image_ids, image_tokens, part_base, birds,
                   coco_data):

    annotation_id = 0
    for image_key in image_keys:
        nab_id = birds.image_class_labels[image_key] # bird species class id (str)
        # Add bird bbox
        bird_bbox = birds.image_bboxes[image_key] # [x,y,width,height]
        coco_data.add_annotation({
            "id"         : annotation_id,
            "image_id"   : image_ids[image_key],
            "category_id": 1,
            "bbox"       : bird_bbox,
            "iscrowd"    : 0,
            "area"       : bird_bbox[2]*bird_bbox[3], # w * h
            "tokens"     : [image_tokens[image_key][1]]
            })
        annotation_id = annotation_id + 1
        
        # Add visible part bboxes
        for part_num in range(len(birds.image_parts[image_key])):
            [px,py,pv] = birds.image_parts[image_key][part_num]

            if pv==0: # Part not visible
                continue

            part_bbox = [px-PBOXSZ, py-PBOXSZ, 2*PBOXSZ, 2*PBOXSZ]

            part_descr = parts.NABParts.parts[part_num].description(int(nab_id))
            category_id = part_id_base+part_num
            coco_data.add_annotation({
                "id"         : annotation_id,
                "image_id"   : image_ids[image_key],
                "category_id": category_id,
                "bbox"       : part_bbox,
                "iscrowd"    : 0,
                "area"       : part_bbox[2]*part_bbox[3], # w * h
                "name"       : part_descr,
                "tokens"     : [image_tokens[image_key][category_id]]
            })
            annotation_id = annotation_id + 1
        

def process_categories(output_path, birds):
    coco_data = CocoAnnotationsBuilder()
    coco_data.add_category({
        "id"  : 1,
        "name": "bird"
    })

    process_parts(part_id_base, birds, coco_data)
    with open(output_path,'w') as fd:
        fd.write(coco_data.get())

def process_split(output_path, image_keys, birds):
    
    coco_data = CocoAnnotationsBuilder()

    coco_data.add_info({
        "year"        : 2022,
        "version"     : "0.1",
        "description" : "Bird part points given bounding boxes (40px h/w)",
        "contributor" : "NABirds, Jerod Weinman",
        "date_created": "2022/07/15"
    })

    # bird + part categories
    coco_data.add_category({"id"  : 1,"name": "bird"})
    process_parts(part_id_base, birds, coco_data)

    image_ids,image_tokens = process_images(image_keys, birds, coco_data)

    process_bboxes(image_keys, image_ids, image_tokens, part_id_base,
                   birds, coco_data)

    with open(output_path,'w') as fd:
        fd.write(coco_data.get())
        
        
if __name__ == '__main__':
    
    deps = getdeps()
    birds = NABirds( deps['nab'],
                     os.path.join(deps['split'],'train_val_test_split.txt'),
                     'images')
    parts.Part.cub = CUB( deps['cub'],
                          os.path.join(deps['map'],'nabird_cub.csv'))
    process_split('../data/train.json', birds.train_images, birds)
    process_split('../data/val.json', birds.val_images, birds)
    process_split('../data/test.json', birds.test_images, birds)
    process_categories('../data/categories.json',birds)
