from tensorflow.python.keras.losses import binary_crossentropy
from tensorflow.python.keras.regularizers import Regularizer
import numpy as np
import tensorflow as tf


class WatermarkRegularizer(Regularizer):

    def __init__(self, scale, b, randseed='none'):
        self.scale = tf.constant(scale, dtype=tf.float32)
        self.b = b  # 密钥
        self.x = None  # 投影矩阵
        self.v_x = None  # 投影矩阵设为可训练参数
        self.randseed = randseed

    def __call__(self, weight):
        if self.x is None:
            x_rows = np.prod(weight.shape[0:3])
            x_cols = self.b.shape[1]
            if self.randseed != 'none':
                np.random.seed(int(self.randseed))
            self.x = np.random.randn(x_rows, x_cols)
            self.v_x = tf.Variable(self.x, trainable=False, dtype=tf.float32)
        weight_mean = tf.reduce_mean(weight, axis=3)
        w = tf.reshape(weight_mean, (1, tf.reduce_prod(weight_mean.shape)))
        regularized_loss = self.scale * tf.reduce_sum(
            binary_crossentropy(tf.sigmoid(tf.matmul(w, self.v_x)), tf.cast(self.b, tf.float32)))
        return regularized_loss

    # 获取投影矩阵
    def get_matrix(self):
        return self.x

    # 获取签名值
    def get_signature(self):
        return self.b