# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

import os

import yaml
from detectron2.config import CfgNode as CN


def add_dataset_config(cfg):
    cfg.DATASETS.AUG = False
    cfg.DATASETS.RANDOM_CROP_PROB = 0.0
    cfg.DATASETS.IGNORE_DIFFICULT = False
    cfg.DATASETS.FIX_CROP = False
    cfg.DATASETS.CROP_SIZE = (512, 512)
    cfg.DATASETS.MAX_ROTATE_THETA = 30
    cfg.DATASETS.FIX_ROTATE = False


def add_glass_config(cfg):
    """
    Add config for Hybrid head.
    """
    ## Orientation related fields
    cfg.MODEL.ROTATED_BOXES_ON = False
    cfg.MODEL.ORIENTATION_ON = False

    ## Localization old Hybrid head config
    cfg.MODEL.ROI_HYBRID_HEAD = CN()
    cfg.MODEL.ROI_HYBRID_HEAD.NAME = "ResBlockHybridHead"
    cfg.MODEL.ROI_HYBRID_HEAD.POOLER_RESOLUTION = 64
    cfg.MODEL.ROI_HYBRID_HEAD.NUM_FEATURES = 256
    cfg.MODEL.ROI_HYBRID_HEAD.DEPTH = 3
    cfg.MODEL.ROI_HYBRID_HEAD.NORM_IMG_CROPS = False
    ## Filter dont care boxes in the RPN not recommended to use untested yet
    cfg.MODEL.FILTERED_RPN = CN()
    cfg.MODEL.FILTERED_RPN.IGNORE_TEXT = ["###", ""]

    ## E2E Feature Fusion
    cfg.MODEL.LOCAL_FEATURE_EXTRACTOR = CN()
    cfg.MODEL.LOCAL_FEATURE_EXTRACTOR.NAME = "ResNet_FeatureExtractor"
    cfg.MODEL.LOCAL_FEATURE_EXTRACTOR.NUM_FEATURES = 256

    cfg.MODEL.HYBRID_FUSION = CN()
    cfg.MODEL.HYBRID_FUSION.NAME = "MultiAspectGCAttention"
    cfg.MODEL.HYBRID_FUSION.NUM_FEATURES = 256

    ## MultiAspectGCAttention
    cfg.MODEL.HYBRID_FUSION.RATIO = 0.5
    cfg.MODEL.HYBRID_FUSION.HEADERS = 8
    cfg.MODEL.HYBRID_FUSION.FUSION_TYPE = 'channel_add'
    cfg.MODEL.ROI_MASK_HEAD.LOSS_WEIGHT = 0.005
    cfg.MODEL.ROI_HEADS.CLASS_NAMES = ['word']
    cfg.MODEL.ROI_ORIENTATION_HEAD = CN()
    cfg.MODEL.ROI_ORIENTATION_HEAD.LOSS_WEIGHT = 1.0
    cfg.MODEL.ROI_ORIENTATION_HEAD.APPLY_TO_BOXES = False
    cfg.MODEL.ROI_ORIENTATION_HEAD.APPLY_TO_BOXES_DURING_TRAINING = True

    # cfg.MODEL.ROI_MASK_HEAD.MASK_INFERENCE = True

    # cfg.TEST.MIN_SIZE_TEST = 1600
    # cfg.TEST.MIN_SIZE_TEST = 1600
    cfg.INPUT.MIN_SIZE_TEST = 1600
    cfg.INPUT.MAX_SIZE_TEST = 1600

    # Maximum ratio to upscale small images, until they reach INPUT.MIN_TEST
    cfg.INPUT.MAX_UPSCALE_RATIO = 2

    cfg.INPUT.ROTATION = CN()
    cfg.INPUT.ROTATION.ENABLED = False
    cfg.INPUT.ROTATION.ANGLES = [0]

    cfg.TEST.IOU_THRESHOLD = 0.5
    cfg.TEST.USE_FILTERED_METRICS = True
    cfg.TEST.DONT_CARE_GT_LABELS = ['###', '']


