import tensornet as tn
import tensorflow as tf
def columns_builder():
    columns = {}
    for slot in set(C.WIDE_SLOTS + C.DEEP_SLOTS):
        columns[slot] = tn.feature_column.category_column(key=slot)
 
    wide_columns = []
    for slot in C.WIDE_SLOTS:
        feature_column = tf.feature_column.embedding_column(columns[slot], dimension=1)
        wide_columns.append(feature_column)
 
    deep_columns = []
    for slot in C.DEEP_SLOTS:
        feature_column = tf.feature_column.embedding_column(columns[slot], dimension=8)
        deep_columns.append(feature_column)
 
    return wide_columns, deep_columns
 
def create_model(wide_columns, deep_columns):
    wide, deep = None, None
 
    inputs = {}
    for slot in set(C.WIDE_SLOTS + C.DEEP_SLOTS):
        inputs[slot] = tf.keras.layers.Input(name=slot, shape=(None,), dtype="int64", sparse=True)
 
    sparse_opt = tn.core.AdaGrad(learning_rate=0.01, initial_g2sum=0.1, initial_scale=0.1)
    if wide_columns:
        wide = tn.layers.EmbeddingFeatures(wide_columns, sparse_opt, name='wide_inputs', is_concat=True)(inputs)
 
    if deep_columns:
        deep = tn.layers.EmbeddingFeatures(deep_columns, sparse_opt, name='deep_inputs', is_concat=True)(inputs)
        for i, unit in enumerate(C.DEEP_HIDDEN_UNITS):
            deep = tf.keras.layers.Dense(unit, activation='relu', name='dnn_{}'.format(i))(deep)
 
    if wide_columns and not deep_columns:
        output = tf.keras.layers.Dense(1, activation='sigmoid', name='pred')(wide)
    elif deep_columns and not wide_columns:
        output = tf.keras.layers.Dense(1, activation='sigmoid', name='pred')(deep)
    else:
        both = tf.keras.layers.concatenate([deep, wide], name='both')
        output = tf.keras.layers.Dense(1, activation='sigmoid', name='pred')(both)
 
    model = tn.model.Model(inputs, output)
 
    dense_opt = tn.core.Adam(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8)
    model.compile(optimizer=tn.optimizer.Optimizer(dense_opt),
                  loss='binary_crossentropy',
                  metrics=['accuracy', tf.keras.metrics.AUC(),])
 
    return model
def create_model():
    if C.MODEL_TYPE == "DeepFM":
        return DeepFM(C.LINEAR_SLOTS, C.DEEP_SLOTS, C.DEEP_HIDDEN_UNITS)
    elif C.MODEL_TYPE == "WideDeep":
        return WideDeep(C.LINEAR_SLOTS, C.DEEP_SLOTS, C.DEEP_HIDDEN_UNITS)
    elif C.MODEL_TYPE == "DCN":
        return DCN(C.DEEP_SLOTS, C.DEEP_HIDDEN_UNITS)
    else:
        sys.exit("unsupported model type: " + C.MODEL_TYPE)
