from yacs.config import CfgNode
from yacs.config import CfgNode as CN
from yacs_stubgen import build_pyi

# YACS overwrite these settings using YAML, all YAML variables MUST BE defined here first
# as this is the master list of ALL attributes.
_C = CN()

# importing default as a global singleton
_C.DESCRIPTION = "Default config"

_C.EXPERIMENTAL = CN()
# _C.EXPERIMENTAL.USE_PREV_IMAGE = False
# _C.EXPERIMENTAL.USE_EVEN_PREV_IMAGE = False
_C.EXPERIMENTAL.SHUFFLE_IMAGES = False
_C.EXPERIMENTAL.BLANK_IMAGE = False
_C.EXPERIMENTAL.T_IMAGE = 0
_C.EXPERIMENTAL.USE_RETINA_MAPPER = True
_C.EXPERIMENTAL.USE_LAYER_SELECTOR = True
_C.EXPERIMENTAL.USE_BHV = True
_C.EXPERIMENTAL.USE_BHV_PASSTHROUGH = True
_C.EXPERIMENTAL.BEHV_ONLY = False
_C.EXPERIMENTAL.BEHV_SELECTION = [-1]
_C.EXPERIMENTAL.BACKBONE_NOGRAD = False
_C.EXPERIMENTAL.STRAIGHT_FORWARD = False  # deprecated, use BACKBONE_NOGRAD
_C.EXPERIMENTAL.STRAIGHT_FORWARD_BUT_KEEP_BACKBONE_GRAD = False  # deprecated
_C.EXPERIMENTAL.ANOTHER_SPLIT = False
_C.EXPERIMENTAL.SHUFFLE_VAL = True
_C.EXPERIMENTAL.NO_SPLIT = False
_C.EXPERIMENTAL.USE_DEV_MODEL = False
_C.EXPERIMENTAL.USE_FEAT = False
_C.EXPERIMENTAL.CENTER_FRAME = -1
_C.EXPERIMENTAL.REPLACE_BLANK_WITH_RANDOM = False
_C.EXPERIMENTAL.REPLACE_NONBLANK_WITH_RANDOM = False
_C.EXPERIMENTAL.USE_PREV_FRAME = True
_C.EXPERIMENTAL.USE_PREV_BEHV = True
_C.EXPERIMENTAL.USE_CURRENT_FRAME = True
_C.EXPERIMENTAL.USE_CURRENT_BEHV = True
_C.EXPERIMENTAL.USE_FTR_FRAME = True
_C.EXPERIMENTAL.USE_FTR_BEHV = True
_C.EXPERIMENTAL.USE_COORDS = True
_C.EXPERIMENTAL.MAX_DATASET_SIZE = -1

_C.DATAMODULE = CN()
_C.DATAMODULE.BATCH_SIZE = 32
_C.DATAMODULE.NUM_WORKERS = 8
_C.DATAMODULE.PIN_MEMORY = True

_C.DATAMODULE.FEATURE_EXTRACTOR_MODE = False

_C.DATASET = CN()
_C.DATASET.IMAGE_RESOLUTION = [224, 224]
_C.DATASET.N_PREV_FRAMES = 32
_C.DATASET.N_FTR_FRAMES = 32
_C.DATASET.CACHE_DIR = "/data/cache"
_C.DATASET.SUBJECT_LIST = ["subj01"]
_C.DATASET.ROIS = ["all"]
_C.DATASET.FMRI_SPACE = "fsaverage"
_C.DATASET.FILTER_BY_SESSION = [-1]
# _C.DATASET.FILTER_BY_SESSION = 'all'
_C.DATASET.ROOT = "/data/ALG23"
_C.DATASET.DARK_POSTFIX = ""

_C.POSITION_ENCODING = CN()
_C.POSITION_ENCODING.IN_DIM = 3
_C.POSITION_ENCODING.MAX_STEPS = 1000
_C.POSITION_ENCODING.FEATURES = 32
_C.POSITION_ENCODING.PERIODS = 10000


_C.MODEL = CN()

_C.MODEL.WIDTH_RATIO = 1.0

