# Copyright (c) Tianqi Shen. All Rights Reserved.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from yacs.config import CfgNode as CN

__C = CN()

cfg = __C

__C.META_ARC = "Topology-GS"

__C.CUDA = True

# ------------------------------------------------------------------------ #
# Pipeline options
# ------------------------------------------------------------------------ #
__C.PIPE = CN()
__C.PIPE.CONVERT_SHS_PYTHON = False
__C.PIPE.COMPUTE_COV3D_PYTHON = False
__C.PIPE.DEBUG = False

# ------------------------------------------------------------------------ #
# Optimization options
# ------------------------------------------------------------------------ #
__C.OPTIM = CN(new_allowed=True)

__C.OPTIM.ITERATIONS = 30_000
__C.OPTIM.PERCENT_DENSE = 0.01
__C.OPTIM.LAMBDA_DSSIM = 0.2

# persistence loss
__C.OPTIM.USE_PERS_LOSS = False
__C.OPTIM.PERSISTENCE = CN()
__C.OPTIM.PERSISTENCE.RATIO = 0.1
__C.OPTIM.PERSISTENCE.MODE = 'area'
__C.OPTIM.PERSISTENCE.LAMBDA = 10.0
__C.OPTIM.PERSISTENCE.START = 1500
__C.OPTIM.PERSISTENCE.END = 15_000
__C.OPTIM.PERSISTENCE.INTERVAL = 200
__C.OPTIM.PERSISTENCE.DIMS = [0, 1, 2]
__C.OPTIM.PERSISTENCE.KS = [100, 100, 100]

__C.OPTIM.TEST_ITERATIONS = [30_000]
__C.OPTIM.SAVE_ITERATIONS = [30_000]

# for anchor densification
__C.OPTIM.START_STAT = 500
__C.OPTIM.UPDATE_FROM = 1500
__C.OPTIM.UPDATE_INTERVAL = 100
__C.OPTIM.UPDATE_UNTIL = 15_000

__C.OPTIM.MIN_OPACITY = 0.005
__C.OPTIM.SUCCESS_THRESHOLD = 0.8
__C.OPTIM.DENSIFY_GRAD_THRESHOLD = 0.0002

# for learning rate
__C.OPTIM.FEATURE_LR = 0.0075
__C.OPTIM.OPACITY_LR = 0.02
__C.OPTIM.SCALING_LR = 0.007
__C.OPTIM.ROTATION_LR = 0.002

__C.OPTIM.POSITION = CN()
__C.OPTIM.POSITION.LR_INIT = 0.0
__C.OPTIM.POSITION.LR_FINAL = 0.0
__C.OPTIM.POSITION.LR_DELAY_MULT = 0.01

__C.OPTIM.OFFSET = CN()
__C.OPTIM.OFFSET.LR_INIT = 0.01
__C.OPTIM.OFFSET.LR_FINAL = 0.0001
__C.OPTIM.OFFSET.LR_DELAY_MULT = 0.01

__C.OPTIM.MLP_OPACITY = CN()
__C.OPTIM.MLP_OPACITY.LR_INIT = 0.002
__C.OPTIM.MLP_OPACITY.LR_FINAL = 0.00002
__C.OPTIM.MLP_OPACITY.LR_DELAY_MULT = 0.01

__C.OPTIM.MLP_COV = CN()
__C.OPTIM.MLP_COV.LR_INIT = 0.004
__C.OPTIM.MLP_COV.LR_FINAL = 0.004
__C.OPTIM.MLP_COV.LR_DELAY_MULT = 0.01

__C.OPTIM.MLP_COLOR = CN()
__C.OPTIM.MLP_COLOR.LR_INIT = 0.008
__C.OPTIM.MLP_COLOR.LR_FINAL = 0.00005
__C.OPTIM.MLP_COLOR.LR_DELAY_MULT = 0.01

# optional
__C.OPTIM.MLP_FEATUREBANK = CN()
__C.OPTIM.MLP_FEATUREBANK.LR_INIT = 0.01
__C.OPTIM.MLP_FEATUREBANK.LR_FINAL = 0.00001
__C.OPTIM.MLP_FEATUREBANK.LR_DELAY_MULT = 0.01