def add_e2e_config(cfg):
    """
    Add config for cotrain.
    """
    _C = cfg
    _C.MODEL.RECOGNIZER_ON = True
    _C.MODEL.ROI_MASK_HEAD.SAMPLE_WORDS_STRATEGY = 'random'  # 'random' | 'long_first'
    _C.MODEL.ROI_MASK_HEAD.SAMPLE_WORDS_STRATEGY_PROB = 0.3
    _C.MODEL.ROI_MASK_HEAD.CLASS_IND = 0
    _C.MODEL.ROI_MASK_HEAD.IGNORE_EMPTY_TEXT = True

    # Recognizer Head
    # ==========================
    _C.MODEL.RECOGNIZER_ON = False
    _C.MODEL.ROI_MASK_HEAD.LABELS_TYPE = "attention"
    _C.MODEL.ROI_MASK_HEAD.MAX_WORD_LENGTH = 50
    _C.MODEL.ROI_MASK_HEAD.CHARACTER_SET = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&''()*+,-./:;<=>?@[\]^_`{|}~ '
    _C.MODEL.ROI_MASK_HEAD.UNK_SYMBOL_PRED = False
    _C.MODEL.ROI_MASK_HEAD.IGNORE_EMPTY_TEXT = True
    _C.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION_WIDTH = 32
    _C.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION_HEIGHT = 32
    _C.MODEL.ROI_MASK_HEAD.IN_FEATURES = ['p2', 'p3', 'p4', 'p5', 'p6']
    _C.MODEL.ROI_MASK_HEAD.CLASS_IND = 0
    _C.MODEL.ROI_MASK_HEAD.PAD_SAMPLER = ''  # 'FeatPadV1' 'FeatPadV2'
    _C.MODEL.ROI_MASK_HEAD.MAX_BATCH_SIZE = 256
    _C.MODEL.ROI_MASK_HEAD.LOSS_WEIGHT = 2.0
    _C.MODEL.ROI_MASK_HEAD.IGNORE_TEXT = ["###"]
    _C.MODEL.ROI_MASK_HEAD.SENSITIVE = True  # sensitive character mode

    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD = CN()

    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.POOLER_PAD = CN()
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.POOLER_PAD.NAME = ""

    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.BACKBONE = CN()
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.BACKBONE.NAME = "CNN_V1"

    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.ENCODER = CN()
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.ENCODER.NAME = "BiLSTMBlock"  # | TransformerEncoderBlock
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.ENCODER.NUM_OF_LAYERS = 2
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.ENCODER.HEIGHT_REDUCTION = "mean"  # "flatten" | "mean"
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.ENCODER.N_HEAD = 8  # for TransformerEncoderBlock (multi-attention heads)

    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.DECODER = CN()
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.DECODER.NAME = "ASTER"
    _C.MODEL.ROI_MASK_HEAD.RECOGNIZER_HEAD.DECODER.POS_ENC_HEIGHT_WIDTH = None

    ################# config duplication #################
    _C.MODEL.ROI_RECOGNIZER_HEAD = CN()
    _C.MODEL.ROI_RECOGNIZER_HEAD.SAMPLE_WORDS_STRATEGY = 'random'  # 'random' | 'long_first'
    _C.MODEL.ROI_RECOGNIZER_HEAD.SAMPLE_WORDS_STRATEGY_PROB = 0.3

    _C.MODEL.ROI_RECOGNIZER_HEAD.NAME = ""
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD = CN()

    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.POOLER_PAD = CN()
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.POOLER_PAD.NAME = ""

    _C.MODEL.ROI_RECOGNIZER_HEAD.LABELS_TYPE = "attention"
    _C.MODEL.ROI_RECOGNIZER_HEAD.MAX_WORD_LENGTH = 50
    _C.MODEL.ROI_RECOGNIZER_HEAD.CHARACTER_SET = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&''()*+,-./:;<=>?@[\]^_`{|}~ '
    _C.MODEL.ROI_RECOGNIZER_HEAD.UNK_SYMBOL_PRED = False
    _C.MODEL.ROI_RECOGNIZER_HEAD.IGNORE_EMPTY_TEXT = True
    _C.MODEL.ROI_RECOGNIZER_HEAD.POOLER_RESOLUTION_WIDTH = 32
    _C.MODEL.ROI_RECOGNIZER_HEAD.POOLER_RESOLUTION_HEIGHT = 32
    _C.MODEL.ROI_RECOGNIZER_HEAD.IN_FEATURES = ['p2', 'p3', 'p4', 'p5', 'p6']
    _C.MODEL.ROI_RECOGNIZER_HEAD.CLASS_IND = 0
    _C.MODEL.ROI_RECOGNIZER_HEAD.PAD_SAMPLER = ''  # 'FeatPadV1' 'FeatPadV2'
    _C.MODEL.ROI_RECOGNIZER_HEAD.MAX_BATCH_SIZE = 256
    _C.MODEL.ROI_RECOGNIZER_HEAD.LOSS_WEIGHT = 2.0
    _C.MODEL.ROI_RECOGNIZER_HEAD.IGNORE_TEXT = ["###"]
    _C.MODEL.ROI_RECOGNIZER_HEAD.SENSITIVE = True  # sensitive character mode

    _C.MODEL.ROI_RECOGNIZER_HEAD.POOLER_TYPE = "ROIAlignRotated"
    _C.MODEL.ROI_RECOGNIZER_HEAD.NORM = "BN"
    _C.MODEL.ROI_RECOGNIZER_HEAD.CHARACTER_SET = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&''()*+,-./:;<=>?@[\]^_`{|}~ '
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.BACKBONE = CN()
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.BACKBONE.NAME = "CNN_V1_2"
    _C.MODEL.ROI_RECOGNIZER_HEAD.POOLER_SAMPLING_RATIO = 0
    _C.MODEL.ROI_RECOGNIZER_HEAD.CONV_DIM = 256

    _C.MODEL.ROI_RECOGNIZER_HEAD.SAMPLING_RATIO = 0
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.ENCODER = CN()
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.ENCODER.NAME = "BiLSTMBlockV2"  # | TransformerEncoderBlock
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.ENCODER.NUM_OF_LAYERS = 2
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.ENCODER.HEIGHT_REDUCTION = "mean"  # "flatten" | "mean"
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.ENCODER.N_HEAD = 8  # for TransformerEncoderBlock (multi-attention heads)

    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.DECODER = CN()
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.DECODER.NAME = "ASTER_V2"
    _C.MODEL.ROI_RECOGNIZER_HEAD.RECOGNIZER_HEAD.DECODER.POS_ENC_HEIGHT_WIDTH = None

    _C.MODEL.ROI_MASK_HEAD.MASK_INFERENCE = False


def add_post_process_config(cfg):
    # Post-Processing Constants:
    # ===========================
    cfg.POST_PROCESSING = CN()
    cfg.POST_PROCESSING.NAME = 'PostProcessorAcademic'
    # Master flag for skipping everything
    cfg.POST_PROCESSING.SKIP_ALL = False
    # How much to inflate the boxes in percentage
    cfg.POST_PROCESSING.BOX_INFLATE_RATIO = 0.05

    # How many pixels to add (or subtract) to each box [left, top, right, bottom]
    cfg.POST_PROCESSING.BOX_PX_PADDING = [0, 0, 0, 0]

    # We drop boxes that have less than this number of pixels in one of their dimensions
    cfg.POST_PROCESSING.MIN_BOX_DIMENSION = 2

    # We drop boxes that span by more than this amount of pixels outside the image regions
    cfg.POST_PROCESSING.MAX_OUTSIDE_BOX_MARGIN_PX = 5

    # Above this IoU threshold we merge two boxes (if the additional merge conditions are met)
    # Lower number -> more merges
    cfg.POST_PROCESSING.MERGE_IOA_THRESH = 0.3

    # Threshold for overlap width / smaller box height to determine if we should merge
    # Lower number -> more merges
    cfg.POST_PROCESSING.OVERLAP_WIDTH_PER_HEIGHT_THRESH = 0.3

    # What's the max tolerance for height ratio differences in order to merge
    # Lower number -> more merges
    cfg.POST_PROCESSING.PAIRS_HEIGHT_RATIO_THRESH = 0.35

    # Filtered out boxes get this confidence level
    cfg.POST_PROCESSING.LOW_CONFIDENCE = 0.01
    # What's the minimal confidence above it a box is considered valid for merging / contain other boxes
    cfg.POST_PROCESSING.VALID_CONFIDENCE = 0.15

    # Thresholds for filtering
    cfg.POST_PROCESSING.DETECT_THRESHOLD = 0.25
    cfg.POST_PROCESSING.TEXT_THRESHOLD = 0.25

    # Max angle difference between boxes, used for rotated boxes
    cfg.POST_PROCESSING.MAX_ANGLE_DIFF = 15


def merge_from_dataset_config(cfg, dataset_config_path):
    _C = cfg
    # Parsing the dataset_config_path
    with open(dataset_config_path, 'r') as fp:
        dataset_config = yaml.safe_load(fp)
    # Updating the fields in the config
    _C.DATASETS.TRAIN = dataset_config.get('DATASETS', [])
    _C.DATASETS.TEST = dataset_config.get('VAL_DATASETS', [])
    _C.DATASETS.ROOT = dataset_config.get('ROOT', '.')
    _C.DATASETS.CONFIG = os.path.basename(dataset_config_path)
