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


class InterHandHyperparams:
    """Interhand model hyperparameters
    """
    # input image size
    IMG_SHAPE = (256, 256)
    # heatmap generation dimensions
    DEPTH_DIM = 400
    JOINT_HM_DIMS = (64, 64, 64)
    ROOT_HM_DEPTH = 64
    SIGMA = 2.5

    BATCH_SIZE = 8


class HandsHyperparams:
    """Hand pose estimation model hyperparameters
    """
    # input image size
    IMG_SHAPE = (256, 256)
    # minimal bbox are in input image
    MIN_BBOX_AREA = 30 * 30
    # heatmap generation dimensions
    HM_DIMS = (64, 64)
    SIGMA = 2.5
    # classification cross entropy weights
    LOSS_WEIGHTS = (0.55886041, 4.74733678)  # negative, positive (single hands, interacting hands)
    # training parameters
    INIT_EPOCHS = 1
    FINETUNE_EPOCHS = 100
    BATCH_SIZE = 4
    INIT_LEARNING_RATE = 5e-3
    FINETUNE_LEARNING_RATE = 1e-3
    # inference parameters
    DELTA = 4
    HANDNESS_THOLD = 0.9


class WFLWHyperparams:
    """Face pose estimation model hyperparameters
    """
    # input image size
    IMG_SHAPE = (256, 256)
    # heatmap generation dimensions
    HM_DIMS = (64, 64)
    SIGMA = 2.5
    # training parameters
    INIT_EPOCHS = 5
    FINETUNE_EPOCHS = 100
    BATCH_SIZE = 4
    INIT_LEARNING_RATE = 1e-3
    FINETUNE_LEARNING_RATE = 1e-3
    # inference parameters
    DELTA = 4


