from mindspore.ops import operations as P
class QDENN(nn.Cell):
    def __init__(self,
                 vocab,
                 embed_size,
                 fix_length,
                 qembed_layers=2,
                 n_threads=10,):
        super(QDENN, self).__init__()
        self.quantum_embedding = QEmbedding(vocab, embed_size, fix_length,
                                    qembed_layers, n_threads)
        self.bn=nn.BatchNorm1d(num_features=embed_size,use_batch_statistics=True)
        self.w0 = Parameter(Tensor(np.random.random(1), mindspore.float32), name="w0", requires_grad=True)
        self.w1 = Parameter(Tensor(np.random.random(1), mindspore.float32), name="w1", requires_grad=True)


    def construct(self, inputs):#输入(batchsize, 630)，一条数据中的每个值都为1.0或0.0
        [self.w1[0] if i > 0 else 0.0  for i in inputs[0][:30]]
        expand_dims = P.ExpandDims()
        tile = P.Tile()
        cast = P.Cast()
        # 将其转换成bool类型，True对应非零元素，False对应零元素
        inputs_bool = cast(inputs, mindspore.bool_)
        # 使用tile和reshape操作扩展权重以匹配inputs的shape
        w0_expanded = tile(expand_dims(self.w0, 1), (inputs.shape[0], inputs.shape[1]))
        w1_expanded = tile(expand_dims(self.w1, 1), (inputs.shape[0], inputs.shape[1]))
        # 使用MindSpore的Select操作根据inputs_bool选择w0或w1
        selected_weights = P.Select()(inputs_bool, w0_expanded, w1_expanded)
        embeddings = self.quantum_embedding(selected_weights)
        outputs=self.bn(embeddings)
        return outputs