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

import numpy as np
from configs._base_.models.faster_rcnn_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 = 0.001 * BATCH_SIZE * NUM_GPU
SAVE_WEIGHTS_INTE = 11725 * 2
DECAY_EPOCH = [16, 22, 40]
MAX_EPOCH = 24
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
IMG_SHORT_SIDE_LEN = [800, 450, 500, 640, 700, 900, 1000, 1100, 1200]
IMG_MAX_LENGTH = 1200
# data augmentation
IMG_ROTATE = True
RGB2GRAY = True
VERTICAL_FLIP = True
HORIZONTAL_FLIP = True
IMAGE_PYRAMID = True

# 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
LEVEL = ['P2', 'P3', 'P4', 'P5', 'P6']
BASE_ANCHOR_SIZE_LIST = [32, 64, 128, 256, 512]
ANCHOR_STRIDE = [4, 8, 16, 32, 64]
ANCHOR_SCALES = [1.0]
ANCHOR_RATIOS = [0.5, 1., 2.0]

# loss
FAST_RCNN_LOCATION_LOSS_WEIGHT = 1.0
FAST_RCNN_CLASSIFICATION_LOSS_WEIGHT = 1.0

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

VERSION = 'FPN_Res50D_DIOR_R_R2CNN_KL_4x_20211028'

