import torch
import torch.nn as nn
import torchvision
import numpy as np
import os
from PIL import Image

# 定义LSTM模型
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)

        out, (h_n,c_n) = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out

# 设置超参数
input_size = 28
sequence_length = 28
hidden_size = 128
num_layers = 2
num_classes = 10
batch_size = 100
num_epochs = 1
learning_rate = 0.001

# 加载MNIST数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=torchvision.transforms.ToTensor())

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 创建LSTM模型
model = LSTMModel(input_size, hidden_size, num_layers, num_classes)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.reshape(-1, sequence_length, input_size)
        outputs = model(images)
        loss = criterion(outputs, labels)

        predictions = torch.argmax(outputs,dim=1)
        # acc = torch.eq(predictions,labels).sum().item()
        # print(acc)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))

# 保存模型
torch.save(model.state_dict(), 'model.pth')

# 加载模型
model.load_state_dict(torch.load('model.pth'))
with torch.no_grad():
    for i, (images, labels) in enumerate(test_loader):
        images = images.reshape(-1, sequence_length, input_size)
        outputs = model(images)
        predictions = torch.argmax(outputs,dim=1)
        acc = torch.eq(predictions,labels).sum().item()
        print(acc)

# folder_path = './mnist_pi'  # 替换为图片所在的文件夹路径
# def infer_images_in_folder(folder_path):
#     with torch.no_grad():
#         for file_name in os.listdir(folder_path):
#             file_path = os.path.join(folder_path, file_name)
#             if os.path.isfile(file_path) and file_name.endswith(('.jpg', '.jpeg', '.png')):
#                 image = Image.open(file_path)
#                 label = file_name.split(".")[0].split("_")[1]
#                 image = np.array(image)/255.0
#                 image = np.expand_dims(image,axis=0)
#                 image= torch.tensor(image).to(torch.float32)
#                 logits = model(image)
#                 predicted_class = torch.argmax(logits)
#                 print("file_path:",file_path,"img size:",image.shape,"label:",label,'Predicted class:', predicted_class)
#             break

# infer_images_in_folder(folder_path)



# 保存模型参数为numpy的数组格式
model_params = {}
# print(list(model.parameters()))
for name, param in model.named_parameters():
    model_params[name] = param.detach().numpy()
    print(name,param.shape)

np.savez('model.npz', **model_params)

import numpy as np
import os
from PIL import Image

# 加载模型参数
model_data = np.load('model.npz')

'''
weight_ih_l[k] : the learnable input-hidden weights of the :math:`\text{k}^{th}` layer
    `(W_ii|W_if|W_ig|W_io)`, of shape `(4*hidden_size, input_size)` for `k = 0`.
    Otherwise, the shape is `(4*hidden_size, num_directions * hidden_size)`. If
    ``proj_size > 0`` was specified, the shape will be
    `(4*hidden_size, num_directions * proj_size)` for `k > 0`
weight_hh_l[k] : the learnable hidden-hidden weights of the :math:`\text{k}^{th}` layer
    `(W_hi|W_hf|W_hg|W_ho)`, of shape `(4*hidden_size, hidden_size)`. If ``proj_size > 0``
    was specified, the shape will be `(4*hidden_size, proj_size)`.
bias_ih_l[k] : the learnable input-hidden bias of the :math:`\text{k}^{th}` layer
    `(b_ii|b_if|b_ig|b_io)`, of shape `(4*hidden_size)`
bias_hh_l[k] : the learnable hidden-hidden bias of the :math:`\text{k}^{th}` layer
    `(b_hi|b_hf|b_hg|b_ho)`, of shape `(4*hidden_size)`
'''


# 提取模型参数
lstm_weight_ih_l0 = model_data['lstm.weight_ih_l0']
lstm_weight_hh_l0 = model_data['lstm.weight_hh_l0']
lstm_bias_ih_l0 = model_data['lstm.bias_ih_l0']
lstm_bias_hh_l0 = model_data['lstm.bias_hh_l0']
lstm_weight_ih_l1 = model_data['lstm.weight_ih_l1']
lstm_weight_hh_l1 = model_data['lstm.weight_hh_l1']
lstm_bias_ih_l1 = model_data['lstm.bias_ih_l1']
lstm_bias_hh_l1 = model_data['lstm.bias_hh_l1']
fc_weight = model_data['fc.weight']
fc_bias = model_data['fc.bias']

