import numpy as np
import tritonclient.http as httpclient

# 1. 初始化客户端
client = httpclient.InferenceServerClient(url="localhost:8000")

# 2. 准备两个 2x2 的随机矩阵（为了显示清晰，用较小的尺寸）
A = np.random.randint(1, 10, size=(2, 2)).astype(np.int32)
B = np.random.randint(1, 10, size=(2, 2)).astype(np.int32)

# 3. 封装 Triton 输入
inputs = [
    httpclient.InferInput("INPUT0", A.shape, "INT32"),
    httpclient.InferInput("INPUT1", B.shape, "INT32")
]
inputs[0].set_data_from_numpy(A)
inputs[1].set_data_from_numpy(B)

# 4. 执行推理
response = client.infer(model_name="add", inputs=inputs)
C = response.as_numpy("OUTPUT0")

# 5. 格式化打印 A + B = C
print("--- 矩阵加法结果 ---")
# 使用 numpy 的 array2string 配合自定义格式
def fmt(mat): return np.array2string(mat, precision=1, separator=', ')

print(f"{fmt(A)}\n\n  +\n\n{fmt(B)}\n\n  =\n\n{fmt(C)}")

# 如果你想要更紧凑的同行对比 (仅限小矩阵)
print("\n--- 同行对比 (A + B = C) ---")
for i in range(len(A)):
    plus = "+" if i == len(A)//2 else " "
    equal = "=" if i == len(A)//2 else " "
    print(f"{A[i]}  {plus}  {B[i]}  {equal}  {C[i]}")