import tensorflow as tf
import sys
import os
from object_detection.utils import dataset_util
from object_detection.utils import label_map_util
import numpy as np
import sys
import cv2
import progress

MIN_SIZE = 30

feature = {
    'image/format' :    tf.FixedLenFeature((), tf.string, default_value='jpeg'),
    'image/filename':   tf.FixedLenFeature((), tf.string, default_value=''),
    'image/key/sha256': tf.FixedLenFeature((), tf.string, default_value=''),
    'image/source_id':  tf.FixedLenFeature((), tf.string, default_value=''),
    'image/height':     tf.FixedLenFeature((), tf.int64, default_value=1),
    'image/width':      tf.FixedLenFeature((), tf.int64, default_value=1),
    'image/encoded':    tf.FixedLenFeature((), tf.string, default_value=''),
    'image/object/bbox/xmin':   tf.VarLenFeature(tf.float32),
    'image/object/bbox/xmax':   tf.VarLenFeature(tf.float32),
    'image/object/bbox/ymin':   tf.VarLenFeature(tf.float32),
    'image/object/bbox/ymax':   tf.VarLenFeature(tf.float32),
    'image/object/class/label': tf.VarLenFeature(tf.int64),
    'image/object/class/text':  tf.VarLenFeature(tf.string),
    'image/object/area':        tf.VarLenFeature(tf.float32),
    'image/object/is_crowd':    tf.VarLenFeature(tf.int64),
    'image/object/difficult':   tf.VarLenFeature(tf.int64),
    'image/object/group_of':    tf.VarLenFeature(tf.int64),
    'image/object/weight':      tf.VarLenFeature(tf.float32),
}

