from utils.convnext_utils import CNBlockConfig

config = dict()

config['gpu'] = '4'
config['name'] = 'abl_iccvw_all_4'
config['notes'] = 'fixed collision loss, train from zero'
config['tags'] = ['iccv', 'abl_all', 'torch']
config['script'] = 'scripts/cvae.py'

config['pretrained_model'] = 'checkpoints/iccvw/abl_iccvw_all_3'

config['use_skinning'] =              True # and deshaping
config['use_normals'] =               True
config['use_extra_losses'] =          True # edge, collision, 3d mesh
config['use_norm_loss'] =             True

ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False)
assert ablation_count == 1 or ablation_count == 0

config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
if config['use_skinning']:
  assert (config['avg_garment_path'] != '')

config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision

config['batch_size'] = 4
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['device'] = 'cuda'

config['train_dir'] = '/data-local/127/cloth3d_processed'
config['valid_dir'] = '/data-local/127/cloth3d_processed'
config['test_dir'] = '/data-local/127_2/cloth3d_processed_test'

if config['is_skirt']:
  config['train_frames'] = 'frames_train_skirt.txt'
  config['valid_frames'] = 'frames_valid_skirt.txt'
  config['test_frames'] = 'frames_test_skirt.txt'
else:
  config['train_frames'] = 'frames_train.txt'
  config['valid_frames'] = 'frames_valid.txt'
  config['test_frames'] = 'frames_test.txt'

config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'

config['smpl_res'] = 14475
if config['is_cape']:
  config['smpl_res'] = 3638

config['latent_dim'] = 128

###############    
config['lr'] = 0.0001#                 0.0001
###############
config['starting_balancing_coeff'] = 0.0001
config['balancing_activation_limit'] = 0.01

config['signature']: set = { 'id',
                            'frame_nr',
                            'pose_shape',
                            'uv_body_posed',
                            'mesh_body_posed',
                            'mesh_posed',
                            'uv_normals',
                            'uv_static',
                            'mesh_mask',
                            'uv_mask',
                        }

if config['use_skinning']:
  config['signature'].add('uv_unposed')
  config['signature'].add('mesh_unposed')
else:
  config['signature'].add('uv_posed')

config['upscale_enhancer'] = ''
config['freeze_encoders'] = True
config['use_one_hot'] = False
config['bottleneck_scale'] = 2

config['cvae'] = dict()
config['cvae']['device'] = config['device']
config['cvae']['is_label'] = False
config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
config['cvae']['kernel_size'] = 7
config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']

block_setting = [
    CNBlockConfig(32, 64, 3),
    CNBlockConfig(64, 128, 3),
    CNBlockConfig(128, 256, 3),
    CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
    CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
]

block_setting_reverse = [
    CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
    CNBlockConfig(256, 128, 9),
    CNBlockConfig(128, 64, 9),
    CNBlockConfig(64, 32, 3),
    CNBlockConfig(32, 3, 3),
]

config['cvae']['encoder_block_settings'] = block_setting
config['cvae']['decoder_block_settings'] = block_setting_reverse



####################### cond_ae_norm ############################
config['cond_ae_norm'] = dict()
config['cond_ae_norm']['device'] = config['device']
config['cond_ae_norm']['latent_dim'] = config['latent_dim']
config['cond_ae_norm']['kernel_size'] = 7
config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']

block_setting = [
    CNBlockConfig(16, 32, 3),
    CNBlockConfig(32, 64, 3),
    CNBlockConfig(64, 128, 9),
    CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
]

block_setting_reverse = [
    CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
    CNBlockConfig(128, 64, 9),
    CNBlockConfig(64, 32, 9),
    CNBlockConfig(32, 16, 3),
    CNBlockConfig(16, 3, 3),
]

config['cond_ae_norm']['encoder_block_settings'] = block_setting
config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_norm']['activation'] = 'tanh'

####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']

block_setting = [
    CNBlockConfig(16, 32, 3),
    CNBlockConfig(32, 64, 3),
    CNBlockConfig(64, 128, 9),
    CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9),
    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]

block_setting_reverse = [
    CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
    CNBlockConfig(128, 64, 9),
    CNBlockConfig(64, 32, 9),
    CNBlockConfig(32, 16, 3),
    CNBlockConfig(16, 3, 3),
]

config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'

####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']

block_setting = [
    CNBlockConfig(8, 16, 3),
    CNBlockConfig(16, 32, 3),
    CNBlockConfig(32, 64, 3),
    CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]

block_setting_reverse = [
    CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
    CNBlockConfig(64, 32, 9),
    CNBlockConfig(32, 16, 3),
    CNBlockConfig(16, 8, 3),
    CNBlockConfig(8, 3, 3),
]

config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'