"""
Configuration file containing information about project directory three structure and model hyperparameters.
"""
import sys
import os

import tensorflow as tf


def get_anchors(path):
    """Load object detector anchors from file

    Args:
        path (str): path to anchor file

    Returns:
        tf.Tensor: tensor of [N, 3, 2] shape, where N is number of model heads (3 for yolov5 and 5 for efficientdet)
    """
    # if file not found
    if not os.path.isfile(path):
        return None
    # read file lines and convert to floats
    with open(path, 'r') as file:
        lines = file.readlines()
    anchors = list()
    for line in lines:
        nums = line.split(' ')
        anchors.append([float(nums[0]), float(nums[1])])
    # convert to tensorflow tensor and reshape to [N, 3, 2]
    return tf.reshape(tf.convert_to_tensor(anchors, dtype=tf.float32), (-1, 3, 2))


class FacesHyperparams:
    """
    Hyperparameters and anchor path for face detection model
    """
    BASE_PATH = os.path.dirname(os.path.realpath(__file__))
    # anchor file paths - yolov5 anchors and efficientdet anchors
    ANCHORS_PATH = os.path.join(BASE_PATH, '..', 'data', 'anchors', 'faces', 'anchors.txt')
    EFFDET_ANCHORS_PATH = os.path.join(BASE_PATH, '..', 'data', 'anchors', 'faces', 'efficientdet_anchors.txt')
    # model input image shape
    IMG_SHAPE = (512, 512)
    # load anchor file
    ANCHORS = get_anchors(ANCHORS_PATH)
    # training hyperparameters
    EPOCHS = 50
    BATCH_SIZE = 8
    LEARNING_RATE = 0.001
    # inference and testing hyperparameters
    MAX_DETECTIONS = 200
    IOU_THRESHOLD = 0.1
    SCORE_THRESHOLD = 0.5


class HandsHyperparams:
    """
    Hyperparameters and anchor path for hand detection model
    """
    BASE_PATH = os.path.dirname(os.path.realpath(__file__))
    # anchor file paths - yolov5 anchors and efficientdet anchors
    ANCHORS_PATH = os.path.join(BASE_PATH, '..', 'data', 'anchors', 'hands', 'anchors.txt')
    EFFDET_ANCHORS_PATH = os.path.join(BASE_PATH, '..', 'anchors', 'data', 'hands', 'efficientdet_anchors.txt')
    # model input image shape
    IMG_SHAPE = (512, 512)
    # load anchor file
    ANCHORS = get_anchors(ANCHORS_PATH)
    # training hyperparameters
    EPOCHS = 50
    BATCH_SIZE = 8
    LEARNING_RATE = 0.001
    # inference and testing hyperparameters
    MAX_DETECTIONS = 100
    IOU_THRESHOLD = 0.1
    SCORE_THRESHOLD = 0.1


class CombinedHyperparams:
    """
    Hyperparameters and anchor path for combined hand and face detection model
    """
    BASE_PATH = os.path.dirname(os.path.realpath(__file__))
    # anchor file paths - yolov5 anchors and efficientdet anchors
    YOLO_ANCHORS_PATH = os.path.join(BASE_PATH, '..', 'data', 'anchors', 'combined', 'anchors.txt')
    EFFDET_ANCHORS_PATH = os.path.join(BASE_PATH, '..', 'data', 'anchors', 'combined', 'efficientdet_anchors.txt')
    # model input image shape
    IMG_SHAPE = (512, 512)
    # load anchor files
    YOLO_ANCHORS = get_anchors(YOLO_ANCHORS_PATH)
    EFFDET_ANCHORS = get_anchors(EFFDET_ANCHORS_PATH)
    # training hyperparameters
    EPOCHS = 50
    BATCH_SIZE = 4
    LEARNING_RATE = 0.001
    # inference and testing hyperparameters
    MAX_DETECTIONS = 10
    IOU_THRESHOLD = 0.2
    SCORE_THRESHOLD = 0.75


