class CustomGRU_layer(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(CustomGRU_layer, self).__init__()
        # 初始化参数
        self.W_xz = nn.Parameter(torch.randn(input_size, hidden_size))  # 更新门的输入到隐藏层的权重
        self.W_hz = nn.Parameter(torch.randn(hidden_size, hidden_size))  # 更新门的隐藏层到隐藏层的权重

        self.W_xr = nn.Parameter(torch.randn(input_size, hidden_size))  # 重置门的输入到隐藏层的权重
        self.W_hr = nn.Parameter(torch.randn(hidden_size, hidden_size))  # 重置门的隐藏层到隐藏层的权重

        self.W_xh = nn.Parameter(torch.randn(input_size, hidden_size))  # 候选隐藏状态的输入到隐藏层的权重
        self.W_hh = nn.Parameter(torch.randn(hidden_size, hidden_size))  # 候选隐藏状态的隐藏层到隐藏层的权重

        self.hb_z = nn.Parameter(torch.zeros(hidden_size))  # 更新门的偏置
        self.hb_r = nn.Parameter(torch.zeros(hidden_size))  # 重置门的偏置
        self.hb_h = nn.Parameter(torch.zeros(hidden_size))  # 候选隐藏状态的偏置
        
        self.xb_z = nn.Parameter(torch.zeros(hidden_size))
        self.xb_r = nn.Parameter(torch.zeros(hidden_size))
        self.xb_h = nn.Parameter(torch.zeros(hidden_size))

    def forward(self, x, h):
        # 前向传播
        z = torch.sigmoid((torch.matmul(x, self.W_xz) + self.xb_z) + (torch.matmul(h, self.W_hz) + self.hb_z))  # 更新门
        r = torch.sigmoid((torch.matmul(x, self.W_xr) + self.xb_r) + (torch.matmul(h, self.W_hr) + self.hb_r))  # 重置门
        h_tilda = torch.tanh((torch.matmul(x, self.W_xh) + self.xb_h) + r * (torch.matmul(h, self.W_hh) + self.hb_h))  # 候选隐藏状态
        h = z * h + (1 - z) * h_tilda  # 更新隐藏状态
        return h

class CustomGRU(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(CustomGRU, self).__init__()
        self.input_size = input_size  # 输入特征的维度
        self.hidden_size = hidden_size  # 隐藏层的维度

        # 初始化自定义的GRU层
        self.gru = CustomGRU_layer(self.input_size, self.hidden_size)

    def forward(self, X, h0 = None):
        # x.shape = (seq_length, batch_size, input_size)
        # h0.shape = (1, batch_size, hidden_size)
        # output.shape = (seq_length, batch_size, hidden_size)

        # 获取批次大小
        batch_size = X.shape[1]
        # 获取序列长度
        self.seq_length = X.shape[0]

        # 如果没有提供初始隐藏状态，则初始化为零张量
        if h0 is None:
            prev_h = torch.zeros([batch_size, self.hidden_size]).to(device)
        else:
            prev_h = torch.squeeze(h0, 0)

        # 初始化输出张量
        output = torch.zeros([self.seq_length, batch_size, self.hidden_size]).to(device)

        # 循环处理序列中的每个时间步
        for i in range(self.seq_length):
            # 通过GRU层处理当前时间步的数据，并更新隐藏状态
            prev_h = self.gru(X[i], prev_h)
            # 将当前时间步的输出存储在输出张量中
            output[i] = prev_h

        # 返回最终的输出和隐藏状态
        return output, torch.unsqueeze(prev_h, 0)

# 将nn.GRU中的4个随机初始化的可学习参数进行保存，并用来替换CustomGRU中CustomGRU_layer随机初始化的可学习参数，并通过torch.allclose判断输出是否相等，若相等则证明MyGRU的实现与官方的nn.GRU是一致的

# 初始化nn.GRU
gru = nn.GRU(input_size=input_size, hidden_size=hidden_size).to(device)
weight_ih_l0 = gru.weight_ih_l0.T
weight_hh_l0 = gru.weight_hh_l0.T
bias_ih_l0 = gru.bias_ih_l0
bias_hh_l0 = gru.bias_hh_l0

# 初始化CustomGRU
custom_gru = CustomGRU(input_size=input_size, hidden_size=hidden_size).to(device)

# 替换CustomGRU中的参数
custom_gru.gru.W_xr = nn.Parameter(weight_ih_l0[:, :custom_gru.gru.W_xr.size(1)])  # 更新门的输入权重
custom_gru.gru.W_hr = nn.Parameter(weight_hh_l0[:, :custom_gru.gru.W_hr.size(1)])  # 更新门的隐藏权重

custom_gru.gru.W_xz = nn.Parameter(weight_ih_l0[:, custom_gru.gru.W_xr.size(1):custom_gru.gru.W_xr.size(1) + custom_gru.gru.W_xz.size(1)])  # 重置门的输入权重
custom_gru.gru.W_hz = nn.Parameter(weight_hh_l0[:, custom_gru.gru.W_hr.size(1):custom_gru.gru.W_hr.size(1) + custom_gru.gru.W_hz.size(1)])  # 重置门的隐藏权重

custom_gru.gru.W_xh = nn.Parameter(weight_ih_l0[:, custom_gru.gru.W_xr.size(1) + custom_gru.gru.W_xz.size(1):])  # 候选隐藏状态的输入权重
custom_gru.gru.W_hh = nn.Parameter(weight_hh_l0[:, custom_gru.gru.W_hr.size(1) + custom_gru.gru.W_hz.size(1):])  # 候选隐藏状态的隐藏权重

custom_gru.gru.hb_r = nn.Parameter(bias_hh_l0[:custom_gru.gru.hb_r.size(0)])  # 更新门的偏置
custom_gru.gru.hb_z = nn.Parameter(bias_hh_l0[custom_gru.gru.hb_r.size(0):custom_gru.gru.hb_z.size(0) + custom_gru.gru.hb_r.size(0)])  # 重置门的偏置
custom_gru.gru.hb_h = nn.Parameter(bias_hh_l0[custom_gru.gru.hb_z.size(0) + custom_gru.gru.hb_r.size(0):])  # 候选隐藏状态的偏置

custom_gru.gru.xb_r = nn.Parameter(bias_ih_l0[:custom_gru.gru.xb_r.size(0)])
custom_gru.gru.xb_z = nn.Parameter(bias_ih_l0[custom_gru.gru.xb_r.size(0):custom_gru.gru.xb_z.size(0) + custom_gru.gru.xb_r.size(0)])
custom_gru.gru.xb_h = nn.Parameter(bias_ih_l0[custom_gru.gru.xb_z.size(0) + custom_gru.gru.xb_r.size(0):])

# 初始化输入数据
x = torch.rand(seq_length, batch_size, input_size).to(device)

# 获取CustomGRU和nn.GRU的输出
output1, h1 = custom_gru(x)
output2, h2 = gru(x)

# 使用torch.allclose比较输出是否相等
print("output1 == output2 ? {}".format(torch.allclose(output1.to('cpu'), output2.to('cpu'), atol=1e-6)))
print("h1 == h2 ? {}".format(torch.allclose(h1.to('cpu'), h2.to('cpu'), atol=1e-6)))

# 读取文本文件
with open('data/msra_train_bio.txt', 'r', encoding='utf-8') as file:
    lines = file.readlines()

# 用于存储最终的实体和标签
texts = []

text = []
for line in lines:
    if line == '\n':
        texts.append(text)

        text = []
        label = []
    else:
        parts = line.split('\t')
        text.append(parts[0])

print(len(texts))
# print(texts)

# 统计词频
# 展平二维列表
all_tokens = [token for sublist in texts for token in sublist]

# 统计词频
token_counts = Counter(all_tokens)

print(len(token_counts))

# 创建一个字典，将每个字符映射到一个唯一的整数
char_to_idx = {char: i + 1 for i, char in enumerate(set(''.join(all_tokens)).union({'<PAD>'}))}

class MyDataset(Dataset):
    def __init__(self, file):
        super(MyDataset, self).__init__()
        # 读取文本文件
        with open(file, 'r', encoding='utf-8') as file:
            lines = file.readlines()

        # 用于存储最终的实体和标签
        texts = []
        labels = []

        text = []
        label = []
        for line in lines:
            if line == '\n':
                # 将每个字符转换为整数
                encoded_texts = [char_to_idx[char] if char in char_to_idx else 0 for char in text]
                texts.append(encoded_texts)
                try:
                    labels.append([label_dict[num] if num != '' else 0 for num in label])
                except:
                    pass
                text = []
                label = []
            else:
                parts = line.split('\t')
                text.append(parts[0])
                label.append(parts[-1].strip() if parts[-1].strip() == 'O' else parts[-1].strip()[2:])

        self.texts = texts
        self.labels = labels

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, item):
        return self.texts[item], self.labels[item]

# 收集所有文本长度
lengths = [len(text) for text, label in dataset]

# 计算最大长度
max_length = max(lengths)
print("最大文本长度:", max_length)

# 绘制直方图
plt.hist(lengths, bins=100, edgecolor='black')
plt.title('文本长度分布')
plt.xlabel('文本长度')
plt.ylabel('频数')
plt.grid(True)
plt.show()

# 自定义 GRU 层和自定义 GRU 模型直接使用前文的定义
# class CustomGRU_layer
# classCustomGRU


# 定义模型
class Model(nn.Module):
    def __init__(self, input_size, num_class, hidden_size=768):
        super(Model, self).__init__()
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.gru = CustomGRU(input_size, hidden_size)
        # self.gru = nn.GRU(hidden_size, hidden_size)
        self.fc1 = nn.Linear(hidden_size, 256)
        self.fc2 = nn.Linear(256, num_class)
        self.relu = nn.ReLU()

    def forward(self, x):
        embed = self.embedding(x)
        x = self.relu(embed)
        out, _ = self.gru(x)
        out = self.relu(out)
        fc1 = self.fc1(out)
        fc1 = self.relu(fc1)
        out = self.fc2(fc1)
        return out


def collate_fn(data, max_length=seq_length):
    tokens = [i[0] for i in data]  # 提取所有文本
    labels = [i[1] for i in data]  # 提取所有标签
    
    # 处理文本
    new_tokens = []
    for token in tokens:
        # 截断或填充标签以匹配最大序列长度
        new_token = token[:max_length] + [input_size - 1] * (max_length - len(token[:max_length]))
        new_tokens.append(new_token)

    # 将标签列表转换为LongTensor
    tokens_tensor = torch.LongTensor(new_tokens)

    # 处理标签
    new_labels = []
    for label in labels:
        # 截断或填充标签以匹配最大序列长度
        new_label = label[:max_length] + [num_class - 1] * (max_length - len(label[:max_length]))
        new_labels.append(new_label)

    # 将标签列表转换为LongTensor
    labels_tensor = torch.LongTensor(new_labels)

    return tokens_tensor, labels_tensor

def get_correct_and_total_count(labels, outs):
    outs = outs.argmax(dim=2)                   # 获取概率最高的类别作为预测结果
    correct = (outs == labels).sum().item()     # 计算总的正确预测数量
    total = labels.size(0) * labels.size(1)     # 计算总的标签数量

    # 计算除了补充元素以外的正确率
    select = labels != num_class - 1                  # 选择非补充的标签

    labels_non_supplementary = labels[select]                     # 选择非补充的真实标签
    outs_non_supplementary = outs[select]                         # 选择非补充的预测结果

    # 计算F1分数
    # F1分数是衡量分类模型精确度的一种指标，它是精确率和召回率的调和平均数，用于在模型的精确性和完整性之间取得平衡。
    f1 = f1_score(labels_non_supplementary.cpu().numpy(), outs_non_supplementary.cpu().numpy(), average='macro')

    # 由于labels和outs在过滤后可能不再是二维的，需要调整它们的形状
    labels_non_supplementary = labels_non_supplementary.view(-1)  # 展平为一维
    outs_non_supplementary = outs_non_supplementary.view(-1)  # 展平为一维
    correct_content = (outs_non_supplementary == labels_non_supplementary).sum().item()      # 计算非补充部分的正确预测数量
    total_content = len(labels_non_supplementary)                 # 计算非补充部分的标签数量

    return correct, total, correct_content, total_content, f1

def train(dataloader, model, criterion, optimizer, device, scheduler):
    size = len(dataloader.dataset)
    avg_loss = 0
    correct, total, correct_content, total_content = 0, 0, 0, 0

    model.train()
    # for batch, (inputs, labels) in enumerate(tqdm(dataloader, desc="Processing")):
    for batch, (inputs, labels) in enumerate(dataloader):
        # 将数据存到 device
        inputs, labels = inputs.to(device), labels.to(device)

        # 得到预测的结果
        outs = model(inputs)
        loss = criterion(outs.view(-1, num_class) , labels.view(-1))
        avg_loss += loss

        # 反向传播，更新模型参数
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        counts = get_correct_and_total_count(labels, outs)

        correct += counts[0]
        total += counts[1]
        correct_content += counts[2]
        total_content += counts[3]

    # 更新学习率
    scheduler.step()

    # 当一个epoch完了后返回平均 loss
    avg_loss /= size
    avg_loss = avg_loss.detach().cpu().numpy()
    accuracy = correct / total
    accuracy_content = correct_content / total_content

    print(f"train: correct = {100 * accuracy:.3f}%, Accuracy: {(100 * accuracy_content):>0.3f}%, Avg loss: {avg_loss:>8f}, f1_score: {counts[4]:>5f}")
    return accuracy_content, avg_loss


def validate(dataloader, model, loss_fn, device):
    size = len(dataloader.dataset)

    # 将模型转为验证模式
    model.eval()

    # 初始化test_loss 和 correct， 用来统计每次的误差
    test_loss = 0
    correct, total, correct_content, total_content = 0, 0, 0, 0

    # 测试时模型参数不用更新，no_gard()
    with torch.no_grad():
        # 加载数据加载器，得到里面的X（图片数据）和y(真实标签）
        for batch, (inputs, labels) in enumerate(dataloader):
            # 将数据存到 device
            inputs, labels = inputs.to(device), labels.to(device)

            # 得到预测的结果
            outs = model(inputs)
            loss = loss_fn(outs.view(-1, num_class) , labels.view(-1))
            test_loss += loss

            counts = get_correct_and_total_count(labels, outs)

            correct += counts[0]
            total += counts[1]
            correct_content += counts[2]
            total_content += counts[3]

    test_loss /= size
    test_loss = test_loss.detach().cpu().numpy()
    accuracy = correct / total
    accuracy_content = correct_content / total_content
    print(f"valid: correct = {100 * accuracy:.3f}%, Accuracy: {(100 * accuracy_content):>0.3f}%, Avg loss: {test_loss:>8f}, f1_score: {counts[4]:>5f}")
    return accuracy_content, test_loss


def predict(tokens, model, device, max_length=seq_length):
    model = model.to(device)

    inputs = []
    for sentence in tokens:
        # 将每个字符转换为整数
        part = [char_to_idx[char] if char in char_to_idx else 0 for char in sentence]
        inputs.append(part)

    # 处理文本
    new_inputs = []
    for token in inputs:
        # 截断或填充标签以匹配最大序列长度
        new_input = token[:max_length] + [input_size - 1] * (max_length - len(token[:max_length]))

        # 将标签列表转换为LongTensor 
        new_input = torch.LongTensor (new_input)
        new_inputs.append(new_input)

    preds = []
    with torch.no_grad():
        for batch, inputs in enumerate(new_inputs):
            # 将数据存到显卡
            inputs = inputs.to(device)
            size = inputs.size(0)

            # 得到预测的结果
            outs = model(inputs)

            # 获取概率最高的类别作为预测结果
            outs = outs.argmax(dim=-1)     
            preds.append(outs)
            # preds = torch.cat(preds, dim=0)

    preds_txt = []
    for index, pred in enumerate(preds):
        pred = pred[:len(tokens[index])]

        # 创建一个从数字索引到标签名称的映射
        idx_to_label = {v: k for k, v in label_dict.items()}

        # 使用映射将数字索引转换为标签名称
        label_strings = [idx_to_label[idx.item()] for idx in pred]
        preds_txt.append(label_strings)
    return pred, preds_txt

# # 给训练集和测试集分别创建一个数据集加载器
train_data = MyDataset('data/msra_train_bio.txt')
valid_data = MyDataset('data/msra_test_bio.txt')

train_dataloader = DataLoader(
    dataset=train_data,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=collate_fn,
    drop_last=True
)
valid_dataloader = DataLoader(
    dataset=valid_data,
    batch_size=batch_size,
    shuffle=False,
    collate_fn=collate_fn,
    drop_last=True
)

# 输出使用设备
print(f"Using {device} device")

# 创建模型
model = Model(input_size, num_class, hidden_size)
model = model.to(device)

# 定义标签损失权重
class_weights = torch.tensor([1, 10, 10, 10, 0.01]).to(device)
loss_fn = nn.CrossEntropyLoss(weight=class_weights)

# 定义优化器
lr = 1e-5
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

# 定义学习率策略（余弦退火）
# scheduler = CosineAnnealingLR(optimizer, T_max=200, eta_min=1e-10)

scheduler = StepLR(optimizer, step_size=5, gamma=0.95)

epochs = 50
loss_ = 100
save_root = "models/"
avg_loss = 0

train_accs = []
train_losses = []
val_accs = []
val_losses = []

for t in range(epochs):
    print(f"\n----------------------Training Epoch {t} ----------------------------")

    train_accuracy, avg_loss = train(train_dataloader, model, loss_fn, optimizer, device, scheduler)
    val_accuracy, val_loss = validate(valid_dataloader, model, loss_fn, device)

    train_accs.append(train_accuracy) 
    train_losses.append(avg_loss)
    val_accs.append(val_accuracy)
    val_losses.append(val_loss)

    if avg_loss < loss_:
        loss_ = avg_loss
        torch.save(model.state_dict(), save_root + "model_best.pth")

    # torch.save(model.state_dict(), save_root + "model_last.pth")

# 绘制训练损失和验证损失曲线
plt.figure(figsize=(4, 3))
plt.plot(train_losses, label='Train Loss', color='blue')
plt.plot(val_losses, label='Validation Loss', color='red')

plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

# 绘制训练损失和验证正确率曲线
plt.figure(figsize=(4, 3))
plt.plot(train_accs, label='Train Accuracy', color='blue')
plt.plot(val_accs, label='Validation Accuracy', color='red')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Acc')
plt.legend()

# 显示图形
plt.show()

tokens = [
    '在2024年的春天，张伟、李娜和王强三位科学家在中国科学院的资助下，前往云南的西双版纳热带植物园进行生物多样性研究。他们与当地的傣族村民合作，共同探索了这片土地上丰富的植物资源，希望能找到对抗气候变化的新方法。' ,
    '正当朱镕基当选政府总理后第一次在中外记者招待会上，回答外国记者的提问：中国农村是否实行民主选举制度的时候，一位中国电视编导正携带着她的反映中国农村民主选举村委会领导的电视纪录片《村民的选择》（北京电视台摄制，仝丽编导），出现在法国的真实电影节上。',
    '中共中央致中国致公党十一大的贺词各位代表、各位同志：在中国致公党第十一次全国代表大会隆重召开之际，中国共产党中央委员会谨向大会表示热烈的祝贺，向致公党的同志们致以亲切的问候！'
]

save_root = "models/"
model = Model(input_size, num_class, hidden_size).to(device)
model_state_dict = torch.load(save_root + "model_best.pth")
model.load_state_dict(model_state_dict)

preds, preds_txt = predict(tokens, model, device)
for idx, token in enumerate(tokens):
    for num in range(len(token)):
        print(token[num], preds_txt[idx][num])
    print('\n\n')