AUTOSAVE_TRAFFIC_SIGNS_NAME_MAP = {
    '1_1'                : '',
    '1_2'                : '',
    '1_5'                : '',
    '1_7'                : 'warning_roundabout_ahead',
    '1_8'                : '',
    '1_11_1'             : '',
    '1_11_2'             : '',
    '1_12_1'             : '',
    '1_12_2'             : '',
    '1_13'               : '',
    '1_14'               : '',
    '1_15'               : '',
    '1_16'               : '',
    '1_17'               : 'warning_uneven_road_ahead',
    '1_18'               : '',
    '1_19'               : '',
    '1_20_1'             : 'warning_road_narrows_on_both',
    '1_20_2'             : 'warning_road_narrows_on_right',
    '1_20_3'             : 'warning_road_narrows_on_left',
    '1_21'               : '',
    '1_22'               : 'warning_pedestrian_crossing_ahead',
    '1_23'               : 'warning_children',
    '1_25'               : '',
    '1_27'               : '',
    '1_33'               : '',
    '2_1'                : 'priority_priority_road',
    '2_2'                : 'priority_eof_priority_road',
    '2_3_1'              : 'warning_crossroads_minor_road_r_l',
    '2_3_2'              : 'warning_crossroads_minor_road_r',
    '2_3_3'              : 'warning_crossroads_minor_road_l',
    '2_4'                : 'priority_give_way',
    '2_5'                : 'priority_stop',
    '2_6'                : '',
    '3_1'                : 'prohibitory_no_entry',
    '3_2'                : 'prohibitory_no_vehicles',
    '3_4'                : 'prohibitory_no_heavy_goods_vehicles',
    '3_10'               : '',
    '3_11_n5'            : 'prohibitory_max_weight',
    '3_13_r3.5'          : 'prohibitory_max_height',
    '3_13_r4.5'          : 'prohibitory_max_height',
    '3_18_1'             : 'prohibitory_no_right_turn',
    '3_18_2'             : 'prohibitory_no_left_turn',
    '3_19'               : 'prohibitory_no_uturn',
    '3_20'               : 'prohibitory_no_overtaking',
    '3_24_n5'            : 'prohibitory_max_speed_5',
    '3_24_n10'           : 'prohibitory_max_speed_10',
    '3_24_n20'           : 'prohibitory_max_speed_20',
    '3_24_n30'           : 'prohibitory_max_speed_30',
    '3_24_n40'           : 'prohibitory_max_speed_40',
    '3_24_n50'           : 'prohibitory_max_speed_50',
    '3_24_n60'           : 'prohibitory_max_speed_65',
    '3_24_n70'           : 'prohibitory_max_speed_70',
    '3_24_n80'           : 'prohibitory_max_speed_80',
    '3_25_n20'           : 'prohibitory_eof_max_speed_20',
    '3_27'               : 'prohibitory_no_parking_or_stopping',
    '3_28'               : 'prohibitory_no_parking',
    '3_29'               : '',
    '3_30'               : '',
    '3_32'               : '',
    '4_1'                : '',
    '4_1_1'              : 'mandatory_proceed_straight',
    '4_1_2'              : 'mandatory_turn_right_ahead',
    '4_1_2_p'            : '',
    '4_1_3'              : 'mandatory_turn_left_ahead',
    '4_1_3_p'            : '',
    '4_1_4'              : 'mandatory_proceed_straight_or_turn_right',
    '4_1_5'              : 'mandatory_proceed_straight_or_turn_left',
    '4_1_6'              : '',
    '4_2_1'              : '',
    '4_2_2'              : '',
    '4_2_3'              : '',
    '4_3'                : 'mandatory_roundabout',
    '5_3'                : '',
    '5_5'                : 'prescription_one_way_road',
    '5_6'                : 'prescription_eof_one_way_road',
    '5_7_1'              : 'prescription_entry_to_one_way_road_on_the_right',
    '5_7_2'              : 'prescription_entry_to_one_way_road_on_the_left',
    '5_8'                : '',
    '5_11'               : 'prescription_bus_lane',
    '5_14'               : '',
    '5_15_1'             : '',
    '5_15_2'             : '',
    '5_15_2_B'           : 'prescription_lane_direction_b',
    '5_15_2_F'           : 'prescription_lane_direction_f',
    '5_15_2_FL'          : 'prescription_lane_direction_fl',
    '5_15_2_FR'          : 'prescription_lane_direction_fr',
    '5_15_2_FR_FL'       : 'prescription_lane_direction_fr_fl',
    '5_15_2_F_FL'        : 'prescription_lane_direction_f_fl',
    '5_15_2_F_FR'        : 'prescription_lane_direction_f_fr',
    '5_15_2_F_L'         : 'prescription_lane_direction_f_l',
    '5_15_2_F_R'         : 'prescription_lane_direction_f_r',
    '5_15_2_L'           : 'prescription_lane_direction_l',
    '5_15_2_R'           : 'prescription_lane_direction_r',
    '5_15_2_R_L'         : 'prescription_lane_direction_r_l',
    '5_15_3'             : 'prescription_start_new_line_right',
    '5_15_4'             : 'prescription_start_new_line_left',
    '5_15_5'             : 'prescription_end_line_right',
    '5_15_6'             : 'prescription_end_line_left',
    '5_15_7'             : '',
    '5_16'               : '',
    '5_18'               : '',
    '5_19_1'             : '',
    '5_20'               : '',
    '6_2_n20'            : '',
    '6_2_n50'            : '',
    '6_2_n60'            : '',
    '6_2_n70'            : '',
    '6_3_1'              : '',
    '6_4'                : 'information_parking',
    '6_6'                : '',
    '6_7'                : '',
    '6_16'               : '',
    '7_1'                : '',
    '7_2'                : '',
    '7_3'                : '',
    '7_4'                : '',
    '7_5'                : '',
    '7_6'                : '',
    '7_7'                : '',
    '7_11'               : '',
    '7_12'               : '',
    '8_15'               : '',
    '8_18'               : '',
    '8_1_1'              : '',
    '8_1_3'              : 'information_distance_object_right',
    '8_1_4'              : 'information_distance_object_left',
    '8_2_1'              : '',
    '8_3_1'              : '',
    '8_3_2'              : '',
    '8_4_1'              : 'information_heavy_vehicle',
    '8_4_4'              : '',
    '8_5_2'              : 'information_working_day',
    '8_5_4'              : '',
    '8_6_2'              : '',
    '8_6_4'              : '',
    '8_8'                : 'information_paid_services'
}