# optional
__C.OPTIM.APPEARANCE = CN()
__C.OPTIM.APPEARANCE.LR_INIT = 0.05
__C.OPTIM.APPEARANCE.LR_FINAL = 0.0005
__C.OPTIM.APPEARANCE.LR_DELAY_MULT = 0.01

# ------------------------------------------------------------------------ #
# Model options
# ------------------------------------------------------------------------ #
__C.MODEL = CN(new_allowed=True)

__C.MODEL.SH_DEGREE = 3
__C.MODEL.FEAT_DIM = 32
__C.MODEL.N_OFFSETS = 10
__C.MODEL.VOXEL_SIZE = 0.001
__C.MODEL.UPDATE_DEPTH = 3
__C.MODEL.UPDATE_INIT_FACTOR = 16
__C.MODEL.UPDATE_HIERACHY_FACTOR = 4

__C.MODEL.USE_FEAT_BANK = False
__C.MODEL.SOURCE_PATH = ""
__C.MODEL.MODEL_PATH = ""
__C.MODEL.IMAGES = "images"
__C.MODEL.RESOLUTION = -1
__C.MODEL.WHITE_BACKGROUND = False
__C.MODEL.EVAL = False
__C.MODEL.LOD = 0
__C.MODEL.DATA_DEVICE = "cuda"

__C.MODEL.APPEARANCE_DIM = 0
__C.MODEL.LOWPOLY = False
__C.MODEL.DS = 1
__C.MODEL.RATIO = 1
__C.MODEL.UNDISTORTED = False

__C.MODEL.ADD_OPACITY_DIST = False
__C.MODEL.ADD_COV_DIST = False
__C.MODEL.ADD_COLOR_DIST = False

# ------------------------------------------------------------------------ #
# Datasets options
# ------------------------------------------------------------------------ #
__C.DATA = CN()
__C.DATA.INTERPOLATION = False  # for point clouds interpolation
__C.DATA.INTERPOLATION_TYPE = 'LPVI'  # for point clouds interpolation
__C.DATA.NAMES = ('MIPNERF360', 'TANDT', 'DB', 'NERFSYNTHETIC', 'BUNGEENERF', 'IMW2020')

__C.DATA.MIPNERF360 = CN()

__C.DATA.MIPNERF360.BICYCLE = CN()
__C.DATA.MIPNERF360.BICYCLE.PATH = 'data/mipnerf360/bicycle'
__C.DATA.MIPNERF360.BICYCLE.K_MAX = 7
__C.DATA.MIPNERF360.BICYCLE.K_MIN = 3
__C.DATA.MIPNERF360.BICYCLE.THRESHOLD = 0.2

__C.DATA.MIPNERF360.BONSAI = CN()
__C.DATA.MIPNERF360.BONSAI.PATH = 'data/mipnerf360/bonsai'
__C.DATA.MIPNERF360.BONSAI.K_MAX = 7
__C.DATA.MIPNERF360.BONSAI.K_MIN = 3
__C.DATA.MIPNERF360.BONSAI.THRESHOLD = 0.2

__C.DATA.MIPNERF360.COUNTER = CN()
__C.DATA.MIPNERF360.COUNTER.PATH = 'data/mipnerf360/counter'
__C.DATA.MIPNERF360.COUNTER.K_MAX = 7
__C.DATA.MIPNERF360.COUNTER.K_MIN = 3
__C.DATA.MIPNERF360.COUNTER.THRESHOLD = 0.2

__C.DATA.MIPNERF360.FLOWERS = CN()
__C.DATA.MIPNERF360.FLOWERS.PATH = 'data/mipnerf360/flowers'
__C.DATA.MIPNERF360.FLOWERS.K_MAX = 7
__C.DATA.MIPNERF360.FLOWERS.K_MIN = 3
__C.DATA.MIPNERF360.FLOWERS.THRESHOLD = 0.2

__C.DATA.MIPNERF360.GARDEN = CN()
__C.DATA.MIPNERF360.GARDEN.PATH = 'data/mipnerf360/garden'
__C.DATA.MIPNERF360.GARDEN.K_MAX = 7
__C.DATA.MIPNERF360.GARDEN.K_MIN = 3
__C.DATA.MIPNERF360.GARDEN.THRESHOLD = 0.2