_C.MODEL.BACKBONE = CN()
_C.MODEL.BACKBONE.NAME = "dinov2_vit_l"
_C.MODEL.BACKBONE.CACHE_DIR = "/data/cache"
_C.MODEL.BACKBONE.LAYERS = [5, 11, 17, 23]
_C.MODEL.BACKBONE.FEATURE_DIMS = [1024, 1024, 1024, 1024]
_C.MODEL.BACKBONE.CLS_DIMS = [2048, 2048, 2048, 1024]

_C.MODEL.BACKBONE.LORA = CN()
_C.MODEL.BACKBONE.LORA.SCALE = 0.2
_C.MODEL.BACKBONE.LORA.RANK = 4

_C.MODEL.BACKBONE.ADAPTIVE_LN = CN()
_C.MODEL.BACKBONE.ADAPTIVE_LN.SCALE = 0.5

_C.MODEL.BACKBONE_SMALL = CN()
_C.MODEL.BACKBONE_SMALL.NAME = "dinov2_vit_b"
_C.MODEL.BACKBONE_SMALL.LAYERS = [5, 8, 11]
_C.MODEL.BACKBONE_SMALL.CLS_DIMS = [1536, 1536, 768]
_C.MODEL.BACKBONE_SMALL.T_DIM = 128
_C.MODEL.BACKBONE_SMALL.WIDTH = 128
_C.MODEL.BACKBONE_SMALL.MERGE_WIDTH = 128

_C.MODEL.BACKBONE_SMALL.LORA = CN()
_C.MODEL.BACKBONE_SMALL.LORA.SCALE = 0.2
_C.MODEL.BACKBONE_SMALL.LORA.RANK = 4

_C.MODEL.BACKBONE_SMALL.ADAPTIVE_LN = CN()
_C.MODEL.BACKBONE_SMALL.ADAPTIVE_LN.SCALE = 0.5

_C.MODEL.PREV_FEAT = CN()
_C.MODEL.PREV_FEAT.DIM = 1024

_C.MODEL.CONV_HEAD = CN()
_C.MODEL.CONV_HEAD.MAX_DIM = 1024
_C.MODEL.CONV_HEAD.KERNEL_SIZES = [5, 5, 5, 5]
_C.MODEL.CONV_HEAD.DEPTHS = [3, 3, 3, 3]
_C.MODEL.CONV_HEAD.WIDTH = 256
_C.MODEL.CONV_HEAD.SIMPLE = True

_C.MODEL.COND = CN()
_C.MODEL.COND.USE = True
_C.MODEL.COND.DROPOUT = 0.2
_C.MODEL.COND.IN_DIM = 35
_C.MODEL.COND.DIM = 256
_C.MODEL.COND.PASSTHROUGH_DIM = 64

_C.MODEL.MAX_TRAIN_VOXELS = 25600
_C.MODEL.CHUNK_SIZE = 25600

_C.MODEL.COORDS_MLP = CN()
_C.MODEL.COORDS_MLP.WIDTH = 128
_C.MODEL.COORDS_MLP.DEPTH = 3
_C.MODEL.COORDS_MLP.LOG = True

_C.MODEL.RETINA_MAPPER = CN()
_C.MODEL.RETINA_MAPPER.CONSTANT_SIGMA = 0.01

_C.MODEL.LAYER_SELECTOR = CN()

_C.MODEL.BOTTLENECK = CN()
_C.MODEL.BOTTLENECK.RANK = -1
_C.MODEL.BOTTLENECK.OUT_DIM = 64

_C.MODEL.VOXEL_OUTS = CN()
_C.MODEL.VOXEL_OUTS.SHARED = CN()
_C.MODEL.VOXEL_OUTS.SHARED.USE = False
_C.MODEL.VOXEL_OUTS.SHARED.MLP = CN()
_C.MODEL.VOXEL_OUTS.SHARED.MLP.DEPTH = 3
_C.MODEL.VOXEL_OUTS.SHARED.MLP.WIDTH = 1024

_C.LOSS = CN()
_C.LOSS.NAME = "SmoothL1Loss"
_C.LOSS.SMOOTH_L1_BETA = 0.01

