from configs.config_base import config

config['gpu'] = '1'
config['name'] = 'cvae_pyr_4_4'
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!!!!!!!!!\
common condition encoders for all levels, only run once\
full res loss on each level\
all higher levels are added to the output of first level\
using custom MLP upscale enhancement\
layers 32 and 512\
try loading weights from previous trainings\
'
config['dataset_in_memory'] = False
config['resolution'] = None

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, 7),
                        }
config['signature_static'] = { 'uv_static': (512, 256, 3),
                              'mesh_mask': (14475,),
                              'uv_mask': (512, 256,),
                              }

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

config['pyramid']['resolutions'] = [32, 512]

for res in config['pyramid']['resolutions']:
  if res != 512:
      res_modifier = '_' + str(res)
      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)


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

config['avg_garment_path'] = ''


config['pyramid']['levels'] = []
for res in config['pyramid']['resolutions']:
  if res == 512:
    level = dict()
    level['resolution'] = 512
    level['latent_dim'] = config['latent_dim']
    level['kernel_size'] = 7
    level['encoder_depths'] = [3, 3, 9, 3]
    level['encoder_dims'] = [32, 64, 128, 256]
    level['decoder_input'] = (16, 8, 80)
    level['decoder_depths'] = [3, 3, 9, 3]
    level['decoder_dims'] = [256, 128, 64, 16]
    level['avg_garment_path'] = ''
    level['upscale_enhancer'] = ''
    config['pyramid']['levels'].append(level)
  elif res == 128:
    level = dict()
    level['resolution'] = 128
    level['latent_dim'] = config['latent_dim']
    level['kernel_size'] = 5
    level['encoder_depths'] = [16, 16, 9, 9]
    level['encoder_dims'] = [64, 128, 256, 256]
    level['decoder_input'] = (4, 2, 128)
    level['decoder_depths'] = [16, 16, 9, 9]
    level['decoder_dims'] = [256, 256, 128, 64]
    level['avg_garment_path'] = ''
    level['upscale_enhancer'] = 'checkpoints/up_2_3_d128'
    config['pyramid']['levels'].append(level)
  elif res == 32:
    level = dict()
    level['resolution'] = 32
    level['latent_dim'] = config['latent_dim']
    level['kernel_size'] = 3
    level['encoder_depths'] = [27, 27]
    level['encoder_dims'] = [256, 256]
    level['decoder_input'] = (4, 2, 256)
    level['decoder_depths'] = [27, 27]
    level['decoder_dims'] = [256, 256]
    level['avg_garment_path'] = ''
    level['upscale_enhancer'] = 'checkpoints/up_2_3'
    config['pyramid']['levels'].append(level)
  else:
    raise NotImplementedError