from configs.config_base import config

config['gpu'] = '0'
config['name'] = 'cvae_19_2'
config['notes'] = '\
offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
retrained body encoder with deoriented bodies \
use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
fixed gender issue and retrained body encoder, FINAL\
remove shape from unposed garment\
training on 100 garments with offsets over the average of these\
fixed autoencoders!!!!!!!!!\
cvae only loads encoders instead AEs\
with edge and collision loss\
'

# config['resolution'] = 32


if config['resolution'] != 512:
  config['notes'] += '\
using custom enhancer for upscaling, using not filled versions\
'
    
config['lr'] = 0.0001
config['starting_balancing_coeff'] = 0.0001

config['signature_dynamic'] = { 'id': (),
                          'pose_shape': (72 + 11,),
                          'uv_body_posed': (512, 256, 3),
                          'mesh_body_posed': (14475, 3,),
                          'mesh_posed': (14475, 3,),
                          'mesh_unposed': (14475, 3,),
                          'uv_unposed': (512, 256, 3),
                          'uv_curvature': (512, 256),
                        }
config['signature_static'] = { 'uv_static': (512, 256, 3),
                              'mesh_mask': (14475,),
                              'uv_mask': (512, 256,),
                              }

if config['resolution'] != 512:
  res_modifier = '_' + str(config['resolution'])
  config['signature_dynamic']['uv_body_posed' + res_modifier] = (config['resolution'], config['resolution']//2, 3)
  config['signature_dynamic']['uv_unposed' + res_modifier] = (config['resolution'], config['resolution']//2, 3)
  config['signature_dynamic']['uv_curvature' + res_modifier] = (config['resolution'], config['resolution']//2)
  config['signature_static']['uv_static' + res_modifier] = (config['resolution'], config['resolution']//2, 3)
  config['signature_static']['uv_mask' + res_modifier] = (config['resolution'], config['resolution']//2)

if config['resolution'] == 512:
  config['avg_garment_path'] = 'custom_data/train_100_no_shape_avg_uv.npy'
else:
  config['avg_garment_path'] = ''

config['static_model'] = 'checkpoints/ae_static_7_2_x'
config['body_model'] = 'checkpoints/ae_body_9_x'
config['curvature_model'] = 'checkpoints/ae_curv_9_x'
config['freeze_encoders'] = True

config['cvae'] = dict()
config['cvae']['is_label'] = False
config['cvae']['latent_dim'] = config['latent_dim']

if config['resolution'] == 512:
  config['cvae']['kernel_size'] = 7
  config['cvae']['encoder_depths'] = [3, 3, 9, 3]
  config['cvae']['encoder_dims'] = [32, 64, 128, 256]
  config['cvae']['decoder_input'] = (16, 8, 80)
  config['cvae']['decoder_depths'] = [3, 3, 9, 3]
  config['cvae']['decoder_dims'] = [256, 128, 64, 16]
elif config['resolution'] == 128:
  config['cvae']['kernel_size'] = 5
  config['cvae']['encoder_depths'] = [16, 16, 9, 9]
  config['cvae']['encoder_dims'] = [64, 128, 256, 256]
  config['cvae']['decoder_input'] = (4, 2, 128)
  config['cvae']['decoder_depths'] = [16, 16, 9, 9]
  config['cvae']['decoder_dims'] = [256, 256, 128, 64]
elif config['resolution'] == 32:
  config['cvae']['kernel_size'] = 3
  config['cvae']['encoder_depths'] = [27, 27]
  config['cvae']['encoder_dims'] = [256, 256]
  config['cvae']['decoder_input'] = (4, 2, 256)
  config['cvae']['decoder_depths'] = [27, 27]
  config['cvae']['decoder_dims'] = [256, 256]
else:
  raise NotImplementedError