import tensorflow as tf


def save_model_params(sess,
                      file_path):
    # assert file_path.endswith('.npz')
    param_dict = {v.name: v.eval(sess) for v in tf.global_variables()}
    np.savez_compressed(file_path, **param_dict)


def load_model_params(net,
                      param_dict,
                      sess,
                      ignore_missing=False):
    for param_name, param_data in param_dict:
        with tf.variable_scope(param_name, reuse=True):
            try:
                var = tf.get_variable(param_name)
                sess.run(var.assign(param_data))
            except ValueError:
                if not ignore_missing:
                    raise


def prepare_model(model_name,
                  use_pretrained,
                  pretrained_model_file_path):
    kwargs = {'pretrained': use_pretrained}

    net = get_model(model_name, **kwargs)
    input_image_size = net.in_size[0] if hasattr(net, 'in_size') else 224

