import numpy as np

from tf_watermark.tf_watermark_regularizers import WatermarkRegularizer

RESULT_PATH = './result'
fname_x = RESULT_PATH + '/x.npy'  # 投影矩阵
fname_b = RESULT_PATH + '/b.npy'  # 密钥


def get_wmark_regularizers(model):
    ret = []

    for i, layer in enumerate(model.layers):
        if hasattr(layer, "kernel_regularizer"):
            regularizer = layer.kernel_regularizer
            if isinstance(regularizer, WatermarkRegularizer):
                ret.append((i, regularizer))
    return ret


# 保存水印签名的函数
def save_wmark_signatures(model):
    for layer_id, wmark_regularizer in get_wmark_regularizers(model):
        np.save(fname_x, wmark_regularizer.get_matrix())
        np.save(fname_b, wmark_regularizer.get_signature())


# 从文件中获取
def get_layer_weights_and_predicted(model, checkpoint_save_path, target_blk_id):
    x = np.load(fname_x)
    model.load_weights(checkpoint_save_path)
    # get signature from model weight and matrix
    target_layer = model.get_layer(index=target_blk_id)
    layer_weights = target_layer.get_weights()
    weight = (np.array(layer_weights[0])).mean(axis=3)
    pred_bparam = np.dot(weight.reshape(1, weight.size), x)  # dot product
    pred_bparam = np.where(pred_bparam >= 0, 1, 0)
    return layer_weights[0], pred_bparam