"""

 save_the_elephants_survey_A.py

 Convert the .csv file provided for the Save the Elephants Survey A data set to a 
 COCO-camera-traps .json file

"""

#%% Constants and environment

import pandas as pd
import os
import json
import uuid
import time
import humanfriendly
import numpy as np
from tqdm import tqdm

from megadetector.utils.path_utils import find_images

input_base = r'z:/ste_2019_08_drop'
input_metadata_file = os.path.join(input_base,'SURVEY_A.xlsx')

output_base = r'f:/save_the_elephants/survey_a'
output_json_file = os.path.join(output_base,'ste_survey_a.json')
image_directory = os.path.join(input_base,'SURVEY A with False Triggers')
                        
os.makedirs(output_base,exist_ok=True)
assert(os.path.isdir(image_directory))
assert(os.path.isfile(input_metadata_file))

# Handle all unstructured fields in the source data as extra fields in the annotations
mapped_fields = {'No. of Animals in Photo':'num_animals',
                 'No. of new indiviauls (first sighting of new individual)':'num_new_individuals',
                 'Number Adult Males (first sighting of new individual)':'num_adult_males',
                 'Number Adult Females (first sighting of new individual)':'num_adult_females',
                 'Number Adult Unknown (first sighting of new individual)':'num_adult_unknown',
                 'Number Sub-adult Males (first sighting of new individual)':'num_subadult_males',
                 'Number Sub-adult Females (first sighting of new individual)':'num_subadult_females',
                 'Number Sub-adult Unknown (first sighting of new individual)':'num_subadult_unknown',
                 'Number Juvenile (first sighting of new individual)':'num_juvenile',
                 'Number Newborn (first sighting of new individual)':'num_newborn',
                 'Activity':'activity',
                 'Animal ID':'animal_id',
                 'Specific Notes':'notes'}

# photo_type really should be an image property, but there are a few conflicts
# that forced me to handle it as an annotation proprerty
mapped_fields['Photo Type '] = 'photo_type'


#%% Read source data

input_metadata = pd.read_excel(input_metadata_file, sheet_name='9. CT Image')
input_metadata = input_metadata.iloc[2:]

print('Read {} columns and {} rows from metadata file'.format(len(input_metadata.columns),
      len(input_metadata)))


#%% Map filenames to rows, verify image existence

start_time = time.time()

# Maps relative paths to row indices in input_metadata
filenames_to_rows = {}
filenames_with_multiple_annotations = []
missing_images = []

# Build up a map from filenames to a list of rows, checking image existence as we go
for i_row, fn in tqdm(enumerate(input_metadata['Image Name']),total=len(input_metadata)):
    
    # Ignore directories
    if not fn.endswith('.JPG'):
        continue
    
    if fn in filenames_to_rows:
        filenames_with_multiple_annotations.append(fn)
        filenames_to_rows[fn].append(i_row)
    else:
        filenames_to_rows[fn] = [i_row]
        image_path = os.path.join(image_directory, fn)
        if not os.path.isfile(image_path):
            missing_images.append(image_path)
        
elapsed = time.time() - start_time

print('Finished verifying image existence for {} files in {}, found {} filenames with multiple labels, {} missing images'.format(
      len(filenames_to_rows), humanfriendly.format_timespan(elapsed), 
      len(filenames_with_multiple_annotations),len(missing_images)))

        
#%% Make sure the multiple-annotation cases make sense

if False:
    
    #%% 
    
    fn = filenames_with_multiple_annotations[1000] 
    rows = filenames_to_rows[fn]
    assert(len(rows) > 1)
    for i_row in rows:
        print(input_metadata.iloc[i_row]['Species'])
    

#%% Check for images that aren't included in the metadata file

# Enumerate all images
image_full_paths = find_images(image_directory, bRecursive=True)

unannotated_images = []

for iImage, image_path in tqdm(enumerate(image_full_paths),total=len(image_full_paths)):
    relative_path = os.path.relpath(image_path,image_directory)
    if relative_path not in filenames_to_rows:
        unannotated_images.append(relative_path)

print('Finished checking {} images to make sure they\'re in the metadata, found {} unannotated images'.format(
        len(image_full_paths),len(unannotated_images)))


#%% Create CCT dictionaries

images = []
annotations = []
categories = []

image_ids_to_images = {}

category_name_to_category = {}

# Force the empty category to be ID 0
empty_category = {}
empty_category['name'] = 'empty'
empty_category['id'] = 0
category_name_to_category['empty'] = empty_category
categories.append(empty_category)
next_category_id = 1

start_time = time.time()

