"""

 ubc_to_json.py

 Convert the .csv file provided for the UBC data set to a 
 COCO-camera-traps .json file

 Images were provided in eight folders, each of which contained a .csv
 file with annotations.  Those annotations came in two slightly different 
 formats, the two formats corresponding to folders starting with "SC_" and 
 otherwise.

"""

#%% Constants and environment

import pandas as pd
import os
import json
import uuid
import numpy as np
import shutil

from tqdm import tqdm
from PIL import Image

from megadetector.visualization import visualize_db
from megadetector.data_management.databases import integrity_check_json_db
from megadetector.utils.path_utils import find_images

input_base = r'e:\ubc'
assert(os.path.isdir(input_base))

output_base = r'f:\data_staging\ubc'
output_json_file = os.path.join(output_base,'ubc.json')
file_list_file = os.path.join(output_base,'all_files.txt')

os.makedirs(output_base,exist_ok=True)

# Map Excel column names - which vary a little across spreadsheets - to a common set of names
mapped_fields = {"Survey.Name" : "survey_name",
                 "project_id": "survey_name",
                 "Camera.Name": "camera_name",
                 "station_id": "camera_name",
                 "Media.Filename": "filename",
                 "orig_file": "filename",
                 "timestamp_pst": "datetime",
                 "Date.Time": "datetime",
                 "Species": "species",
                 "latin_name": "species",
                 "common.name": "common_name",
                 "common_names": "common_name",
                 "Sighting.Quantity": "species_count"
                }

category_mappings = {
    'bird_spp.':'unknown_bird',
    'dog_dog':'dog',
    'hiker_hiker':'hiker',
    'quad_quad':'quad',
    'skier_skier':'skier',
    'snowshoer_snowshoer':'showshoer',
    'quad_quad':'quad'
    }

target_fields = ['species_count','group_count','behaviour']


#%% Enumerate images

# Load from file if we've already enumerated
if os.path.isfile(file_list_file):
    with open(file_list_file,'r') as f:
        files = f.readlines()
    files = [s.strip() for s in files]
    image_full_paths = files
    print('Loaded {} images from {}'.format(len(image_full_paths),file_list_file))
else:
    image_full_paths = find_images(input_base, recursive=True)
    with open(file_list_file,'w') as f:
        for line in image_full_paths:
            f.write(line + '\n')        
    print('Enumerated {} images from {}'.format(len(image_full_paths),input_base))
    
image_full_paths_set = set(image_full_paths)
image_relative_paths = [os.path.relpath(fn,input_base) for fn in image_full_paths]
image_relative_paths_set = set(image_relative_paths)


#%% Create CCT dictionaries

annotations = []
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
empty_category['common_name'] = 'empty'
category_name_to_category['empty'] = empty_category
next_category_id = 1

latin_to_common = {}

folders = os.listdir(input_base)

# To simplify debugging of the loop below
i_folder = 0; folder = folders[i_folder]; 


##%% Create CCT dictionaries (loop)

invalid_images = []