__C.DATA.MIPNERF360.KITCHEN = CN()
__C.DATA.MIPNERF360.KITCHEN.PATH = 'data/mipnerf360/kitchen'
__C.DATA.MIPNERF360.KITCHEN.K_MAX = 7
__C.DATA.MIPNERF360.KITCHEN.K_MIN = 3
__C.DATA.MIPNERF360.KITCHEN.THRESHOLD = 0.2

__C.DATA.MIPNERF360.ROOM = CN()
__C.DATA.MIPNERF360.ROOM.PATH = 'data/mipnerf360/room'
__C.DATA.MIPNERF360.ROOM.K_MAX = 7
__C.DATA.MIPNERF360.ROOM.K_MIN = 3
__C.DATA.MIPNERF360.ROOM.THRESHOLD = 0.2

__C.DATA.MIPNERF360.STUMP = CN()
__C.DATA.MIPNERF360.STUMP.PATH = 'data/mipnerf360/stump'
__C.DATA.MIPNERF360.STUMP.K_MAX = 7
__C.DATA.MIPNERF360.STUMP.K_MIN = 3
__C.DATA.MIPNERF360.STUMP.THRESHOLD = 0.2

__C.DATA.MIPNERF360.TREEHILL = CN()
__C.DATA.MIPNERF360.TREEHILL.PATH = 'data/mipnerf360/treehill'
__C.DATA.MIPNERF360.TREEHILL.K_MAX = 7
__C.DATA.MIPNERF360.TREEHILL.K_MIN = 3
__C.DATA.MIPNERF360.TREEHILL.THRESHOLD = 0.2


__C.DATA.TANDT = CN()

__C.DATA.TANDT.TRAIN = CN()
__C.DATA.TANDT.TRAIN.PATH = 'data/tandt/train'
__C.DATA.TANDT.TRAIN.K_MAX = 7
__C.DATA.TANDT.TRAIN.K_MIN = 3
__C.DATA.TANDT.TRAIN.THRESHOLD = 0.2

__C.DATA.TANDT.TRUCK = CN()
__C.DATA.TANDT.TRUCK.PATH = 'data/tandt/truck'
__C.DATA.TANDT.TRUCK.K_MAX = 7
__C.DATA.TANDT.TRUCK.K_MIN = 3
__C.DATA.TANDT.TRUCK.THRESHOLD = 0.2


__C.DATA.DB = CN()

__C.DATA.DB.DRJOHNSON = CN()
__C.DATA.DB.DRJOHNSON.PATH = 'data/db/drjohnson'
__C.DATA.DB.DRJOHNSON.K_MAX = 7
__C.DATA.DB.DRJOHNSON.K_MIN = 3
__C.DATA.DB.DRJOHNSON.THRESHOLD = 0.2

__C.DATA.DB.PLAYROOM = CN()
__C.DATA.DB.PLAYROOM.PATH = 'data/db/playroom'
__C.DATA.DB.PLAYROOM.K_MAX = 7
__C.DATA.DB.PLAYROOM.K_MIN = 3
__C.DATA.DB.PLAYROOM.THRESHOLD = 0.2


__C.DATA.NERFSYNTHETIC = CN()

__C.DATA.NERFSYNTHETIC.CHAIR = CN()
__C.DATA.NERFSYNTHETIC.CHAIR.PATH = 'data/nerfsynthetic/chair'
__C.DATA.NERFSYNTHETIC.CHAIR.K_MAX = 7
__C.DATA.NERFSYNTHETIC.CHAIR.K_MIN = 3
__C.DATA.NERFSYNTHETIC.CHAIR.THRESHOLD = 0.2

__C.DATA.NERFSYNTHETIC.DRUMS = CN()
__C.DATA.NERFSYNTHETIC.DRUMS.PATH = 'data/nerfsynthetic/drums'
__C.DATA.NERFSYNTHETIC.DRUMS.K_MAX = 7
__C.DATA.NERFSYNTHETIC.DRUMS.K_MIN = 3
__C.DATA.NERFSYNTHETIC.DRUMS.THRESHOLD = 0.2