class Config:
    """
    Configuration object, contains directory structure and model hyperparameter objects
    """
    BASE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
    # weights root path
    WEIGHTS_PATH = os.path.join(BASE_PATH, 'weights')
    # tfrecords root path
    TFRECORDS_PATH = os.path.join(BASE_PATH, 'data', 'tfrecords')
    # face detection dataset directory tree
    FACES_DATASET_PATH = {
        'images': {
            'train': os.path.join(BASE_PATH, 'data', 'faces', 'WIDER_train', 'images'),
            'val': os.path.join(BASE_PATH, 'data', 'faces', 'WIDER_val', 'images'),
            'test': os.path.join(BASE_PATH, 'data', 'faces', 'WIDER_test', 'images'),
        },
        'annotations': {
            'train': os.path.join(BASE_PATH, 'data', 'faces', 'wider_face_split', 'wider_face_train.mat'),
            'val': os.path.join(BASE_PATH, 'data', 'faces', 'wider_face_split', 'wider_face_val.mat'),
        }
    }
    # face detection tfrecord files
    FACES_TFRECORDS_PATH = {
        'train': os.path.join(TFRECORDS_PATH, 'faces_train.tfrecord'),
        'val': os.path.join(TFRECORDS_PATH,  'faces_val.tfrecord'),
    }
    # face detection weight files
    FACES_WEIGHTS_PATH = {
        'yolov5s': os.path.join(WEIGHTS_PATH, 'faces_yolov5s.tf'),
        'yolov5m': os.path.join(WEIGHTS_PATH, 'faces_yolov5m.tf'),
        'yolov5l': os.path.join(WEIGHTS_PATH, 'faces_yolov5l.tf'),
    }
    # hand detection dataset directory tree
    HANDS_DATASET_PATH = {
        'images': os.path.join(BASE_PATH, 'data', 'hands', 'JPEGImages'),
        'annotations': os.path.join(BASE_PATH, 'data', 'hands', 'Annotations'),
        'split': {
            'train': os.path.join(BASE_PATH, 'data', 'hands', 'ImageSets', 'Main', 'train.txt'),
            'val': os.path.join(BASE_PATH, 'data', 'hands', 'ImageSets', 'Main', 'val.txt'),
            'test': os.path.join(BASE_PATH, 'data', 'hands', 'ImageSets', 'Main', 'test.txt')
        }
    }
    # hand detection tfrecord files
    HANDS_TFRECORDS_PATH = {
        'train': os.path.join(TFRECORDS_PATH, 'hands_train.tfrecord'),
        'val': os.path.join(TFRECORDS_PATH,  'hands_val.tfrecord'),
        'test': os.path.join(TFRECORDS_PATH,  'hands_test.tfrecord'),
    }
    # hand detection weight files
    HANDS_WEIGHTS_PATH = {
        'yolov5s': os.path.join(WEIGHTS_PATH, 'hands_yolov5s.tf'),
        'yolov5m': os.path.join(WEIGHTS_PATH, 'hands_yolov5m.tf'),
        'yolov5l': os.path.join(WEIGHTS_PATH, 'hands_yolov5l.tf'),
    }
    # combined dataset annotation files
    COMBINED_DATASET_PATH = {
        'train': os.path.join(BASE_PATH, 'data', 'combined', 'combined_train.json'),
        'val': os.path.join(BASE_PATH, 'data', 'combined', 'combined_val.json'),
        'test': os.path.join(BASE_PATH, 'data', 'combined', 'combined_test.json'),
    }
    # combined dataset tfrecord files
    COMBINED_TFRECORDS_PATH = {
        'train': os.path.join(TFRECORDS_PATH, 'combined_train.tfrecord'),
        'val': os.path.join(TFRECORDS_PATH,  'combined_val.tfrecord'),
        'test': os.path.join(TFRECORDS_PATH,  'combined_test.tfrecord'),
    }
    # combined dataset weight files
    COMBINED_WEIGHTS_PATH = {
        'yolov5s': os.path.join(WEIGHTS_PATH, 'combined_yolov5s.tf'),
        'yolov5m': os.path.join(WEIGHTS_PATH, 'combined_yolov5m.tf'),
        'yolov5l': os.path.join(WEIGHTS_PATH, 'combined_yolov5l.tf'),
        'efficientdet-d0': os.path.join(WEIGHTS_PATH, 'combined_efficientdet_d0.tf'),
        'efficientdet-d1': os.path.join(WEIGHTS_PATH, 'combined_efficientdet_d1.tf'),
        'efficientdet-d2': os.path.join(WEIGHTS_PATH, 'combined_efficientdet_d2.tf'),
        'efficientdet-d3': os.path.join(WEIGHTS_PATH, 'combined_efficientdet_d3.tf'),
        'efficientdet-d4': os.path.join(WEIGHTS_PATH, 'combined_efficientdet_d4.tf'),
        'efficientdet-d5': os.path.join(WEIGHTS_PATH, 'combined_efficientdet_d5.tf'),
        'efficientdet-d6': os.path.join(WEIGHTS_PATH, 'combined_efficientdet_d6.tf'),
    }
    # transfer learning weights files for yolov5 and efficientdet object detector
    PRETRAINED_TF_WEIGHTS_PATH = {
        'yolov5s': os.path.join(WEIGHTS_PATH, 'pretrained_yolov5s.tf'),
        'yolov5m': os.path.join(WEIGHTS_PATH, 'pretrained_yolov5m.tf'),
        'yolov5l': os.path.join(WEIGHTS_PATH, 'pretrained_yolov5l.tf'),
        'efficientdet-d0': os.path.join(WEIGHTS_PATH, 'pretrained_efficientdet_d0.tf'),
        'efficientdet-d1': os.path.join(WEIGHTS_PATH, 'pretrained_efficientdet_d1.tf'),
        'efficientdet-d2': os.path.join(WEIGHTS_PATH, 'pretrained_efficientdet_d2.tf'),
        'efficientdet-d3': os.path.join(WEIGHTS_PATH, 'pretrained_efficientdet_d3.tf'),
        'efficientdet-d4': os.path.join(WEIGHTS_PATH, 'pretrained_efficientdet_d4.tf'),
        'efficientdet-d5': os.path.join(WEIGHTS_PATH, 'pretrained_efficientdet_d5.tf'),
        'efficientdet-d6': os.path.join(WEIGHTS_PATH, 'pretrained_efficientdet_d6.tf'),
    }
    # pretrained yolov5 model parameters
    YOLO_PRETRAINED_IMG_SHAPE = (640, 640, 3)
    YOLO_PRETRAINED_CLASSES = 80

    HANDS_SPLIT = 0.9
    
    # min relative bbox annotation size to keep
    BBOX_MIN_AREA = 0.04 * 0.04
    # model hyperparameter objects
    FACES_PARAMS = FacesHyperparams()
    HANDS_PARAMS = HandsHyperparams()
    COMBINED_PARAMS = CombinedHyperparams()


CFG = Config()
