# -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import

import numpy as np

from configs._base_.models.retinanet_r50_fpn import *
from configs._base_.datasets.dota_detection import *
from configs._base_.schedules.schedule_1x import *
from alpharotate.utils.pretrain_zoo import PretrainModelZoo

# schedule
BATCH_SIZE = 1
GPU_GROUP = "0,1,2"
NUM_GPU = len(GPU_GROUP.strip().split(','))
LR = 1e-3
SAVE_WEIGHTS_INTE = 11725 * 2
DECAY_EPOCH = [8, 11, 20]
MAX_EPOCH = 12
WARM_EPOCH = 1 / 16.
DECAY_STEP = np.array(DECAY_EPOCH, np.int32) * SAVE_WEIGHTS_INTE
MAX_ITERATION = SAVE_WEIGHTS_INTE * MAX_EPOCH
WARM_SETP = int(WARM_EPOCH * SAVE_WEIGHTS_INTE)

# dataset
DATASET_NAME = 'DIOR-R'
CLASS_NUM = 20

# model
# backbone
pretrain_zoo = PretrainModelZoo()
PRETRAINED_CKPT = pretrain_zoo.pretrain_weight_path(NET_NAME, ROOT_PATH)
TRAINED_CKPT = os.path.join(ROOT_PATH, 'output/trained_weights')

# bbox head
NUM_SUBNET_CONV = 4
LEVEL = ['P3', 'P4', 'P5', 'P6', 'P7']
BASE_ANCHOR_SIZE_LIST = [32, 64, 128, 256, 512]
ANCHOR_STRIDE = [8, 16, 32, 64, 128]
ANCHOR_SCALES = [2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]
ANCHOR_RATIOS = [1, 1 / 2, 2.]

# loss
CLS_WEIGHT = 1.0
REG_WEIGHT = 2.0

BCD_TAU = 2.0
BCD_FUNC = 0   # 0: sqrt  1: log

VERSION = 'RetinaNet_DIOR_R_BCD_2x_20211024'

"""
RetinaNet-H + bcd
FLOPs: 844274461;    Trainable params: 32553441
cls : overpass|| Recall: 0.611672278338945 || Precison: 0.07005141388174807|| AP: 0.4526383158758531
F1:0.5451979074814559 P:0.7122171945701358 R:0.44163860830527496
cls : windmill|| Recall: 0.6751167444963309 || Precison: 0.15721609445393817|| AP: 0.5436856585446005
F1:0.6360982575589883 P:0.7443004023245418 R:0.5553702468312208
cls : baseballfield|| Recall: 0.7641234711706465 || Precison: 0.32019524100061014|| AP: 0.6923536454780171
F1:0.7532419323820472 P:0.9290901324220419 R:0.6333721607454863
cls : harbor|| Recall: 0.5639291465378422 || Precison: 0.03356206394234455|| AP: 0.31420910724950163
F1:0.4095153613078025 P:0.49520328917313844 R:0.34911433172302736
cls : tenniscourt|| Recall: 0.8744382405011576 || Precison: 0.4047528996469995|| AP: 0.8110433253445409
F1:0.8676657392756655 P:0.9367008681925809 R:0.8081165736075173
cls : airplane|| Recall: 0.6412566975158305 || Precison: 0.4418526598422554|| AP: 0.5975349782068682
F1:0.6734722489036558 P:0.8686285824196961 R:0.5499269361909401
cls : golffield|| Recall: 0.8747826086956522 || Precison: 0.0935639880952381|| AP: 0.7502881503861119
F1:0.7904361927386295 P:0.8382066276803118 R:0.7478260869565218
cls : Expressway-Service-area|| Recall: 0.791705069124424 || Precison: 0.05264770777151263|| AP: 0.6786518066431075
F1:0.7454211226836666 P:0.8611111111111112 R:0.6571428571428571
cls : basketballcourt|| Recall: 0.8984156570363467 || Precison: 0.11780520591470121|| AP: 0.8107732610296795
F1:0.8821148811291772 P:0.9413319238900634 R:0.8299161230195713
cls : vehicle|| Recall: 0.3073948948948949 || Precison: 0.19271862938906148|| AP: 0.29007103798551087
F1:0.377051014880233 P:0.6888136268552645 R:0.25957207207207206
cls : Expressway-toll-station|| Recall: 0.5813953488372093 || Precison: 0.05137426149499101|| AP: 0.5359390348236636
F1:0.6461778907214785 P:0.8989637305699482 R:0.5043604651162791
cls : ship|| Recall: 0.7124992894901381 || Precison: 0.4578577298876815|| AP: 0.6549029037108891
F1:0.7036762760133648 P:0.8045195261079421 R:0.6253055192406071
cls : dam|| Recall: 0.5037174721189591 || Precison: 0.03008436944937833|| AP: 0.23852838174675034
F1:0.3550371017028158 P:0.4082125603864734 R:0.3141263940520446
cls : airport|| Recall: 0.5375375375375375 || Precison: 0.03031842818428184|| AP: 0.3047949436673624
F1:0.4312665818744036 P:0.5040160642570282 R:0.3768768768768769
cls : storagetank|| Recall: 0.5097384529771842 || Precison: 0.43921510770138683|| AP: 0.4917128026890021
F1:0.5836446519412516 P:0.8145847711065103 R:0.45473224605111084
cls : bridge|| Recall: 0.399382000772499 || Precison: 0.04917487040471774|| AP: 0.24931145564511167
F1:0.3417464857726594 P:0.4393203883495146 R:0.27964465044418696
cls : stadium|| Recall: 0.8258928571428571 || Precison: 0.12139107611548557|| AP: 0.646367544697447
F1:0.6551390893224737 P:0.7329650092081031 R:0.5922619047619048
cls : trainstation|| Recall: 0.5343811394891945 || Precison: 0.04486227939963714|| AP: 0.3540045872681856
F1:0.4478325857867968 P:0.6353790613718412 R:0.34577603143418467
cls : chimney|| Recall: 0.7914645974781765 || Precison: 0.08828302499188576|| AP: 0.7238040993472865
F1:0.819739972866782 P:0.9572538860103627 R:0.7167798254122212
cls : groundtrackfield|| Recall: 0.8875331564986737 || Precison: 0.16321951219512196|| AP: 0.7318400438478007
F1:0.7716697453802182 P:0.7834882449425916 R:0.7602122015915119
mAP is : 0.5436227542093646
"""
