import math

MAP3DBN = {
    0: {'batch_size': 32, 'batch_split': 1, 'gen_lr': 0.0001, 'disc_lr': 0.0004},
    int(140e3 + 1): {'batch_size': 32, 'batch_split': 1, 'gen_lr': 5e-5, 'disc_lr': 2e-4},
    int(300e3 + 1): {},

    'name': 'map3dbn',
    'trainer': 'PhaseTrainer',
    'phases': [
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": True},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": True},
    ],
    '2d_coords_input': True,
    '2d_semantic_input': False,
    '2d_latent_input': False,
    'neural_field_latent_input': False,
    'use_mixed_precision': True,
    'lock_view_dependence': True,
    'num_steps': 32,
    'render_width': 32,
    'render_height': 64,
    'gen_width': 128,
    'gen_height': 256,
    'ray_start': -0.5,
    'ray_end': 0.55,
    'side_length': 2.85,
    'depth_length': 1.05,
    'vis_rotate': math.pi / 6,
    'fade_steps': 1,
    'sample_dist': 'gaussian',
    'h_stddev': 0.4,
    'v_stddev': 0.1,
    'h_mean': 0,
    'v_mean': 0,
    'coordinate_mode': 'fix_body',
    'betas': (0, 0.9),
    'unique_lr': True,
    'appearance_codes_lr_mul': 1.,
    'mapping_net_lr_mul': 0.05,
    'neural_field_lr_mul': 0.05,
    'weight_decay': 0,
    'gan_lambda': 0,
    'r1_lambda': 0.25,
    'photometric_lambda': 0,
    'perceptual_lambda': [0, 0, 0, 0],
    'latent_lambda': 0,
    'z_lambda': 0,
    'pos_lambda': 0,
    'semantic_lambda': 0,
    'segmentation_lambda': 1,
    'input_dim': 3,
    'output_dim': 3,
    'latent_dim': 384,
    'hidden_dim': 384,
    'feature_dim': 384,
    'semantic_dim': 0,
    'geo_feature_dim': 31,
    'label_dim': 26,
    'grad_clip': 1.,
    'neural_field_cls': 'COORDCONCATSIREN',
    'generator': 'Map3DGenerator',
    'map3d_mode': 'mixed',
    'neural_field_blocks': 4,
    'synthesis_blocks': 9,
    'mod_blocks': list(range(3)),
    'spatial_normalization': 'batch_norm',
    'discriminator': 'UNetDiscriminator',
    'condition_modal_disc_real': 'body_segments',
    'condition_modal_disc_gen': 'rasterized_segments',
    'condition_modal_gen': 'rasterized_segments',
    'ada_aug': dict(xflip=1, rotate90=0, rotate_max=0.05, xint=0,
                    scale=1, rotate=1, aniso=1, xfrac=0, brightness=1, contrast=1, saturation=1),
    'ada_target': 0.6,
    'ada_interval': 0,
    'ada_kimg': 20,
    'ada_alpha_thresh': 0.5,
    'dataset': 'SHHQDataset',
    'dataset_length': 10,
    'dataroot': './datasets/shhq_example_dataset',
    'joints': list(range(24)),
    'white_back': True,
    'clamp_mode': 'relu',
    'z_dist': 'gaussian',
    'hierarchical_sample': False,
    'learnable_dist': False,
    'last_back': False,
    'eval_last_back': True,
}


MAP3DBN512 = {
    0: {'batch_size': 32, 'batch_split': 1, 'gen_lr': 5e-5, 'disc_lr': 2e-4},
    int(300e3 + 1): {},

    'name': 'map3dbn512',
    'trainer': 'PhaseTrainer',
    'phases': [
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": True},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": True},
    ],
    '2d_coords_input': True,
    '2d_semantic_input': False,
    '2d_latent_input': False,
    'neural_field_latent_input': False,
    'use_mixed_precision': True,
    'lock_view_dependence': True,
    'num_steps': 32,
    'render_width': 48,
    'render_height': 96,
    'gen_width': 256,
    'gen_height': 512,
    'ray_start': -0.5,
    'ray_end': 0.55,
    'side_length': 2.85,
    'depth_length': 1.05,
    'vis_rotate': math.pi / 6,
    'fade_steps': 1,
    'sample_dist': 'gaussian',
    'h_stddev': 0.4,
    'v_stddev': 0.1,
    'h_mean': 0,
    'v_mean': 0,
    'coordinate_mode': 'fix_body',
    # 'topk_interval': 2000,
    # 'topk_v': 0,
    'betas': (0, 0.9),
    'unique_lr': True,
    'appearance_codes_lr_mul': 1.,
    'mapping_net_lr_mul': 0.05,
    'neural_field_lr_mul': 0.05,
    'weight_decay': 0,
    'gan_lambda': 0,
    'r1_lambda': 0,
    'photometric_lambda': 0,
    'perceptual_lambda': [0, 0, 0, 0],
    'latent_lambda': 0,
    'z_lambda': 0,
    'pos_lambda': 0,
    'semantic_lambda': 0,
    'segmentation_lambda': 1,
    'input_dim': 3,
    'output_dim': 3,
    'latent_dim': 256,
    'hidden_dim': 256,
    'feature_dim': 256,
    'semantic_dim': 0,
    'geo_feature_dim': 31,
    'label_dim': 26,
    'grad_clip': 1.,
    'neural_field_cls': 'COORDCONCATSIREN',
    'generator': 'Map3DGenerator',
    'map3d_mode': 'mixed',
    'neural_field_blocks': 4,
    'synthesis_blocks': 9,
    'mod_blocks': list(range(3)),
    'spatial_normalization': 'batch_norm',
    'discriminator': 'UNetDiscriminator',
    'condition_modal_disc_real': 'body_segments',
    'condition_modal_disc_gen': 'rasterized_segments',
    'condition_modal_gen': 'rasterized_segments',
    'ada_aug': dict(xflip=1, rotate90=0, rotate_max=0.05, xint=0,
                    scale=1, rotate=1, aniso=1, xfrac=0, brightness=1, contrast=1, saturation=1),
    'ada_target': 0.6,
    'ada_interval': 0,
    'ada_kimg': 20,
    'ada_alpha_thresh': 0.5,
    'dataset': 'SHHQDataset',
    'dataset_length': 10,
    'dataroot': './datasets/shhq_example_dataset',
    'joints': list(range(24)),
    'white_back': True,
    'clamp_mode': 'relu',
    'z_dist': 'gaussian',
    'hierarchical_sample': False,
    'learnable_dist': False,
    'last_back': False,
    'eval_last_back': True,
}


