# 随机挑一个样本观察重构效果
import torch
from models.cae import ConvAutoEncoder
import torch.utils.data as data
from data_utils.data_set import ReconLeadDataSet
import matplotlib.pyplot as plt

cae_model = ConvAutoEncoder()

cae_model.load_state_dict(torch.load("./cae_recon_best_ckpt.pth"))
cae_model.eval()

dataset = ReconLeadDataSet("/Users/wenhan/Develop/PythonProjects/DATA/DataArray/test_x.npy", from_ld=8, to_ld=11)

test_loader = data.DataLoader(dataset, 128, shuffle=True)
ecg_from = None
ecg_to = None
outputs = None
with torch.no_grad():
    for ecg_from, ecg_to in test_loader:
        outputs = cae_model(ecg_from)
        outputs = outputs.numpy()
        print(outputs.shape)
        break

idx = 0
sample_ecg_from = ecg_from[idx][0].numpy()
sample_ecg_to = ecg_to[idx][0].numpy()
sample_ecg_recon = outputs[idx][0]

plt.figure(figsize=[10, 5])
plt.subplot(2, 1, 1)
plt.plot(sample_ecg_from)
plt.title("V3")
plt.grid(True)

plt.subplot(2, 1, 2)
plt.plot(sample_ecg_to, label="target ECG (V6)")  # 重构目标导联的真实ECG
plt.grid(True)
plt.plot(sample_ecg_recon, "r", label="reconstructed ECG (V6)")  # 实际重构出来的ECG
plt.grid(True)
plt.legend()
plt.show()
# 越重合效果越好