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

import numpy as np

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

# 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
REG_LOSS_MODE = 3  # KLD loss

KL_TAU = 1.0
KL_FUNC = 1   # 0: sqrt  1: log

VERSION = 'RetinaNet_DIOR_R_KL_2x_20211022'

"""
RetinaNet-H + kl (fix bug)
FLOPs: 844274470;    Trainable params: 32553441

cls : bridge|| Recall: 0.42217072228659713 || Precison: 0.055893633341856305|| AP: 0.27817479733331735
F1:0.3603380752712262 P:0.4919786096256685 R:0.2842796446504442
cls : airport|| Recall: 0.5870870870870871 || Precison: 0.029587589860007566|| AP: 0.33227801047865985
F1:0.46311952026625164 P:0.4907563025210084 R:0.43843843843843844
cls : harbor|| Recall: 0.5803542673107891 || Precison: 0.03340067839335693|| AP: 0.33289395155926277
F1:0.418530795347702 P:0.49301310043668123 R:0.3636070853462158
cls : storagetank|| Recall: 0.4986515988185437 || Precison: 0.4386744492562606|| AP: 0.43832321724192985
F1:0.5664582196937289 P:0.8012840588787974 R:0.4380805616197937
cls : chimney|| Recall: 0.8215324927255092 || Precison: 0.08113026819923372|| AP: 0.7544661875855199
F1:0.8342371710884863 P:0.9526774595267746 R:0.7419980601357905
cls : groundtrackfield|| Recall: 0.9135278514588859 || Precison: 0.14|| AP: 0.7745722540291036
F1:0.7923255444715267 P:0.806930693069307 R:0.7782493368700265
cls : basketballcourt|| Recall: 0.8998136067101584 || Precison: 0.1706885883496862|| AP: 0.8124658077104565
F1:0.890031915988874 P:0.9426784783741532 R:0.8429636533084809
cls : vehicle|| Recall: 0.31077327327327325 || Precison: 0.15576962877946904|| AP: 0.28520505533325197
F1:0.3686579129287122 P:0.6933043117744611 R:0.25108858858858857
cls : trainstation|| Recall: 0.6286836935166994 || Precison: 0.03669724770642202|| AP: 0.37408720921910726
F1:0.43606819079787673 P:0.5204359673024523 R:0.37524557956778
cls : overpass|| Recall: 0.632996632996633 || Precison: 0.06509695290858726|| AP: 0.4704526143462545
F1:0.5421929857574642 P:0.7020517395182873 R:0.44163860830527496
cls : Expressway-Service-area|| Recall: 0.8617511520737328 || Precison: 0.06569702079820124|| AP: 0.7314873504522715
F1:0.7627762459754709 P:0.8564867967853043 R:0.687557603686636
cls : ship|| Recall: 0.7267094867276758 || Precison: 0.41893308866898227|| AP: 0.6602782024747347
F1:0.7072148968391878 P:0.8143608354317878 R:0.6249928949013812
cls : airplane|| Recall: 0.6335849975645397 || Precison: 0.40317706315381635|| AP: 0.5912174765252652
F1:0.6675602007272211 P:0.8663299009131533 R:0.5429858743302484
cls : dam|| Recall: 0.5650557620817844 || Precison: 0.025595689147091018|| AP: 0.2676282117367076
F1:0.3642335866999512 P:0.4483695652173913 R:0.3066914498141264
cls : stadium|| Recall: 0.8452380952380952 || Precison: 0.10999225406661503|| AP: 0.6555698976390658
F1:0.6600282330014218 P:0.7453183520599251 R:0.5922619047619048
cls : baseballfield|| Recall: 0.7644146767617939 || Precison: 0.27332361516034986|| AP: 0.6891116493831928
F1:0.7497757213676985 P:0.9442724458204335 R:0.6217239370995923
cls : Expressway-toll-station|| Recall: 0.5988372093023255 || Precison: 0.04584399688438856|| AP: 0.5316473000683527
F1:0.6350047063645139 P:0.883419689119171 R:0.4956395348837209
cls : tenniscourt|| Recall: 0.8823369195151846 || Precison: 0.38285174023518287|| AP: 0.811205509624372
F1:0.8714597572153628 P:0.9293427954335082 R:0.8203731444913523
cls : golffield|| Recall: 0.8747826086956522 || Precison: 0.08580689184578642|| AP: 0.765616421580607
F1:0.8022443273804575 P:0.8699186991869918 R:0.7443478260869565
cls : windmill|| Recall: 0.6657771847898599 || Precison: 0.1589171974522293|| AP: 0.5442803882587824
F1:0.6367857570032447 P:0.7703598484848485 R:0.5426951300867244
mAP is : 0.5550480756290106
"""