class Config:
    """Main configuration object, contains project's directory tree structure and trained models' hyperparameter objects
    """
    BASE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
    # weights and tfrecord directory
    WEIGHTS_PATH = os.path.join(BASE_PATH, 'weights')
    TFRECORDS_PATH = os.path.join(BASE_PATH, 'data', 'tfrecords')
    # checkpoints directory
    CHECKPOINTS_PATH = os.path.join(WEIGHTS_PATH, 'checkpoints')
    # interhand dataset directories
    INTERHAND_DATASET_PATH = os.path.join(BASE_PATH, 'data', 'InterHand')
    INTERHAND_SKELETON_PATH = os.path.join(INTERHAND_DATASET_PATH, 'skeleton.txt')
    INTERHAND_DATASET_PATHS = {
        'images': {
            'train': os.path.join(INTERHAND_DATASET_PATH, 'images', 'train'),
            'val': os.path.join(INTERHAND_DATASET_PATH, 'images', 'val'),
            'test': os.path.join(INTERHAND_DATASET_PATH, 'images', 'test'),
        },
        'annotations': {
            'train': {
                'data': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'human_annot', 'InterHand2.6M_train_data.json'),
                'joint': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'human_annot', 'InterHand2.6M_train_joint_3d.json'),
                'camera': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'human_annot', 'InterHand2.6M_train_camera.json')
            },
            'val': {
                'data': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'machine_annot', 'InterHand2.6M_val_data.json'),
                'joint': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'machine_annot', 'InterHand2.6M_val_joint_3d.json'),
                'camera': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'machine_annot', 'InterHand2.6M_val_camera.json')
            },
            'test': {
                'data': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'human_annot', 'InterHand2.6M_test_data.json'),
                'joint': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'human_annot', 'InterHand2.6M_test_joint_3d.json'),
                'camera': os.path.join(INTERHAND_DATASET_PATH, 'annotations', 'human_annot', 'InterHand2.6M_test_camera.json')
            }
        }
    }
    # COCO WholeBody dataset directories
    COCO_DATASET_PATH = os.path.join(BASE_PATH, 'data', 'coco')
    COCO_SKELETON_PATH = os.path.join(COCO_DATASET_PATH, 'skeleton.txt')
    COCO_DATASET_PATHS = {
        'train': {
            'images': os.path.join(COCO_DATASET_PATH, 'images', 'train2017'),
            'annotations': os.path.join(COCO_DATASET_PATH, 'annotations', 'coco_wholebody_train_v1.0.json'),
        },
        'test': {
            'images': os.path.join(COCO_DATASET_PATH, 'images', 'val2017'),
            'annotations': os.path.join(COCO_DATASET_PATH, 'annotations', 'coco_wholebody_val_v1.0.json'),
        },
    }
    # Halpe dataset directories
    HALPE_DATASET_PATH = os.path.join(BASE_PATH, 'data', 'halpe')
    HALPE_DATASET_PATHS = {
        'images': os.path.join(HALPE_DATASET_PATH, 'images', 'train2015'),
        'annotations': os.path.join(HALPE_DATASET_PATH, 'halpe_train_v1.json'),
    }
    # MWB dataset directories
    MWB_DATASET_PATH = os.path.join(BASE_PATH, 'data', 'mwb')
    MWB_DATASET_PATHS = {
        'manual': {
            'train': os.path.join(MWB_DATASET_PATH, 'manual_train'),
            'test': os.path.join(MWB_DATASET_PATH, 'manual_test'),
        },
    }
    # RHD dataset directories
    RHD_DATASET_PATH = os.path.join(BASE_PATH, 'data', 'rhd')
    RHD_SKELETON_PATH = os.path.join(RHD_DATASET_PATH, 'skeleton.txt')
    RHD_DATASET_PATHS = {
        'train': {
            'images': os.path.join(RHD_DATASET_PATH, 'training', 'color'),
            'annotations': os.path.join(RHD_DATASET_PATH, 'training', 'anno_training.pickle')
        },
        'test': {
            'images': os.path.join(RHD_DATASET_PATH, 'evaluation', 'color'),
            'annotations': os.path.join(RHD_DATASET_PATH, 'evaluation', 'anno_evaluation.pickle')
        }
    }
    # WFLW dataset directories
    WFLW_DATASET_PATH = os.path.join(BASE_PATH, 'data', 'WFLW')
    WFLW_DATASET_PATHS = {
        'images': os.path.join(WFLW_DATASET_PATH, 'WFLW_images'),
        'annotations': {
            'train': os.path.join(WFLW_DATASET_PATH, 'wflw_annotations', 'list_98pt_rect_attr_train_test', 'list_98pt_rect_attr_train.txt'),
            'test': os.path.join(WFLW_DATASET_PATH, 'wflw_annotations', 'list_98pt_rect_attr_train_test', 'list_98pt_rect_attr_test.txt'),
        }
    }
    # interhand tfrecords
    INTERHAND_TFRECORD_PATHS = {
        'train': os.path.join(TFRECORDS_PATH, 'interhand_train.tfrecord'),
        'val': os.path.join(TFRECORDS_PATH, 'interhand_val.tfrecord'),
        'test': os.path.join(TFRECORDS_PATH, 'interhand_test.tfrecord'),
    }
    # hand pose estimation tfrecords
    HANDS_TFRECORD_PATHS = {
        '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'),
        'coco_test': os.path.join(TFRECORDS_PATH, 'hands_test_coco.tfrecord'),
        'mwb_test': os.path.join(TFRECORDS_PATH, 'hands_test_mwb.tfrecord'),
        'rhd_test': os.path.join(TFRECORDS_PATH, 'hands_test_rhd.tfrecord'),
    }
    # face pose estimation tfrecords
    WFLW_TFRECORD_PATHS = {
        'train': os.path.join(TFRECORDS_PATH, 'wflw_train.tfrecord'),
        'val': os.path.join(TFRECORDS_PATH, 'wflw_val.tfrecord'),
        'test': os.path.join(TFRECORDS_PATH, 'wflw_test.tfrecord'),
    }
    # interhand tensorflow weights
    INTERHAND_TF_PATH = os.path.join(WEIGHTS_PATH, 'interhand.tf')
    # hand pose estimation weights
    HANDS_WEIGHTS = {
        'efficientnet-b0': os.path.join(WEIGHTS_PATH, 'effnetb0_hands.tf'),
        'efficientnet-b1': os.path.join(WEIGHTS_PATH, 'effnetb1_hands.tf'),
        'efficientnet-b2': os.path.join(WEIGHTS_PATH, 'effnetb2_hands.tf'),
        'efficientnet-b3': os.path.join(WEIGHTS_PATH, 'effnetb3_hands.tf'),
    }
    # face pose estimation weights
    WFLW_WEIGHTS = {
        'efficientnet-b0': os.path.join(WEIGHTS_PATH, 'effnetb0_wflw.tf'),
        'efficientnet-b1': os.path.join(WEIGHTS_PATH, 'effnetb1_wflw.tf'),
        'efficientnet-b2': os.path.join(WEIGHTS_PATH, 'effnetb2_wflw.tf'),
    }
    # interhand evaluation residual file
    INTERHAND_RESIDUALS_PATH = os.path.join(BASE_PATH, 'data', 'InterHand', 'residuals.json')
    # model hyperparameters
    INTERHAND_PARAMS = InterHandHyperparams()
    HANDS_PARAMS = HandsHyperparams()
    WFLW_PARAMS = WFLWHyperparams()


CFG = Config()
