from os.path import join
from app.utils.file_utils import write_json

def save_params_latent(fp_out_dir, tag):
  data = {
    "tag": tag,
    "decay_n": 2, 
    "features": True, 
    "clip": 1.0, 
    "stochastic_clipping": False, 
    "clipping": False, 
    "dataset": "inverses/{}/dataset.hdf5".format(tag),
    "inv_layer": "latent", 
    "decay_lr": True, 
    "inv_it": 15000, 
    "generator_path": "https://tfhub.dev/deepmind/biggan-512/2", 
    "attention_map_layer": "Generator_2/attention/Softmax:0",
    "pre_trained_latent": False, 
    "lambda_dist": 0.0, 
    "likeli_loss": True, 
    "init_hi": 0.001, 
    "lr": 0.1, 
    "norm_loss": False, 
    "generator_fixed_inputs": {
      "truncation": 1.0
    }, 
    "log_z_norm": True, 
    "feature_extractor_path": "https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1", 
    "mse": True, 
    "custom_grad_relu": False, 
    "random_label": False, 
    "lambda_feat": 1.0, 
    "init_gen_dist": False, 
    "log_activation_layer": "Generator_2/GBlock/Relu:0", 
    "batch_size": 4, 
    "fixed_z": False, 
    "feature_extractor_output": "InceptionV3/Mixed_7a", 
    "init_lo": -0.001, 
    "lambda_mse": 1.0, 
    "lambda_reg": 0.1, 
    "dist_loss": False, 
    "sample_size": 4, 
    "dataset": "inverses/{}/dataset.encodings.hdf5".format(tag),
    "save_progress": True,
  }
  fp_out_fn = join(fp_out_dir, "params_latent.json")
  write_json(data, fp_out_fn)

def save_params_dense(fp_out_dir, tag):
  data = {
    "tag": tag,
    "decay_n": 2, 
    "features": True, 
    "clip": 1.0, 
    "stochastic_clipping": False, 
    "clipping": False, 
    "dataset": "inverses/{}/dataset.encodings.hdf5".format(tag),
    "inv_layer": "Generator_2/G_Z/Reshape:0", 
    "decay_lr": True, 
    "inv_it": 15000, 
    "generator_path": "https://tfhub.dev/deepmind/biggan-512/2", 
    "attention_map_layer": "Generator_2/attention/Softmax:0",
    "pre_trained_latent": True, 
    "lambda_dist": 10.0, 
    "likeli_loss": False, 
    "init_hi": 0.001, 
    "lr": 0.01, 
    "norm_loss": False, 
    "generator_fixed_inputs": {
      "truncation": 1.0
    }, 
    "log_z_norm": False, 
    "feature_extractor_path": "https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1", 
    "mse": True, 
    "custom_grad_relu": False, 
    "random_label": False, 
    "lambda_feat": 1.0, 
    "init_gen_dist": False, 
    "log_activation_layer": "Generator_2/GBlock/Relu:0", 
    "batch_size": 4, 
    "fixed_z": True, 
    "feature_extractor_output": "InceptionV3/Mixed_7a", 
    "init_lo": -0.001, 
    "lambda_mse": 1.0, 
    "lambda_reg": 0.1, 
    "dist_loss": True, 
    "sample_size": 4, 
    "dataset": "inverses/{}/dataset.encodings.dense.hdf5".format(tag),
    "save_progress": True,
  }
  fp_out_fn = join(fp_out_dir, "params_dense.json")
  write_json(data, fp_out_fn)