import mindspore
import mindspore.nn as nn
import onnx
import onnxruntime
from mindspore import Tensor
from mindspore.dataset import Cifar10Dataset
from onnx import numpy_helper

from tests.secret_func import verify_secret
from watermark_codec import ModelDecoder

mindspore.set_context(device_target="CPU", max_device_memory="4GB")
train_dataset = Cifar10Dataset(dataset_dir='data/cifar-10-batches-bin', usage='train', shuffle=True)
test_dataset = Cifar10Dataset(dataset_dir='data/cifar-10-batches-bin', usage='test')
batch_size = 32
key_path = './run/train/key.ckpt'
save_path = './run/train/AlexNet.onnx'

# init onnx session and model
ort_session = onnxruntime.InferenceSession(save_path)
model = onnx.load(save_path)

# 获取目标模型权重
# 遍历模型中的所有初始化器
weights = []
for initializer in model.graph.initializer:
    # 检查初始化器的名称是否与卷积层的权重或偏置匹配
    # 注意：这里需要根据实际模型中的命名规则来修改匹配逻辑
    # 将权重或偏置从ONNX的TensorProto格式转换为NumPy数组
    tensor = numpy_helper.to_array(initializer)

    # 根据名称判断是权重还是偏置，并存储到字典中
    if 'weight' in initializer.name:
        weights.append(tensor)
weights = weights[0:2]
weights = [Tensor(x) for x in weights]

# init model decoder
model_decoder = ModelDecoder(weights=weights, key_path=key_path)

secret = model_decoder.decode()
print(f"secret: {secret}")

result = verify_secret(secret)
print(f"result: {result}")