VALID_CLASSES = [
    'warning_roundabout_ahead',
    'warning_uneven_road_ahead',
    'warning_road_narrows_on_both',
    'warning_road_narrows_on_right',
    'warning_road_narrows_on_left',
    'warning_pedestrian_crossing_ahead',
    'warning_children',
    'priority_priority_road',
    'priority_eof_priority_road',
    'warning_crossroads_minor_road_r_l',
    'warning_crossroads_minor_road_r',
    'warning_crossroads_minor_road_l',
    'warning_crossroads_minor_road_fr',
    'warning_crossroads_minor_road_fl',
    'warning_crossroads_minor_road_br',
    'warning_crossroads_minor_road_bl',
    'priority_give_way',
    'priority_stop',
    'prohibitory_no_entry',
    'prohibitory_no_vehicles',
    'prohibitory_no_heavy_goods_vehicles',
    'prohibitory_max_weight',
    'prohibitory_max_weight_per_axle',
    'prohibitory_max_height',
    'prohibitory_max_width',
    'prohibitory_no_right_turn',
    'prohibitory_no_left_turn',
    'prohibitory_no_uturn',
    'prohibitory_no_overtaking',
    'prohibitory_eof_no_overtaking',
    'prohibitory_no_overtaking_by_heavy_vehicle',
    'prohibitory_eof_no_overtaking_by_heavy_vehicle',
    'prohibitory_max_speed_5',
    'prohibitory_max_speed_10',
    'prohibitory_max_speed_15',
    'prohibitory_max_speed_20',
    'prohibitory_max_speed_25',
    'prohibitory_max_speed_30',
    'prohibitory_max_speed_35',
    'prohibitory_max_speed_40',
    'prohibitory_max_speed_45',
    'prohibitory_max_speed_50',
    'prohibitory_max_speed_60',
    'prohibitory_max_speed_65',
    'prohibitory_max_speed_70',
    'prohibitory_max_speed_80',
    'prohibitory_max_speed_90',
    'prohibitory_max_speed_100',
    'prohibitory_max_speed_110',
    'prohibitory_max_speed_120',
    'prohibitory_max_speed_130',
    'prohibitory_eof_max_speed_5',
    'prohibitory_eof_max_speed_10',
    'prohibitory_eof_max_speed_15',
    'prohibitory_eof_max_speed_20',
    'prohibitory_eof_max_speed_25',
    'prohibitory_eof_max_speed_30',
    'prohibitory_eof_max_speed_35',
    'prohibitory_eof_max_speed_40',
    'prohibitory_eof_max_speed_45',
    'prohibitory_eof_max_speed_50',
    'prohibitory_eof_max_speed_60',
    'prohibitory_eof_max_speed_65',
    'prohibitory_eof_max_speed_70',
    'prohibitory_eof_max_speed_80',
    'prohibitory_eof_max_speed_90',
    'prohibitory_eof_max_speed_100',
    'prohibitory_eof_max_speed_110',
    'prohibitory_eof_max_speed_120',
    'prohibitory_eof_max_speed_130',
    'prohibitory_no_parking_or_stopping',
    'prohibitory_no_parking',
    'mandatory_proceed_straight',
    'mandatory_turn_right_ahead',
    'mandatory_turn_left_ahead',
    'mandatory_proceed_straight_or_turn_right',
    'mandatory_proceed_straight_or_turn_left',
    'mandatory_roundabout',
    'prescription_one_way_road',
    'prescription_eof_one_way_road',
    'prescription_entry_to_one_way_road_on_the_right',
    'prescription_entry_to_one_way_road_on_the_left',
    'prescription_bus_lane',
    'prescription_eof_bus_lane',
    'prescription_lane_direction_f',
    'prescription_lane_direction_fr',
    'prescription_lane_direction_r',
    'prescription_lane_direction_b',
    'prescription_lane_direction_l',
    'prescription_lane_direction_fl',
    'prescription_lane_direction_f_fr',
    'prescription_lane_direction_f_r',
    'prescription_lane_direction_fr_r',
    'prescription_lane_direction_f_fl',
    'prescription_lane_direction_f_l',
    'prescription_lane_direction_l_fl',
    'prescription_lane_direction_fr_fl',
    'prescription_lane_direction_r_l',
    'prescription_start_new_line_right',
    'prescription_start_new_line_left',
    'prescription_end_line_right',
    'prescription_end_line_left',
    'prescription_built_up_area',
    'prescription_eof_built_up_area',
    'information_parking',
    'information_distance_stop_ahead',
    'information_distance_object_right',
    'information_distance_object_left',
    'information_start_zone',
    'information_in_zone',
    'information_distance_zone_right',
    'information_distance_zone_left',
    'information_heavy_vehicle',
    'information_light_vehicle',
    'information_holiday',
    'information_working_day',
    'information_holiday_time',
    'information_working_day_time',
    'information_paid_services',
]