# print(lstm_weight_ih_l0.shape,lstm_weight_hh_l0.shape)
# print(lstm_bias_ih_l0.shape,lstm_bias_hh_l0.shape)
# 定义LSTM模型
def lstm_model(inputs):
    '''
        踩到两个坑，一个是矩阵形状都是这种(4*hidden_size, hidden_size)合并的，需要拆分。
        另一个坑是，两层的lstm层需要每个时间步的输出都输入到下一层，而不是最后一个时间步的数据给下一层
    '''

    batch_size, sequence_length, input_size = inputs.shape
    hidden_size = lstm_weight_hh_l0.shape[1]
    num_classes = fc_weight.shape[0]


    h0 = np.zeros((batch_size, hidden_size))
    c0 = np.zeros((batch_size, hidden_size))

    # 第一层LSTM
    h_l0, c_l0 = np.zeros_like(h0), np.zeros_like(c0)
    out_0 = []
    for t in range(sequence_length):
        x = inputs[:, t, :]
        '''
        i_t = \sigma(W_{ii} x_t + b_{ii} + W_{hi} h_{t-1} + b_{hi}) \\
        f_t = \sigma(W_{if} x_t + b_{if} + W_{hf} h_{t-1} + b_{hf}) \\
        g_t = \tanh(W_{ig} x_t + b_{ig} + W_{hg} h_{t-1} + b_{hg}) \\
        o_t = \sigma(W_{io} x_t + b_{io} + W_{ho} h_{t-1} + b_{ho}) \\
        c_t = f_t \odot c_{t-1} + i_t \odot g_t \\
        h_t = o_t \odot \tanh(c_t) \\
        '''
        # 输入门
        i_t = sigmoid(np.dot(x, lstm_weight_ih_l0[:128].T) + np.dot(h_l0, lstm_weight_hh_l0[:128].T) + lstm_bias_ih_l0[:128] + lstm_bias_hh_l0[:128])
        # 遗忘门
        f_t = sigmoid(np.dot(x, lstm_weight_ih_l0[128:256].T) + np.dot(h_l0, lstm_weight_hh_l0[128:256].T) + lstm_bias_ih_l0[128:256] + lstm_bias_hh_l0[128:256])
        # 候选向量
        g_t = np.tanh(np.dot(x, lstm_weight_ih_l0[256:256+128].T) + np.dot(h_l0, lstm_weight_hh_l0[256:256+128].T) + lstm_bias_ih_l0[256:256+128] + lstm_bias_hh_l0[256:256+128])
        # 输出门
        o_t = sigmoid(np.dot(x, lstm_weight_ih_l0[256+128:512].T) + np.dot(h_l0, lstm_weight_hh_l0[256+128:512].T) + lstm_bias_ih_l0[256+128:512] + lstm_bias_hh_l0[256+128:512])
        # 细胞状态
        c_l0 = f_t * c_l0 + i_t * g_t
        # 隐藏状态
        h_l0 = o_t * np.tanh(c_l0)
        out_0.append(h_l0)

    # 第二层LSTM
    h_l1, c_l1 = np.zeros_like(h0), np.zeros_like(c0)
    out_1 = []
    for t in range(sequence_length):
        x = out_0[t]
        # 输入门
        i_t = sigmoid(np.dot(x, lstm_weight_ih_l1[:128].T) + np.dot(h_l1, lstm_weight_hh_l1[:128].T) + lstm_bias_ih_l1[:128] + lstm_bias_hh_l1[:128])
        # 遗忘门
        f_t = sigmoid(np.dot(x, lstm_weight_ih_l1[128:256].T) + np.dot(h_l1, lstm_weight_hh_l1[128:256].T) + lstm_bias_ih_l1[128:256] + lstm_bias_hh_l1[128:256])
        # 候选向量
        g_t = np.tanh(np.dot(x, lstm_weight_ih_l1[256:256+128].T) + np.dot(h_l1, lstm_weight_hh_l1[256:256+128].T) + lstm_bias_ih_l1[256:256+128] + lstm_bias_hh_l1[256:256+128])
        # 输出门
        o_t = sigmoid(np.dot(x, lstm_weight_ih_l1[256+128:512].T) + np.dot(h_l1, lstm_weight_hh_l1[256+128:512].T) + lstm_bias_ih_l1[256+128:512] + lstm_bias_hh_l1[256+128:512])
        # 细胞状态
        c_l1 = f_t * c_l1 + i_t * g_t
        # 隐藏状态
        h_l1 = o_t * np.tanh(c_l1)
        out_1.append(h_l1)

    # 全连接层
    fc_output = np.dot(h_l1, fc_weight.T) + fc_bias
    predictions = np.argmax(fc_output, axis=1)
    return predictions

# Sigmoid激活函数
def sigmoid(x):
    return 1 / (1 + np.exp(-x))


folder_path = './mnist_pi'  # 替换为图片所在的文件夹路径
def infer_images_in_folder(folder_path):
    for file_name in os.listdir(folder_path):
        file_path = os.path.join(folder_path, file_name)
        if os.path.isfile(file_path) and file_name.endswith(('.jpg', '.jpeg', '.png')):
            image = Image.open(file_path)
            label = file_name.split(".")[0].split("_")[1]
            image = np.array(image)/255.0
            image = np.expand_dims(image,axis=0)
            predicted_class = lstm_model(image)
            print("file_path:",file_path,"img size:",image.shape,"label:",label,'Predicted class:', predicted_class)
            

infer_images_in_folder(folder_path)