"""
R2CNN + KL
FLOPs: 1024386459;    Trainable params: 41791120
cls : airplane|| Recall: 0.8555772040915733 || Precison: 0.8483458101907752|| AP: 0.8100530969944938
F1:0.8712282525398065 P:0.9311538994320543 R:0.8185582075012178
cls : airport|| Recall: 0.6696696696696697 || Precison: 0.18353909465020576|| AP: 0.4531321023884684
F1:0.5892525175807356 P:0.6216666666666667 R:0.56006006006006
cls : baseballfield|| Recall: 0.8459522422830519 || Precison: 0.7149889244400689|| AP: 0.8027218249372241
F1:0.8415900234734689 P:0.9088145896656535 R:0.783634245777519
cls : basketballcourt|| Recall: 0.9268406337371855 || Precison: 0.6338432122370937|| AP: 0.8997873182733191
F1:0.9250733001828865 P:0.9576059850374065 R:0.8946877912395154
cls : bridge|| Recall: 0.5948242564696794 || Precison: 0.12874101320849357|| AP: 0.4280630858998197
F1:0.501623874146289 P:0.6307782328847279 R:0.4163769795287756
cls : chimney|| Recall: 0.7807953443258971 || Precison: 0.5028107432854466|| AP: 0.7266402645307253
F1:0.8359326135352704 P:0.9842312746386334 R:0.7264791464597479
cls : dam|| Recall: 0.6598513011152416 || Precison: 0.12050237610319077|| AP: 0.35186626069368476
F1:0.4635965373183971 P:0.4782608695652174 R:0.44981412639405205
cls : Expressway-Service-area|| Recall: 0.9078341013824884 || Precison: 0.3335590924483576|| AP: 0.8321678407649732
F1:0.8657245335955783 P:0.9130879345603272 R:0.8230414746543778
cls : Expressway-toll-station|| Recall: 0.8284883720930233 || Precison: 0.2844311377245509|| AP: 0.7553217366361984
F1:0.7900987992098923 P:0.8761061946902655 R:0.7194767441860465
cls : golffield|| Recall: 0.8782608695652174 || Precison: 0.39177657098525986|| AP: 0.7910863841137586
F1:0.8404967967852546 P:0.866913123844732 R:0.8156521739130435
cls : groundtrackfield|| Recall: 0.946949602122016 || Precison: 0.44259856186461694|| AP: 0.8395162701079597
F1:0.8148886170873214 P:0.8170666666666667 R:0.8127320954907162
cls : harbor|| Recall: 0.6231884057971014 || Precison: 0.17629373177842567|| AP: 0.46797251420262653
F1:0.5366525408732694 P:0.629004329004329 R:0.4679549114331723
cls : overpass|| Recall: 0.6857463524130191 || Precison: 0.16125626814462918|| AP: 0.5591116560607017
F1:0.623946885341405 P:0.7756463719766472 R:0.5218855218855218
cls : ship|| Recall: 0.8871141931450008 || Precison: 0.7684959499716867|| AP: 0.8066332166451309
F1:0.8869926751604764 P:0.9307728337236534 R:0.8471551185130449
cls : stadium|| Recall: 0.8601190476190477 || Precison: 0.5026086956521739|| AP: 0.7468239465113075
F1:0.741319939131293 P:0.7885906040268457 R:0.6994047619047619
cls : storagetank|| Recall: 0.8175163734429177 || Precison: 0.717214961694457|| AP: 0.7908083452520556
F1:0.8305175214308274 P:0.8964337086454045 R:0.7736398270621977
cls : tenniscourt|| Recall: 0.9118888737573199 || Precison: 0.7783331396024643|| AP: 0.8919399332691789
F1:0.908645431226984 P:0.9529218353011508 R:0.8683099550592401
cls : trainstation|| Recall: 0.7426326129666012 || Precison: 0.16050955414012738|| AP: 0.5726280023650793
F1:0.6345242477856353 P:0.7389033942558747 R:0.555992141453831
cls : vehicle|| Recall: 0.6368993993993994 || Precison: 0.23898529494619417|| AP: 0.5365339085238283
F1:0.5848837102499993 P:0.7427745664739884 R:0.48235735735735735
cls : windmill|| Recall: 0.8162108072048032 || Precison: 0.48910653607835297|| AP: 0.7175211669245801
F1:0.788887867095528 P:0.8458015267175573 R:0.7391594396264176
mAP is : 0.6890164437547557

ms
cls : airplane|| Recall: 0.9119581100828057 || Precison: 0.616327874249033|| AP: 0.8777170256954362
F1:0.8837579035898456 P:0.9384236453201971 R:0.8351193375547978
cls : airport|| Recall: 0.6021021021021021 || Precison: 0.16919831223628692|| AP: 0.4281768703953834
F1:0.5924082646931765 P:0.6404886561954625 R:0.551051051051051
cls : baseballfield|| Recall: 0.8980780430984275 || Precison: 0.46933495662760616|| AP: 0.8068954493528616
F1:0.858361485694677 P:0.9039948453608248 R:0.8171228887594641
cls : basketballcourt|| Recall: 0.9431500465983225 || Precison: 0.36344047405279223|| AP: 0.9001243694997731
F1:0.9265156910966668 P:0.9694501018329938 R:0.8872320596458527
cls : bridge|| Recall: 0.6183854770181537 || Precison: 0.09139171138257791|| AP: 0.4489023056826419
F1:0.5150009008754076 P:0.6328828828828829 R:0.43414445731942836
cls : chimney|| Recall: 0.7953443258971872 || Precison: 0.28731604765241764|| AP: 0.7270266945503753
F1:0.8408790748492599 P:0.9768934531450578 R:0.7381183317167799
cls : dam|| Recall: 0.5371747211895911 || Precison: 0.08266590389016018|| AP: 0.3458757546578058
F1:0.49599502892982217 P:0.5367965367965368 R:0.46096654275092935
cls : Expressway-Service-area|| Recall: 0.9142857142857143 || Precison: 0.232427366447985|| AP: 0.8587819872837339
F1:0.8871769578622174 P:0.9258517034068137 R:0.8516129032258064
cls : Expressway-toll-station|| Recall: 0.8372093023255814 || Precison: 0.1453444360333081|| AP: 0.7619892654925362
F1:0.7958647810988424 P:0.8774080560420315 R:0.7281976744186046
cls : golffield|| Recall: 0.8243478260869566 || Precison: 0.3062015503875969|| AP: 0.7741286806124805
F1:0.8219647366181174 P:0.9022869022869023 R:0.7547826086956522
cls : groundtrackfield|| Recall: 0.9718832891246685 || Precison: 0.24092582851130984|| AP: 0.8418129397110325
F1:0.8148098166875944 P:0.8307607497243661 R:0.7994694960212202
cls : harbor|| Recall: 0.5922705314009662 || Precison: 0.21786518185049164|| AP: 0.45459932605957964
F1:0.562612666097958 P:0.676005422503389 R:0.4818035426731079
cls : overpass|| Recall: 0.6750841750841751 || Precison: 0.10949303722581233|| AP: 0.5629170981671051
F1:0.6332725069475366 P:0.7896060352053647 R:0.5286195286195287
cls : ship|| Recall: 0.9224975842664697 || Precison: 0.6661125818301217|| AP: 0.8885049148118623
F1:0.9056001559637421 P:0.9335022930243785 R:0.8793270050588302
cls : stadium|| Recall: 0.9181547619047619 || Precison: 0.32051948051948054|| AP: 0.7863583393351388
F1:0.7519210463744598 P:0.7795527156549521 R:0.7261904761904762
cls : storagetank|| Recall: 0.864218141346689 || Precison: 0.5641591683898731|| AP: 0.7980589554402355
F1:0.8525992410779699 P:0.9029386426725375 R:0.8075852917255254
cls : tenniscourt|| Recall: 0.9340868854691543 || Precison: 0.5549801763896756|| AP: 0.8932557056302098
F1:0.9104291151829436 P:0.9417758369723436 R:0.881111262426801
cls : trainstation|| Recall: 0.7131630648330058 || Precison: 0.12795206203736342|| AP: 0.5725679750397001
F1:0.6446054963142697 P:0.705607476635514 R:0.593320235756385
cls : vehicle|| Recall: 0.6862987987987988 || Precison: 0.17915552027907614|| AP: 0.5601817375108809
F1:0.6089219538432563 P:0.742772223723318 R:0.5159534534534534
cls : windmill|| Recall: 0.8042028018679119 || Precison: 0.44238532110091744|| AP: 0.7218977247116922
F1:0.7936179980439408 P:0.8357933579335793 R:0.7555036691127418
mAP is : 0.7004886559820231
"""
