from __future__ import annotations

# Learning rates
DEFAULT_ACTOR_LR = 0.0003
DEFAULT_CRITIC_LR = 0.003
DEFAULT_CRITIC_LR_END = None

# Optimizer
DEFAULT_OPTIMIZER_STRING = 'adam'

# Estimator
DEFAULT_ESTIMATOR_STRING = '<default>'
DEFAULT_ANN_N_NEURONS_BY_LAYER = (8,)

# Ego-attention estimator
DEFAULT_EGO_ATTENTION_EMBEDDING_SIZES = (64, 64)
DEFAULT_EGO_ATTENTION_FEATURE_SIZE = 64
DEFAULT_EGO_ATTENTION_HEADS = 4
DEFAULT_EGO_ATTENTION_N_LAYERS = 2
DEFAULT_EGO_ATTENTION_OUTPUT_SIZES = (64, 64)
DEFAULT_EGO_ATTENTION_CRITIC_OUTPUT_SIZES = (192, 192)

# PPO
DEFAULT_PPO_CLIP_EPSILON = 0.2
DEFAULT_MAX_GRAD_NORM = 0.5
DEFAULT_N_SWEEPS_PER_EPOCH = 6
DEFAULT_N_MTS_PER_MINIBATCH = 64
DEFAULT_MINIBATCH_SEED_NUMBER = 0

# Value function
DEFAULT_DISCOUNT = 0.95
DEFAULT_VANILLA_GAE_LAMBDA = 0.9
DEFAULT_CRITIC_LAMBDA = 1
DEFAULT_N_CRITIC_ITERATIONS = 10

# Entropy
DEFAULT_ENTROPY_TEMPERATURE = 0.05
DEFAULT_LOGIT_BASED_ENTROPY = False
DEFAULT_LOGIT_VARIANCE_THRESHOLD = 0
DEFAULT_LOGIT_VARIANCE_LAMBDA = 1.0
DEFAULT_ACTOR_LOGIT_CLIP = 5.0

# Noise
DEFAULT_NOISE = 0.0

# KLD
DEFAULT_TARGET_KLD = None
DEFAULT_TARGET_MAX_KLD = None
DEFAULT_TARGET_MIN_KLD = None
DEFAULT_KLD_METHOD = 'binary-search'

# Freezing
DEFAULT_FREEZE_FROM_EPOCH = None
DEFAULT_FREEZE_ACTOR_FROM_EPOCH = None
DEFAULT_FREEZE_CRITIC_FROM_EPOCH = None

# Training
DEFAULT_N_EPOCHS = 300
DEFAULT_N_TS_PER_E = 40
DEFAULT_N_ES_PER_EPOCH = 400
DEFAULT_SEED_NUMBER = 0
DEFAULT_USE_JIT = True
DEFAULT_REPLAY_BUFFER_FACTOR = 1

# Progress bar
DEFAULT_TQDM_KWARGS = {
    'colour': 'cyan',
    'ascii': '_▒▓█',
    'ncols': 70,
}