_C.LOSS.SYNC = CN()
_C.LOSS.SYNC.USE = False
_C.LOSS.SYNC.STAGE = "VAL"
_C.LOSS.SYNC.SKIP_EPOCHS = 10
_C.LOSS.SYNC.EMA_BETA = 0.9
_C.LOSS.SYNC.EMA_BIAS_CORRECTION = False
_C.LOSS.SYNC.UPDATE_RULE = "exp"
_C.LOSS.SYNC.EXP_SCALE = 10.0
_C.LOSS.SYNC.EXP_SHIFT = 0.0
_C.LOSS.SYNC.LOG_SHIFT = 10.0
_C.LOSS.SYNC.EMA_KEY = "running_grad"

_C.LOSS.DARK = CN()
_C.LOSS.DARK.USE = False
_C.LOSS.DARK.MAX_EPOCH = 100
_C.LOSS.DARK.GT_ROIS = ["htroi_1"]
_C.LOSS.DARK.GT_SCALE_UP_COEF = 1.0

_C.LOSS.DARK.ANNEAL = CN()
_C.LOSS.DARK.ANNEAL.T = 30

_C.REGULARIZER = CN()
_C.REGULARIZER.LAYER = 3e-5

_C.OPTIMIZER = CN()
_C.OPTIMIZER.NAME = "AdamW"
_C.OPTIMIZER.LR = 3e-4
_C.OPTIMIZER.WEIGHT_DECAY = 3e-4

_C.OPTIMIZER.SCHEDULER = CN()
_C.OPTIMIZER.SCHEDULER.T_INITIAL = 1
_C.OPTIMIZER.SCHEDULER.T_MULT = 1.0
_C.OPTIMIZER.SCHEDULER.CYCLE_DECAY = 0.5
_C.OPTIMIZER.SCHEDULER.CYCLE_LIMIT = 3
_C.OPTIMIZER.SCHEDULER.WARMUP_T = 3
_C.OPTIMIZER.SCHEDULER.K_DECAY = 1.5
_C.OPTIMIZER.SCHEDULER.LR_MIN = 3e-4
_C.OPTIMIZER.SCHEDULER.LR_MIN_WARMUP = 1e-4

_C.TRAINER = CN()
_C.TRAINER.DDP = False
_C.TRAINER.PRECISION = 16
_C.TRAINER.GRADIENT_CLIP_VAL = 0.5
_C.TRAINER.MAX_EPOCHS = 1000
_C.TRAINER.MAX_STEPS = -1
_C.TRAINER.ACCUMULATE_GRAD_BATCHES = 1
_C.TRAINER.VAL_CHECK_INTERVAL = 1.0
_C.TRAINER.LIMIT_TRAIN_BATCHES = 0.1
_C.TRAINER.LIMIT_VAL_BATCHES = 0.5
_C.TRAINER.LOG_TRAIN_N_STEPS = 100

_C.TRAINER.CALLBACKS = CN()

_C.TRAINER.CALLBACKS.EARLY_STOP = CN()
_C.TRAINER.CALLBACKS.EARLY_STOP.PATIENCE = 30

_C.TRAINER.CALLBACKS.CHECKPOINT = CN()
_C.TRAINER.CALLBACKS.CHECKPOINT.SAVE_TOP_K = 10
_C.TRAINER.CALLBACKS.CHECKPOINT.REMOVE = True
_C.TRAINER.CALLBACKS.CHECKPOINT.LOAD_BEST_ON_VAL = False
_C.TRAINER.CALLBACKS.CHECKPOINT.LOAD_BEST_ON_END = False

_C.MODEL_SOUP = CN()
_C.MODEL_SOUP.USE = True
_C.MODEL_SOUP.RECIPE = "greedy"
_C.MODEL_SOUP.GREEDY_TARGET = "heldout"

_C.RESULTS_DIR = "/nfscc/alg23/ray_results/"
_C.CHECKPOINT_DIR = "/data/ckpt/"

_C.ANALYSIS = CN()
_C.ANALYSIS.SAVE_NEURON_LOCATION = False
_C.ANALYSIS.DRAW_NEURON_LOCATION = False

# this alias ensure you can import `AutoConfig` and use something like `isinstance`
AutoConfig = CN
# _C is the CfgNode object, "_C" should be its varname correctly
# AutoConfig is an alias of CfgNode, "AutoConfig" should be its varname correctly
build_pyi(_C, __file__, var_name="_C")