# tensornet/examples/main.py
def main():
    strategy = tn.distribute.PsStrategy()
    model, sub_model = create_model()
    dense_opt = tn.core.Adam(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8)
    model.compile(optimizer=tn.optimizer.Optimizer(dense_opt),loss='binary_crossentropy',metrics=['acc', "mse", "mae", 'mape',  tf.keras.metrics.AUC(),tn.metric.CTR(), tn.metric.PCTR(), tn.metric.COPC()])

    train_dataset = read_dataset(...）
    model.fit(train_dataset, epochs=1, verbose=1, callbacks=[cp_cb, tb_cb])
if __name__ == "__main__":
    main()

tn.distribute.PsStrategy()  # 执行 __init__
    super(OneDeviceStrategy, self).__init__(PsExtend(self))
        tn.core.init() # PsExtend() 执行
        # py_wrapper.cc 中由 PYBIND11_MODULE 注册，对应一个匿名函数
        []() {
            PsCluster* cluster = PsCluster::Instance(); 	// 获取ps集群的实例
            if (cluster->IsInitialized()) {                 // 判断集群是否初始化
                return true;
            }
            if (cluster->Init() < 0) {                      // 对ps集群进行初始化操作
                throw py::value_error("Init tensornet fail");
            }
            tensorflow::BalanceInputDataInfo* data_info = tensorflow::BalanceInputDataInfo::Instance();
            if (data_info->Init(cluster->RankNum(), cluster->Rank()) < 0) {
                throw py::value_error("Init BalanceInputDataInfo fail");
            }
            return true;
        }

# when this layer is been called, all the embedding data of `feature_columns` will be pulled from ps server and return as a tensor list.
# tensornet/tensornet/layers/embedding_features.py
class EmbeddingFeatures(Layer):
    def __init__(self,feature_columns,...):
        self._feature_columns = feature_columns
        self._state_manager = StateManagerImpl(self, name, sparse_opt, dim, self.trainable
    def call(self, features, cols_to_output_tensors=None):
        using_features = self.filter_not_used_features(features)
        transformation_cache = fc.FeatureTransformationCache(using_features)
        self.sparse_pulling_features = self.get_sparse_pulling_feature(using_features)
        pulled_mapping_values = self._state_manager.pull(self.sparse_pulling_features)
        output_tensors = []
        for column in self._feature_columns:
            mapping_value = pulled_mapping_values[column.categorical_column.name]
            with ops.control_dependencies([mapping_value]):
                tensor = column.get_dense_tensor(transformation_cache,self._state_manager)
            processed_tensors = self._process_dense_tensor(column, tensor)
            output_tensors.append(processed_tensors)
        return output_tensors
    def backwards(self, grads_and_vars):
        return self._state_manager.push(grads_and_vars, self.sparse_pulling_features)

class StateManagerImpl(fc.StateManager):
    def __init__(self, layer, name, ...):
        self.sparse_table_handle = tn.core.create_sparse_table(sparse_opt, name if name else "", dimension)
        #  py_wrapper.cc 中由 PYBIND11_MODULE 注册，对应一个匿名函数
        [](py::object obj, std::string name, int dimension) {
            OptimizerBase* opt = static_cast<OptimizerBase*>(PyCapsule_GetPointer(obj.ptr(), nullptr));
            PsCluster* cluster = PsCluster::Instance();
            SparseTable* table = CreateSparseTable(opt, name, dimension, cluster->RankNum(), cluster->Rank());
            return table->GetHandle();
        }
        self.pulled_mapping_values = {}
        self._cols_to_var_map = collections.defaultdict(lambda: None)
        self._var_to_cols_map = collections.defaultdict(lambda: None)
    def create_variable(self,feature_column,name,shape,...)
    def get_variable(self, feature_column, name)
    def pull(self, features):
        ...
        pulled_mapping_values = gen_sparse_table_ops.sparse_table_pull(
            [var.handle for var in vars],
            [f.values for f in feature_values],
            table_handle=self.sparse_table_handle)
        ...
    def push(self, grads_and_vars, features):
        ...
        return gen_sparse_table_ops.sparse_table_push(feature_values, grads,table_handle=self.sparse_table_handle)
    def get_feature_mapping_values(self, column_name)
    def save_sparse_table(self, filepath, mode):
        return tn.core.save_sparse_table(self.sparse_table_handle, filepath, mode)
    def load_sparse_table(self, filepath, mode):
        return tn.core.load_sparse_table(self.sparse_table_handle, filepath, mode)
    def show_decay(self):
        return tn.core.show_decay(self.sparse_table_handle)

# tensorflow/tensorflow/python/training/optimizer.py
class Optimizer(object):
  def minimize(self, loss, global_step=None, var_list=None,...):
    grads_and_vars = self.compute_gradients(loss, var_list=var_list, ...)
    vars_with_grad = [v for g, v in grads_and_vars if g is not None]
    return self.apply_gradients(grads_and_vars, global_step=gl
  def compute_gradients(self, loss, var_list=None,...):
      ...
  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    # 调用update_op
  def _apply_dense(self, grad, var):
      raise NotImplementedError()
  def _apply_sparse_duplicate_indices(self, grad, var):
    ...
    return self._apply_sparse(gradient_no_duplicate_indices, var)
  def _apply_sparse(self, grad, var):
    raise NotImplementedError()
class _RefVariableProcessor(_OptimizableVariable):
  # g ==> 梯度, self._v ==> 待更新的variable
  def update_op(self, optimizer, g):
    if isinstance(g, ops.Tensor):
      update_op = optimizer._apply_dense(g, self._v)
      return update_op
    else:
      assert isinstance(g, ops.IndexedSlices), ("Gradient ", g, " is neither a tensor nor IndexedSlices.")
      return optimizer._apply_sparse_duplicate_indices(g, self._v)
# tensorflow/python/training/gradient_descent.py
class GradientDescentOptimizer(optimizer.Optimizer):
  def __init__(self, learning_rate, use_locking=False, name="GradientDescent"):
    super(GradientDescentOptimizer, self).__init__(use_locking, name)
    self._learning_rate = learning_rate
  def _apply_dense(self, grad, var):
    return training_ops.apply_gradient_descent(var,math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),grad,use_locking=self._use_locking).op
  def _apply_sparse_duplicate_indices(self, grad, var):
    delta = ops.IndexedSlices(grad.values * math_ops.cast(self._learning_rate_tensor, var.dtype.base_dtype),grad.indices, grad.dense_shape)
    return var.scatter_sub(delta, use_locking=self._use_locking)

# tensornet/examples/main.py
dense_opt = tn.core.Adam(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8)
model.compile(optimizer=tn.optimizer.Optimizer(dense_opt),loss='binary_crossentropy',metrics=['acc', "mse", "mae", 'mape',  tf.keras.metrics.AUC(),tn.metric.CTR(), tn.metric.PCTR(), tn.metric.COPC()])

# tensorflow/python/keras/optimizer_v2/optimizer_v2.py
# 实际使用的是其子类 tf.keras.optimizers.SGD`, `tf.keras.optimizers.Adam`
class OptimizerV2(trackable.Trackable):
    def minimize(self, loss, var_list, grad_loss=None, name=None, tape=None):
        grads_and_vars = self._compute_gradients(loss, var_list=var_list, grad_loss=grad_loss, tape=tape)
        return self.apply_gradients(grads_and_vars, name=name) 
    # Add ops to apply sparse gradients to the variable `handle`
    def _resource_apply_sparse(self, grad, handle, indices, apply_state):
        raise NotImplementedError("Must be implemented in subclasses.")
    # Add ops to apply dense gradients to the variable `handle`
    def _resource_apply_dense(self, grad, handle, apply_state):
        raise NotImplementedError("Must be implemented in subclasses.")

# tensornet/optimizer/optimizer.py
class Optimizer(optimizer_v2.OptimizerV2):
    def __init__(self,dense_opt,name='TensornetOptimizer',**kwargs):
        self.dense_table_handle = tn.core.create_dense_table(dense_opt)
        ...
    def save_dense_table(self, filepath):
        return tn.core.save_dense_table(self.dense_table_handle, filepath)
    def load_dense_table(self, filepath):
        return tn.core.load_dense_table(self.dense_table_handle, filepath)
    def _distributed_apply(self, distribution, grads_and_vars, name, apply_state):
        ...
        # 对应 DenseTablePushPullKernel
        gen_dense_table_ops.dense_table_push_pull(vars, grads, table_handle=self.dense_table_handle)
        super(Optimizer, self)._distributed_apply(distribution, grads_and_vars, name, apply_state)
    def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
        return control_flow_ops.no_op()
    def _resource_apply_dense(self, grad, var, apply_state=None):
        return control_flow_ops.no_op()

# tensorflow/python/keras/engine/training.py
class Model(base_layer.Layer, version_utils.ModelVersionSelector):
    def train_step(self, data):
        ...
        data = data_adapter.expand_1d(data)
        x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)
        # Run forward pass.
        with backprop.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compiled_loss(y, y_pred, sample_weight, regularization_losses=self.losses)
        # Run backwards pass.
        self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
        ...

# tensornet/tensornet/model/Model.py
class Model(tf.keras.Model):
    # This method should contain the mathematical logic for one step of training. This typically includes the forward pass,loss calculation, backpropagation,and metric updates.
    def train_step(self, data):
        data = data_adapter.expand_1d(data)
        x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)
        with backprop.GradientTape() as tape:
            y_pred = self(x, training=True) # 计算预测值
            loss = self.compiled_loss(y, y_pred, sample_weight, regularization_losses=self.losses) # 计算损失
        gradients = tape.gradient(loss, self.trainable_variables)                # 计算梯度
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables)) # 更新dense参数
        self.backwards(list(zip(gradients, self.trainable_variables)))           # 更新id特征的参数，更新方法在EmbeddingFeatures中
        ...
    def backwards(self, grads_and_vars):
        backward_ops = []
        for layer in self.layers:
            if not hasattr(layer, 'backwards'):
                continue
            op = layer.backwards(grads_and_vars)
            if op:
                backward_ops.append(op)
            ...