MAP3DBN512L = {
    0: {'batch_size': 32, 'batch_split': 1, 'gen_lr': 5e-5, 'disc_lr': 2e-4},
    int(300e3 + 1): {},

    'name': 'map3dbn512l',
    'legacy_mode': True,
    'trainer': 'PhaseTrainer',
    'phases': [
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": True},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": False, "gen_modal": "rgbs", "do_r1": False},
        {"name": "uncond", "uncond": True, "rotate": True, "gen_modal": "rgbs", "do_r1": True},
    ],
    '2d_coords_input': True,
    '2d_semantic_input': False,
    '2d_latent_input': False,
    'neural_field_latent_input': False,
    'use_mixed_precision': True,
    'lock_view_dependence': True,
    'num_steps': 32,
    'render_width': 48,
    'render_height': 96,
    'gen_width': 256,
    'gen_height': 512,
    'ray_start': -0.5,
    'ray_end': 0.55,
    'side_length': 2.85,
    'depth_length': 1.05,
    'vis_rotate': math.pi / 6,
    'fade_steps': 1,
    'sample_dist': 'gaussian',
    'h_stddev': 0.4,
    'v_stddev': 0.1,
    'h_mean': 0,
    'v_mean': 0,
    'coordinate_mode': 'fix_body',
    # 'topk_interval': 2000,
    # 'topk_v': 0,
    'betas': (0, 0.9),
    'unique_lr': True,
    'appearance_codes_lr_mul': 1.,
    'mapping_net_lr_mul': 0.05,
    'neural_field_lr_mul': 0.05,
    'weight_decay': 0,
    'gan_lambda': 0,
    'r1_lambda': 0,
    'photometric_lambda': 0,
    'perceptual_lambda': [0, 0, 0, 0],
    'latent_lambda': 0,
    'z_lambda': 0,
    'pos_lambda': 0,
    'semantic_lambda': 0,
    'segmentation_lambda': 1,
    'input_dim': 3,
    'output_dim': 3,
    'latent_dim': 420,
    'hidden_dim': 420,
    'feature_dim': 420,
    'semantic_dim': 0,
    'geo_feature_dim': 31,
    'label_dim': 26,
    'grad_clip': 1.,
    'neural_field_cls': 'COORDCONCATSIREN',
    'generator': 'Map3DGenerator',
    'map3d_mode': 'isolated',
    'neural_field_blocks': 4,
    'synthesis_blocks': 9,
    'mod_blocks': list(range(3)),
    'spatial_normalization': 'batch_norm',
    'discriminator': 'UNetDiscriminator',
    'condition_modal_disc_real': 'body_segments',
    'condition_modal_disc_gen': 'rasterized_segments',
    'condition_modal_gen': 'rasterized_segments',
    'ada_aug': dict(xflip=1, rotate90=0, rotate_max=0.05, xint=0,
                    scale=1, rotate=1, aniso=1, xfrac=0, brightness=1, contrast=1, saturation=1),
    'ada_target': 0.6,
    'ada_interval': 0,
    'ada_kimg': 20,
    'ada_alpha_thresh': 0.5,
    'dataset': 'SHHQDataset',
    'dataset_length': 219047,
    'dataroot': './datasets/shhq_train_40000',
    'joints': list(range(24)),
    'white_back': True,
    'clamp_mode': 'relu',
    'z_dist': 'gaussian',
    'hierarchical_sample': False,
    'learnable_dist': False,
    'last_back': False,
    'eval_last_back': True,
}