import tensorflow as tf
from keras import layers


input_config = {
    'category': [
        # {'feature': 'hour', 'dtype': 'int32', 'num_tokens': 24,'vocab': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]},
        {'feature': 'banner_pos', 'dtype': 'int32', 'num_tokens': 8, 'vocab': [0, 1, 2, 3, 4, 5, 6, 7]},
        {'feature': 'device_type', 'dtype': 'int32', 'num_tokens': 6, 'vocab': [0, 1, 2, 3, 4, 5]},
        {'feature': 'device_conn_type', 'dtype': 'int32', 'num_tokens': 6, 'vocab': [0, 1, 2, 3, 4, 5]},
        {'feature': 'C18', 'dtype': 'int32', 'num_tokens': 4, 'vocab': [0, 1, 2, 3]},
    ],
    # hash分桶
    'hash': [
        {'feature': 'site_category', 'num_bins': 1000, 'dtype': 'string'},
        {'feature': 'app_category', 'num_bins': 1000, 'dtype': 'string'},
        {'feature': 'C14', 'num_bins': 1000, 'dtype': 'int32'},
        {'feature': 'C15', 'num_bins': 1000, 'dtype': 'int32'},
        {'feature': 'C16', 'num_bins': 1000, 'dtype': 'int32'},
        {'feature': 'C17', 'num_bins': 1000, 'dtype': 'int32'},
        {'feature': 'C21', 'num_bins': 1000, 'dtype': 'int32'},
    ],
    # 数值分桶
    'int_bucket': [
        # {'feature': 'Age', 'bin_boundaries': [10, 20, 30, 40, 50, 60, 70, 80, 90], 'embedding_dims': 10}
    ],
    # 数值类型（归一化）
    'num': [

    ],
    # 手动交叉
    'cross': [

    ],
    # 原始稠密特征
    # 'dense': [
    #     {'feature': 'site_category', 'dtype': 'float32'}
    # ]
}

voc_size = {
    # 'hour':24,
    'banner_pos': 8,
    'device_type': 6,
    'device_conn_type': 6,
    'C18': 4,
    'site_category': 1000,
    'app_category': 1000,
    'C14': 1000,
    'C15': 1000,
    'C16': 1000,
    'C17': 1000,
    'C21': 1000,

}
spare_features_config = [
    # 'hour',
    'banner_pos', 'device_type', 'device_conn_type', 'C18', 'site_category', 'app_category', 'C14', 'C15', 'C16', 'C17',
    'C21']
dense_features_config = []


class CrossNet(tf.keras.layers.Layer):
    def __init__(self, layer_nums=3):
        super(CrossNet, self).__init__()
        self.layer_nums = layer_nums

    def build(self, input_shape):
        # 计算w的维度，w的维度与输入数据的最后一个维度相同
        self.dim = int(input_shape[-1])

        # 注意，在DCN中W不是一个矩阵而是一个向量，这里根据残差的层数定义一个权重列表
        self.W = [self.add_weight(name='W_' + str(i), shape=(self.dim,)) for i in range(self.layer_nums)]
        self.b = [self.add_weight(name='b_' + str(i), shape=(self.dim,), initializer='zeros') for i in
                  range(self.layer_nums)]

    def call(self, inputs):
        # 进行特征交叉时的x_0一直没有变，变的是x_l和每一层的权重
        x_0 = inputs  # B x dims
        x_l = x_0
        for i in range(self.layer_nums):
            # 将x_l的第一个维度与w[i]的第0个维度计算点积
            xl_w = tf.tensordot(x_l, self.W[i], axes=(1, 0))  # B,
            xl_w = tf.expand_dims(xl_w, axis=-1)  # 在最后一个维度上添加一个维度 # B x 1
            cross = tf.multiply(x_0, xl_w)  # B x dims
            x_l = cross + self.b[i] + x_l

        return x_l