for i_folder,folder in enumerate(folders):
        
    ##%% 
    
    print('\nProcessing folder {} of {}: {}'.format(i_folder,len(folders),folder))
    
    filenames_to_rows = {}
    filenames_with_multiple_annotations = []
    missing_images = []
    image_directory = os.path.join(input_base, folder)
    files = os.listdir(image_directory)
    files = list(filter(lambda f: f.endswith('.csv'), files))
    input_metadata_file = os.path.join(input_base, folder, files[0])
    assert(os.path.isfile(input_metadata_file))
    
    # Read source data for this folder
    input_metadata = pd.read_csv(input_metadata_file)
    
    # Rename columns
    input_metadata.rename(columns=mapped_fields, inplace=True)
    print('Read {} columns and {} rows from metadata file'.format(len(input_metadata.columns), len(input_metadata)))
    
    if folder.startswith("SC_"):
        # Folder name is the first two characters of the filename
        #
        # Create relative path names from the filename itself
        input_metadata['image_relative_path'] = input_metadata['filename'].apply(
            lambda x: os.path.join(folder, x[0:2], x.replace(".JPG", ".jpg")))
    else:
        # Folder name is the camera name
        #
        # Create relative path names from camera name and filename    
        input_metadata['image_relative_path'] = input_metadata[['camera_name', 'filename']].apply(
            lambda x: os.path.join(folder, x[0], x[1]), axis = 1)

    # Which of our images are in the spreadsheet?
    # i_row = 0; fn = input_metadata['image_relative_path'][i_row]        
    for i_row, image_relative_path in enumerate(input_metadata['image_relative_path']):
        
        if image_relative_path in filenames_to_rows:
            filenames_with_multiple_annotations.append(image_relative_path)
            filenames_to_rows[image_relative_path].append(i_row)
        else:
            filenames_to_rows[image_relative_path] = [i_row]
            image_full_path = os.path.join(input_base, image_relative_path)
            
            if not image_full_path in image_full_paths_set:
                assert image_relative_path not in image_relative_paths_set
                missing_images.append(image_full_path)
        
    print('Finished verifying image existence for {} files in {} rows\nFound {} filenames with multiple labels, {} missing images'.format(
        len(input_metadata), len(filenames_to_rows), 
        len(filenames_with_multiple_annotations), len(missing_images)))
    
    
    ##%% Check for images that aren't included in the metadata file

    # Find all the images in this folder
    image_relative_paths_this_folder = [s for s in image_relative_paths if s.startswith(folder)]

    # Which of these aren't in the spreadsheet?
    annotated_files_this_folder = list(filenames_to_rows.keys())
    annotated_files_this_folder_set = set(annotated_files_this_folder)
    unannotated_images = [s for s in image_relative_paths_this_folder if s not in annotated_files_this_folder_set]

    print('Found {} unannotated images (of {}) in this folder'.format(
            len(unannotated_images),len(image_relative_paths_this_folder)))
    
    
    ##%% Create entries in CCT dictionaries
    
    image_relative_path = list(filenames_to_rows.keys())[0]
    
    for image_relative_path in list(filenames_to_rows.keys()):
        
        # Only process images we have on disk
        if image_relative_path not in image_relative_paths_set:
            continue
        
        image_full_path = os.path.join(input_base,image_relative_path)
        
        # This is redundant, but doing this for clarity, at basically no performance
        # cost since we need to *read* the images below to check validity.
        assert os.path.isfile(image_full_path)
        
        img_id = image_relative_path.replace('\\','/').replace('/','_').replace(' ','_')
        row_indices = filenames_to_rows[image_relative_path]
        
        # i_row = row_indices[0]
        for i_row in row_indices:
            
            row = input_metadata.iloc[i_row]
            assert(row['image_relative_path'] == image_relative_path)
            timestamp = row['datetime']            
            location = row['survey_name'] + '_' + row['camera_name']
            
            if img_id in image_ids_to_images:
                im = image_ids_to_images[img_id]
                assert im['file_name'] == image_relative_path
                assert im['location'] == location
            else:
                im = {}
                
                try:
                    pil_image = Image.open(image_full_path)        
                    width, height = pil_image.size
                    im['width'] = width
                    im['height'] = height
                except:
                    # These generally represent zero-byte images in this data set, don't try
                    # to find the very small handful that might be other kinds of failures we 
                    # might want to keep around.
                    # print('Error opening image {}'.format(image_relative_path))
                    invalid_images.append(image_relative_path)
                    continue
                
                im['id'] = img_id
                im['file_name'] = image_relative_path
                im['datetime'] = timestamp
                im['location'] = location

                image_ids_to_images[img_id] = im

            species = row['species'].lower().strip().replace(' ','_')
            
            if (isinstance(species,float) or \
                (isinstance(species,str) and (len(species) == 0))):
                category_name = 'empty'
            else:
                category_name = species
            del species
            
            category_name = category_name.strip().lower()
            
            common_name = row['common_name']
            if isinstance(common_name,float) and np.isnan(common_name):
                common_name = ''
            else:
                common_name = str(common_name).lower().strip().replace(', ',',').replace(' ','_')
            
            for k,v in category_mappings.items():
                common_name = common_name.replace(k,v)
                category_name = category_name.replace(k,v)
            common_name = common_name.replace('.','').replace('spp','species')
            category_name = category_name.replace('.','').replace('spp','species')
                        
            if category_name == 'passerine_species' and common_name != '' and common_name != 'passerine_species':
                category_name = common_name
                
            # If we've seen this category before...
            if category_name in category_name_to_category:
                
                category = category_name_to_category[category_name]
                category_id = category['id'] 
                
                # ...make sure it used the same latin --> common mapping
                #
                # If the previous instance had no mapping, use the new one.
                if category['common_name'] == '':
                    category['common_name'] = common_name
                else:
                    # assert common_name == category['common_name']
                    if common_name != category['common_name']:
                        print('Warning: common name {} used for species {}, previously {}'.format(
                            common_name,category_name,category['common_name']))
                    
            else:
                
                category_id = next_category_id
                category = {}
                category['id'] = category_id
                category['name'] = category_name
                category['common_name'] = common_name
                category_name_to_category[category_name] = category
                next_category_id += 1
            
            # Create an annotation
            ann = {}        
            ann['id'] = str(uuid.uuid1())
            ann['image_id'] = im['id']    
            ann['category_id'] = category_id
                        
            for target_field in target_fields:
                if target_field in input_metadata.columns:
                    val = row[target_field]
                    if isinstance(val,float) and np.isnan(val):
                        val = ''
                    else:
                        val = str(val).strip()
                    ann[target_field] = val
                
            annotations.append(ann)

        # ...for each annotation we found for this image
            
    # ...for each image
             