# i_image = 0; image_name = list(filenames_to_rows.keys())[i_image]
for image_name in tqdm(list(filenames_to_rows.keys())):

    # Example filename:
    #        
    # 'Site 1_Oloisukut_1\Oloisukut_A11_UP\Service_2\100EK113\EK001382.JPG'
    # 'Site 1_Oloisukut_1\Oloisukut_A11_UP\Service_2.1\100EK113\EK001382.JPG'
    img_id = image_name.replace('\\','/').replace('/','_').replace(' ','_')
    
    row_indices = filenames_to_rows[image_name]
    
    # i_row = row_indices[0]
    for i_row in row_indices:
        
        row = input_metadata.iloc[i_row]
        assert(row['Image Name'] == image_name)
        
        timestamp = row['Date'].strftime("%d/%m/%Y")
        station_label = row['Camera Trap Station Label']
        photo_type = row['Photo Type ']
        if isinstance(photo_type,float):
            photo_type = ''
        photo_type = photo_type.strip().lower()
            
        if img_id in image_ids_to_images:
            
            im = image_ids_to_images[img_id]
            assert im['file_name'] == image_name
            assert im['station_label'] == station_label
            
            # There are a small handful of datetime mismatches across annotations
            # for the same image
            # assert im['datetime'] == timestamp
            if im['datetime'] != timestamp:
                print('Warning: timestamp conflict for image {}: {},{}'.format(
                    image_name,im['datetime'],timestamp))
                
        else:
            
            im = {}
            im['id'] = img_id
            im['file_name'] = image_name
            im['datetime'] = timestamp
            im['station_label'] = station_label
            im['photo_type'] = photo_type
            
            image_ids_to_images[img_id] = im
            images.append(im)
    
        species = row['Species']
        
        if (isinstance(species,float) or \
            (isinstance(species,str) and (len(species) == 0))):
            category_name = 'empty'
        else:
            category_name = species
        
        # Special cases based on the 'photo type' field
        if 'vehicle' in photo_type:
            category_name = 'vehicle'
        # Various spellings of 'community'
        elif 'comm' in photo_type:
            category_name = 'human'
        elif 'camera' in photo_type or 'researcher' in photo_type:
            category_name = 'human'
        elif 'livestock' in photo_type:
            category_name = 'livestock'
        elif 'blank' in photo_type:
            category_name = 'empty'
        elif 'plant movement' in photo_type:
            category_name = 'empty'
            
        category_name = category_name.strip().lower()
            
        # Have we seen this category before?
        if category_name in category_name_to_category:
            category_id = category_name_to_category[category_name]['id'] 
        else:
            category_id = next_category_id
            category = {}
            category['id'] = category_id
            category['name'] = category_name
            category_name_to_category[category_name] = category
            categories.append(category)
            next_category_id += 1
        
        # Create an annotation
        ann = {}        
        ann['id'] = str(uuid.uuid1())
        ann['image_id'] = im['id']    
        ann['category_id'] = category_id
        
        # fieldname = list(mapped_fields.keys())[0]
        for fieldname in mapped_fields:
            target_field = mapped_fields[fieldname]
            val = row[fieldname]
            if isinstance(val,float) and np.isnan(val):
                val = ''
            else:
                val = str(val).strip()
            ann[target_field] = val
            
        annotations.append(ann)
        
    # ...for each row
                
# ...for each image
    
print('Finished creating CCT dictionaries in {}'.format(
      humanfriendly.format_timespan(elapsed)))
    

#%% Create info struct

info = {}
info['year'] = 2019
info['version'] = 1
info['description'] = 'Save the Elephants Survey A'
info['contributor'] = 'Save the Elephants'


#%% Write output

json_data = {}
json_data['images'] = images
json_data['annotations'] = annotations
json_data['categories'] = categories
json_data['info'] = info
json.dump(json_data, open(output_json_file, 'w'), indent=2)

print('Finished writing .json file with {} images, {} annotations, and {} categories'.format(
        len(images),len(annotations),len(categories)))


#%% Validate output

from megadetector.data_management.databases import integrity_check_json_db

options = integrity_check_json_db.IntegrityCheckOptions()
options.baseDir = image_directory
options.bCheckImageSizes = False
options.bCheckImageExistence = False
options.bFindUnusedImages = False
    
sortedCategories, data = integrity_check_json_db.integrity_check_json_db(output_json_file,options)


#%% Preview labels

from megadetector.visualization import visualize_db
from megadetector.data_management.databases import integrity_check_json_db

viz_options = visualize_db.DbVizOptions()
viz_options.num_to_visualize = 1000
viz_options.trim_to_images_with_bboxes = False
viz_options.add_search_links = True
viz_options.sort_by_filename = False
viz_options.parallelize_rendering = True
html_output_file,image_db = visualize_db.visualize_db(db_path=output_json_file,
                                                        output_dir=os.path.join(output_base,'preview'),
                                                        image_base_dir=image_directory,
                                                        options=viz_options)
os.startfile(html_output_file)


#%% Scrap

if False:

    pass
    
    #%% Find unique photo types
    
    annotations = image_db['annotations']
    photo_types = set()
    for ann in tqdm(annotations):
        photo_types.add(ann['photo_type'])