def file_lines_cnt(fpath):
    with open(fpath) as f:
        for i, l in enumerate(f):
            pass
    return i + 1

def create_tf_example(filepath, objects):
    img = cv2.imread(filepath)

    height   = img.shape[0] # Image height
    width    = img.shape[1] # Image width
    filename = os.path.basename(filepath) # Filename of the image. Empty if image is not from file
    with open(filepath, 'rb') as file:
        encoded_image_data = file.read() # Encoded image bytes
    image_format = b'jpeg' # b'jpeg' or b'png'

    xmins = [] # List of normalized left x coordinates in bounding box (1 per box)
    xmaxs = [] # List of normalized right x coordinates in bounding box (1 per box)
    ymins = [] # List of normalized top y coordinates in bounding box (1 per box)
    ymaxs = [] # List of normalized bottom y coordinates in bounding box (1 per box)
    classes_text = [] # List of string class name of bounding box (1 per box)
    classes = [] # List of integer class id of bounding box (1 per box)

    for obj in objects:
        classes = classes + [obj[0]]
        classes_text = classes_text + [VALID_CLASSES[obj[0] - 1].encode('utf8')]
        xmins = xmins + [float(obj[1][0]) / float(width)]
        xmaxs = xmaxs + [float(obj[1][2]) / float(width)]
        ymins = ymins + [float(obj[1][1]) / float(height)]
        ymaxs = ymaxs + [float(obj[1][3]) / float(height)]


    tf_example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': dataset_util.int64_feature(height),
        'image/width': dataset_util.int64_feature(width),
        'image/filename': dataset_util.bytes_feature(filename.encode('utf8')),
        'image/source_id': dataset_util.bytes_feature(filename.encode('utf8')),
        'image/encoded': dataset_util.bytes_feature(encoded_image_data),
        'image/format': dataset_util.bytes_feature(image_format),
        'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
        'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
        'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
        'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
        'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
        'image/object/class/label': dataset_util.int64_list_feature(classes),
    }))
    return tf_example

def load_label_map(label_map_path):
    label_map = label_map_util.get_label_map_dict(label_map_path)
    classes_names = [None] * len(label_map)
    for label in label_map.items():
        classes_names[label[1] - 1] = label[0]
    return classes_names

def main(autosave_markup_path, label_map_path, output_tfrecord_path):
    classes_names = load_label_map(label_map_path)
    print(classes_names)

    writer = tf.python_io.TFRecordWriter(output_tfrecord_path)
    lines_cnt = file_lines_cnt(autosave_markup_path)
    line_idx = 0
    with open(autosave_markup_path) as markup:
        for line in markup:
            line = line.split()
            img_filepath = line[0]
            obj_cnt = int(line[1])
            img_objects = []
            progress.printProgressBar(line_idx + 1, lines_cnt, prefix = '', suffix = img_filepath[img_filepath.rfind('/') + 1:])
            line_idx += 1
            for obj_idx in range(obj_cnt):
                if (line[2 + 5 * obj_idx] not in AUTOSAVE_TRAFFIC_SIGNS_NAME_MAP):
                    continue
                temp = AUTOSAVE_TRAFFIC_SIGNS_NAME_MAP[line[2 + 5 * obj_idx]]
                if (temp not in classes_names):
                    continue
                if (MIN_SIZE > int(line[5 + 5 * obj_idx]) or
                    MIN_SIZE > int(line[6 + 5 * obj_idx])):
                    continue
                cl_idx = classes_names.index(temp)
                obj_rect =[int(line[3 + 5 * obj_idx]),  int(line[4 + 5 * obj_idx]),
                           int(line[3 + 5 * obj_idx]) + int(line[5 + 5 * obj_idx]) - 1,
                           int(line[4 + 5 * obj_idx]) + int(line[6 + 5 * obj_idx]) - 1]

                img_objects = img_objects + [(cl_idx + 1, obj_rect)]
            if (len(img_objects) == 0):
                continue
            tf_example = create_tf_example(img_filepath, img_objects)
            writer.write(tf_example.SerializeToString())
    writer.close()

if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2], sys.argv[3])