def build_input(input_config):
    feature_input = []
    feature_map = {}
    input_map = {}
    # 构建连续数值型特征输入
    for num_feature in input_config.get('num', []):
        layer = tf.keras.Input(shape=[1], dtype=num_feature['dtype'], name=num_feature[
            'feature'])
        input_map[num_feature['feature']] = layer
        feature_input.append(layer)  # tf.feature_column.numeric_column(num_feature['feature']))
        feature_map[num_feature['feature']] = layer
    # 构建分类特征输入
    for cate_feature in input_config.get('category', []):
        layer = layers.Input(shape=[1], dtype=cate_feature['dtype'], name=cate_feature['feature'])
        input_map[cate_feature['feature']] = layer
        # 是否数字型
        if cate_feature.get('num_tokens') is None:
            layer = layers.StringLookup(vocabulary=cate_feature['vocabulary'], output_mode="one_hot",
                                                 num_oov_indices=0)(layer)
            input_dim = len(cate_feature['vocabulary'])
        else:
            layer = layers.CategoryEncoding(num_tokens=cate_feature['num_tokens'], output_mode="one_hot")(
                layer)
            input_dim = cate_feature['num_tokens']
        # 是否需要embedding
        if cate_feature.get('embedding_dims') is not None:
            layer = layers.Dense(cate_feature['embedding_dims'], use_bias=False)(layer)
        feature_input.append(layer)
        feature_map[cate_feature['feature']] = layer
    # 需要hash分桶的特征
    for hash_feature in input_config.get('hash', []):
        layer = tf.keras.Input(shape=[1], dtype=hash_feature['dtype'], name=hash_feature['feature'])
        input_map[hash_feature['feature']] = layer
        layer = layers.Hashing(num_bins=hash_feature['num_bins'], output_mode='one_hot',
                                        )(layer)
        if hash_feature.get('embedding_dims') is not None:
            layer = layers.Dense(hash_feature['embedding_dims'], use_bias=False)(layer)
        feature_input.append(layer)
        feature_map[hash_feature['feature']] = layer
    # 连续数值分桶
    for bucket_feature in input_config.get('int_bucket', []):
        layer = layers.Discretization(bin_boundaries=bucket_feature['bin_boundaries'],
                                               name=bucket_feature['feature'])
        if bucket_feature.get('embedding_dims') is not None:
            embedding = layers.Dense(bucket_feature['embedding_dims'], use_bias=False)
            layer = embedding(layer)
        feature_input.append(layer)
        feature_map[bucket_feature['feature']] = layer
        input_map[hash_feature['feature']] = layer
    cross_cate_map = {}
    # 构建交叉特征
    # for cross_feature in input_config.get('cross', []):
    #     col = []
    #     col = col + build_input(cross_feature['features'])
    #     # layer = layers.experimental.preprocessing.HashedCrossing(num_bins=cross_feature['num_bins'],
    #     #                                                                   output_mode='one_hot', sparse=True)(
    #     #     (tuple(col)))
    #     layer=tf.feature_column.indicator_column(tf.feature_column.crossed_column(col, 10000))
    #     feature_input.append(layer)
    #     feature_input_map[cross_feature['feature']] = layer

    return feature_input, feature_map, input_map


def build_embed_features(embedding_dims, spare_features_config, feature_input_map):
    embed_features = []
    for feature_name in spare_features_config:
        embedding = layers.Dense(embedding_dims, use_bias=False)
        embed_features.append(embedding(feature_input_map[feature_name]))
    return embed_features


def build_spare_features(spare_features_config, feature_input_map):
    spare_features = []
    for feature_name in spare_features_config:
        spare_features.append(feature_input_map[feature_name])
    return spare_features


def build_dense_features(dense_features_config, feature_input_map):
    dense_features = []
    for feature_name in spare_features_config:
        dense_features.append(feature_input_map[feature_name])
    return dense_features

def dcn(input_config, spare_features_config, dense_features_config):
    feature_input, feature_map, input_map = build_input(input_config)
    embed_features = build_embed_features(8, spare_features_config, feature_map)
    #spare_features = build_spare_features(spare_features_config, feature_map)
    dense_features = build_dense_features(dense_features_config, feature_map)
    #dnn与cross共享输入
    dnn_input = layers.concatenate(dense_features + embed_features)
    cross_input=dnn_input
    hidden_units = [32, 64, 64]
    dropout_rate = 0.1
    x=dnn_input
    for units in hidden_units:
        x = layers.Dense(units)(x)
        x = layers.BatchNormalization()(x)
        x = layers.ReLU()(x)
        x = layers.Dropout(dropout_rate)(x)
    dnn_output = CrossNet()(dnn_input)
    stack= layers.Concatenate(axis=1)([dnn_output, x])
    output = layers.Dense(1,activation='sigmoid')(stack)
    model = tf.keras.Model(input_map, output)
    model.compile(optimizer="adam",
                  loss="binary_crossentropy",
                  metrics=tf.keras.metrics.BinaryAccuracy()
                  )
    return model




model = dcn(input_config, spare_features_config, dense_features_config)
dataset = tf.data.experimental.make_csv_dataset(
    '/Volumes/Data/oysterqaq/Desktop/Avazu_train_1.csv', batch_size=2, label_name='click'
)
model.summary()
model.fit(dataset,
          batch_size=20, epochs=11)