__C.DATA.NERFSYNTHETIC.FICUS = CN()
__C.DATA.NERFSYNTHETIC.FICUS.PATH = 'data/nerfsynthetic/ficus'
__C.DATA.NERFSYNTHETIC.FICUS.K_MAX = 7
__C.DATA.NERFSYNTHETIC.FICUS.K_MIN = 3
__C.DATA.NERFSYNTHETIC.FICUS.THRESHOLD = 0.2

__C.DATA.NERFSYNTHETIC.HOTDOG = CN()
__C.DATA.NERFSYNTHETIC.HOTDOG.PATH = 'data/nerfsynthetic/hotdog'
__C.DATA.NERFSYNTHETIC.HOTDOG.K_MAX = 7
__C.DATA.NERFSYNTHETIC.HOTDOG.K_MIN = 3
__C.DATA.NERFSYNTHETIC.HOTDOG.THRESHOLD = 0.2

__C.DATA.NERFSYNTHETIC.LEGO = CN()
__C.DATA.NERFSYNTHETIC.LEGO.PATH = 'data/nerfsynthetic/lego'
__C.DATA.NERFSYNTHETIC.LEGO.K_MAX = 7
__C.DATA.NERFSYNTHETIC.LEGO.K_MIN = 3
__C.DATA.NERFSYNTHETIC.LEGO.THRESHOLD = 0.2

__C.DATA.NERFSYNTHETIC.MATERIALS = CN()
__C.DATA.NERFSYNTHETIC.MATERIALS.PATH = 'data/nerfsynthetic/materials'
__C.DATA.NERFSYNTHETIC.MATERIALS.K_MAX = 7
__C.DATA.NERFSYNTHETIC.MATERIALS.K_MIN = 3
__C.DATA.NERFSYNTHETIC.MATERIALS.THRESHOLD = 0.2

__C.DATA.NERFSYNTHETIC.MIC = CN()
__C.DATA.NERFSYNTHETIC.MIC.PATH = 'data/nerfsynthetic/mic'
__C.DATA.NERFSYNTHETIC.MIC.K_MAX = 7
__C.DATA.NERFSYNTHETIC.MIC.K_MIN = 3
__C.DATA.NERFSYNTHETIC.MIC.THRESHOLD = 0.2

__C.DATA.NERFSYNTHETIC.SHIP = CN()
__C.DATA.NERFSYNTHETIC.SHIP.PATH = 'data/nerfsynthetic/ship'
__C.DATA.NERFSYNTHETIC.SHIP.K_MAX = 7
__C.DATA.NERFSYNTHETIC.SHIP.K_MIN = 3
__C.DATA.NERFSYNTHETIC.SHIP.THRESHOLD = 0.2


__C.DATA.BUNGEENERF = CN()

__C.DATA.BUNGEENERF.AMSTERDAM = CN()
__C.DATA.BUNGEENERF.AMSTERDAM.PATH = 'data/bungeenerf/amsterdam'
__C.DATA.BUNGEENERF.AMSTERDAM.K_MAX = 7
__C.DATA.BUNGEENERF.AMSTERDAM.K_MIN = 3
__C.DATA.BUNGEENERF.AMSTERDAM.THRESHOLD = 0.2

__C.DATA.BUNGEENERF.BARCELONA = CN()
__C.DATA.BUNGEENERF.BARCELONA.PATH = 'data/bungeenerf/barcelona'
__C.DATA.BUNGEENERF.BARCELONA.K_MAX = 7
__C.DATA.BUNGEENERF.BARCELONA.K_MIN = 3
__C.DATA.BUNGEENERF.BARCELONA.THRESHOLD = 0.2

__C.DATA.BUNGEENERF.BILBAO = CN()
__C.DATA.BUNGEENERF.BILBAO.PATH = 'data/bungeenerf/bilbao'
__C.DATA.BUNGEENERF.BILBAO.K_MAX = 7
__C.DATA.BUNGEENERF.BILBAO.K_MIN = 3
__C.DATA.BUNGEENERF.BILBAO.THRESHOLD = 0.2

