import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from datasets import load_dataset
from transformers import BertTokenizer, BertModel
from torch.utils.data import DataLoader

# 0.测试load_dataset方法
def dm01():
    # data = load_dataset('csv', data_files='../data/train.csv')
    # print("data-->", data)
    # data = load_dataset('csv', data_files='../data/train.csv', split='train')
    # print("data-->", data)
    # 注意：split，只能是train
    data = load_dataset('csv', data_files='../data/test.csv', split='train')
    print("data-->", data)
    print(data[0:4])


# 加载预训练模型
tokenizer = BertTokenizer.from_pretrained("../models/bert-base-chinese")
# 这里导入模型要使用BertModel，而不是BertModelForSequenceClassification
# 因为，这里使用预训练模型进行特征提取，下游自己指定多分类输出层（线性层）
bert_model = BertModel.from_pretrained("../models/bert-base-chinese")


# 1.数据预处理
def collate_fn1(batch):
    # 打印数据
    # print("batch-->", batch)
    # 从样本中提取文本和标签
    sentences = [row['text'] for row in batch]
    labels = [row['label'] for row in batch]
    # print("sentences-->", sentences)
    # print("labels-->", labels)
    # 对文本进行数值化+张量化，对标签进行张量化
    data = tokenizer.batch_encode_plus(sentences,
                                        padding='max_length',
                                        max_length=300,
                                        truncation=True,
                                        return_tensors='pt')
    # print("data-->", data)
    # input_ids：数据数值化+张量化后的结果
    # token_type_ids：文本（句子）类型，0代表一个句子
    # attention_mask：注意力掩码，0代表填充的，1代表有值（需要掩码）
    input_ids = data['input_ids']
    token_type_ids = data['token_type_ids']
    attention_mask = data['attention_mask']

    # 对label进行张量化
    labels = torch.tensor(labels)

    # 返回数据
    return input_ids, token_type_ids, attention_mask, labels


def load_data():
    # 使用load_dataset加载数据，得到的就是一个Dataset对象
    mydataset = load_dataset('csv', data_files='../data/train.csv', split='train')
    # 把Dataset转换为DataLoader对象
    # collate_fn：对数据进行二次处理的函数，就是这一批次的数据，需要接受一个函数地址
    mydataloader = DataLoader(dataset=mydataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=collate_fn1)
    # print("mydataloader-->", mydataloader)
    # 遍历
    for input_ids, token_type_ids, attention_mask, labels in mydataloader:
        print("input_ids-->", input_ids)
        print("token_type_ids-->", input_ids)
        print("attention_mask-->", input_ids)
        print("labels-->", input_ids)
        break

# 2.模型构建
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 自定义下游输出层（线性层）
        # in_features：输入维度，就是Bert模型的768维度
        # out_features：输出类别，也就是下游的任务类别，这里是2分类
        self.linear = nn.Linear(in_features=768, out_features=2)

    def forward(self, input_ids, token_type_ids, attention_mask):
        # 先试用Bert模型进行特征提取(可以冻结，也可以不冻结，如果效果不好，就不冻结)
        with torch.no_grad():
            out = bert_model(input_ids=input_ids,
                             token_type_ids=token_type_ids,
                             attention_mask=attention_mask)
        # print("out1-->", out)
        # print("out2-->", out.last_hidden_state.shape) # [4, 300, 768]
        # ：代表取4句话
        # 0：300个token中，取第0个，也就是CLS
        # out.last_hidden_state[:, 0]
        # out.last_hidden_state[:, 0, :]
        # print("out3-->", out.pooler_output.shape)     # [4, 768]
        # 然后把特征送给下游的自定义输出线性层，实现二分类任务
        output = self.linear(out.pooler_output)
        # print("output-->", output, output.shape) # [4, 2]
        return output


def test_model():
    # 数据
    mydataset = load_dataset('csv', data_files='../data/train.csv', split='train')
    mydataloader = DataLoader(dataset=mydataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=collate_fn1)
    # 模型
    model = MyModel()
    # 遍历
    for input_ids, token_type_ids, attention_mask, labels in mydataloader:
        model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
        break


# 3.模型训练
def train_model():
    # 3+2+5
    # 3：模型，损失函数、优化器
    # 2：内外两层循环
    # 5：内层循环中，5个步骤：前向传播、损失计算、梯度清零、反向传播、梯度更新
    # 模型
    model = MyModel()
    # 损失函数
    criterion = nn.CrossEntropyLoss()
    # 优化器
    optimizer = optim.AdamW(params=model.parameters(), lr=5e-4)
    # 轮次
    epochs = 1
    # 冻结预训练模型的参数(Bert模型)
    # 双重保险
    for param in bert_model.parameters():
        param.requires_grad_(False)
    # 设置模型为训练模型
    model.train()
    # 加载数据
    mydataset = load_dataset('csv', data_files='../data/train.csv', split='train')
    # 把数据转换为DataLoader
    mydataloader = DataLoader(dataset=mydataset,
                              shuffle=True,
                              drop_last=True,
                              batch_size=4,
                              collate_fn=collate_fn1)
    # 外层变量轮次
    for epoch in range(epochs):
    # 内层变量批次
    # tqdm：进度条，pip install tqdm
        for i, (input_ids, token_type_ids, attention_mask, labels) in enumerate(tqdm(mydataloader), start=1):
            # 前向传播
            output = model(input_ids=input_ids,
                           token_type_ids=token_type_ids,
                           attention_mask=attention_mask)
            # output.shape->[4,2]
            # 损失计算
            loss = criterion(output, labels)
            # print("loss-->", loss)
            # 梯度清零
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 梯度更新
            optimizer.step()
            # 打印日志
            if i % 5 == 0:
                pred = torch.argmax(output, dim=-1)
                accuracy = (pred == labels).sum().item() / len(labels)
                print("epoch:", epoch + 1, "loss:", loss.item(), "accuracy:", accuracy)
        # 模型保存
        torch.save(model.state_dict(), '../save_model/model_%d.pth' % (epoch + 1))


# 4.模型评估
def model_eval():
    # 加载数据
    mydataset = load_dataset("csv", data_files='../data/test.csv', split='train')
    # 转换为DataLoader
    mydataloader = DataLoader(dataset=mydataset,
                              shuffle=False,
                              drop_last=True,
                              batch_size=4,
                              collate_fn=collate_fn1)
    # 准备模型
    model = MyModel()
    # 加载刚刚预训练好的模型参数
    model.load_state_dict(torch.load('../save_model/model_1.pth'))
    # 准备参数：total，accuracy
    num = 0
    total = 0
    # 设置模型为评估模式
    model.eval()
    # 遍历
    for i, (input_ids, token_type_ids, attention_mask, labels) in enumerate(tqdm(mydataloader), start=1):
        # 把数据送给模型，得到预测值
        output = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
        # 基于预测结果，得到预测值
        y_pred = torch.argmax(output, dim=-1)
        num += (y_pred == labels).sum().item()
        total += len(labels)
    # 计算准确率
    print("accuracy-->", num / total)


if __name__ == '__main__':
    # dm01()
    # load_data()
    # test_model()
    # train_model()
    model_eval()
    pass