# ...for each dataset
            
images = list(image_ids_to_images.values())
categories = list(category_name_to_category.values())

# Print all of our species mappings
for c in categories:
    print(c['name'].ljust(30) + c['common_name'])

print('Finished creating CCT dictionaries, loaded {} images of {} total on disk ({} invalid)'.format(
    len(images), len(image_relative_paths_set), len(invalid_images)))


#%% Copy images for which we actually have annotations to a new folder, lowercase everything

# im = images[0]
for im in tqdm(images):
    relative_filename = im['file_name']
    input_filename = os.path.join(input_base,relative_filename)
    output_filename = os.path.join(output_base,relative_filename).lower()
    os.makedirs(os.path.dirname(output_filename),exist_ok=True)
    
    shutil.copy(input_filename, output_filename)
    im['file_name'] = im['file_name'].lower()
    im['id'] = im['id'].lower()
    

#%% Create info struct

info = {}
info['year'] = 2020
info['version'] = 1
info['description'] = 'UBC Camera Traps'
info['contributor'] = 'UBC'


#%% Convert image IDs to lowercase in annotations, tag as sequence level

# While there isn't any sequence information, the nature of false positives
# here leads me to believe the images were labeled at the sequence level, so
# we should trust labels more when positives are verified.  Overall false
# positive rate looks to be between 1% and 5%.
    
for ann in annotations:
    ann['image_id'] = ann['image_id'].lower()
    ann['sequence_level_annotation'] = True


#%% 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

options = integrity_check_json_db.IntegrityCheckOptions()
options.baseDir = output_base
options.bCheckImageSizes = False
options.bCheckImageExistence = False
options.bFindUnusedImages = True

sortedCategories, data, errors = integrity_check_json_db.integrity_check_json_db(
    output_json_file, options)


#%% Preview labels

viz_options = visualize_db.DbVizOptions()
viz_options.num_to_visualize = 2000
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=output_base,
                                                         options=viz_options)
os.startfile(html_output_file)