__C.DATA.BUNGEENERF.CHICAGO = CN()
__C.DATA.BUNGEENERF.CHICAGO.PATH = 'data/bungeenerf/chicago'
__C.DATA.BUNGEENERF.CHICAGO.K_MAX = 7
__C.DATA.BUNGEENERF.CHICAGO.K_MIN = 3
__C.DATA.BUNGEENERF.CHICAGO.THRESHOLD = 0.2

__C.DATA.BUNGEENERF.HOLLYWOOD = CN()
__C.DATA.BUNGEENERF.HOLLYWOOD.PATH = 'data/bungeenerf/hollywood'
__C.DATA.BUNGEENERF.HOLLYWOOD.K_MAX = 7
__C.DATA.BUNGEENERF.HOLLYWOOD.K_MIN = 3
__C.DATA.BUNGEENERF.HOLLYWOOD.THRESHOLD = 0.2

__C.DATA.BUNGEENERF.POMPIDOU = CN()
__C.DATA.BUNGEENERF.POMPIDOU.PATH = 'data/bungeenerf/pompidou'
__C.DATA.BUNGEENERF.POMPIDOU.K_MAX = 7
__C.DATA.BUNGEENERF.POMPIDOU.K_MIN = 3
__C.DATA.BUNGEENERF.POMPIDOU.THRESHOLD = 0.2

__C.DATA.BUNGEENERF.QUEBEC = CN()
__C.DATA.BUNGEENERF.QUEBEC.PATH = 'data/bungeenerf/quebec'
__C.DATA.BUNGEENERF.QUEBEC.K_MAX = 7
__C.DATA.BUNGEENERF.QUEBEC.K_MIN = 3
__C.DATA.BUNGEENERF.QUEBEC.THRESHOLD = 0.2

__C.DATA.BUNGEENERF.ROME = CN()
__C.DATA.BUNGEENERF.ROME.PATH = 'data/bungeenerf/rome'
__C.DATA.BUNGEENERF.ROME.K_MAX = 7
__C.DATA.BUNGEENERF.ROME.K_MIN = 3
__C.DATA.BUNGEENERF.ROME.THRESHOLD = 0.2


__C.DATA.IMW2020 = CN()

__C.DATA.IMW2020.BRANDENBURG_GATE = CN()
__C.DATA.IMW2020.BRANDENBURG_GATE.PATH = 'data/imw2020/brandenburg_gate'
__C.DATA.IMW2020.BRANDENBURG_GATE.K_MAX = 7
__C.DATA.IMW2020.BRANDENBURG_GATE.K_MIN = 3
__C.DATA.IMW2020.BRANDENBURG_GATE.THRESHOLD = 0.2

__C.DATA.IMW2020.BUCKINGHAM_PALACE = CN()
__C.DATA.IMW2020.BUCKINGHAM_PALACE.PATH = 'data/imw2020/buckingham_palace'
__C.DATA.IMW2020.BUCKINGHAM_PALACE.K_MAX = 7
__C.DATA.IMW2020.BUCKINGHAM_PALACE.K_MIN = 3
__C.DATA.IMW2020.BUCKINGHAM_PALACE.THRESHOLD = 0.2

__C.DATA.IMW2020.COLOSSEUM_EXTERIOR = CN()
__C.DATA.IMW2020.COLOSSEUM_EXTERIOR.PATH = 'data/imw2020/colosseum_exterior'
__C.DATA.IMW2020.COLOSSEUM_EXTERIOR.K_MAX = 7
__C.DATA.IMW2020.COLOSSEUM_EXTERIOR.K_MIN = 3
__C.DATA.IMW2020.COLOSSEUM_EXTERIOR.THRESHOLD = 0.2

__C.DATA.IMW2020.FLORENCE_CATHEDRAL_SIDE = CN()
__C.DATA.IMW2020.FLORENCE_CATHEDRAL_SIDE.PATH = 'data/imw2020/florence_cathedral_side'
__C.DATA.IMW2020.FLORENCE_CATHEDRAL_SIDE.K_MAX = 7
__C.DATA.IMW2020.FLORENCE_CATHEDRAL_SIDE.K_MIN = 3
__C.DATA.IMW2020.FLORENCE_CATHEDRAL_SIDE.THRESHOLD = 0.2
