# coding:utf8

import torch
import torch.nn as nn
import numpy as np
import random
import json
import matplotlib.pyplot as plt

"""

构建一个模型, 继承torch.nn并初始值

"""


class TorchModel(nn.Module):
    def __init__(self, input_dim, sentence_length, vocab):
        super(TorchModel, self).__init__()
        self.embedding = nn.Embedding(len(vocab) + 1, input_dim)  # 根据字表长度+1和字的维度初始化训练模型
        self.layer = nn.Linear(input_dim, input_dim)  # 字长维度构建线性层
        self.pool = nn.AvgPool1d(sentence_length)  # 池化层,降低维度,缩减模型大小,防止过拟合
        self.classify = nn.Linear(input_dim, 1)  # 指定输出层为1
        self.activation = torch.sigmoid  # 使用激活函数,注意值
        self.dropout = nn.Dropout(0.1)  # 随机丢弃神经元, 减少过拟合,其余元素乘于1/(1-p)进行放大
        self.loss = nn.functional.mse_loss  # loss采用均方差损失

    def forward(self, x, y=None):
        x = self.embedding(x)  # input shape:(batch_size, sen_len)
        x = self.layer(x)  # input shape:(batch_size, sen_len, input_dim)
        x = self.dropout(x)  # input shape:(batch_size, sen_len, input_dim)
        x = self.activation(x)  # input shape:(batch_size, sen_len, input_dim)
        x = self.pool(x.transpose(1, 2)).squeeze()  # input shape:(batch_size, sen_len, input_dim)
        x = self.classify(x)  # input shape:(batch_size, input_dim)
        y_pred = self.activation(x)  # input shape:(batch_size, 1)
        if y is not None:
            return self.loss(y_pred, y)
        else:
            return y_pred


# 建立字表
def build_vocab():
    chars = "abcdefghijklmnopqrstuvwxyz"
    vocab = {}
    for index, char in enumerate(chars):
        vocab[char] = index + 1
    vocab['unk'] = len(vocab) + 1
    return vocab


# 简单随机对比
def build_simple(vocab, sentence_length):
    x = [random.choice(list(vocab.keys())) for _ in range(sentence_length)]
    if set("abc") & set(x):
        y = 1
    else:
        y = 0
    x = [vocab.get(word, vocab['unk']) for word in x]
    return x, y


# 构建数据集
def build_dataset(sample_length, vocab, sentence_length):
    dataset_x = []
    dataset_y = []
    for i in range(sample_length):
        x, y = build_simple(vocab, sentence_length)
        dataset_x.append(x)
        dataset_y.append([y])
    return torch.LongTensor(dataset_x), torch.FloatTensor(dataset_y)


# 根据字表,字的维度,字长建立模型
def build_model(vocab, char_dim, sentence_length):
    model = TorchModel(char_dim, sentence_length, vocab)
    return model


# 用来测试每轮模型的准确率
def evaluate(model, vocab, sample_length):
    model.eval()
    x, y = build_dataset(200, vocab, sample_length)
    print("本次预测集中共有%d个正样本, %d个负样本" % (sum(y), 200 - sum(y)))
    correct, wrong = 0, 0
    with torch.no_grad():
        y_pred = model(x)  # 模型预测
        for y_p, y_t in zip(y_pred, y):  # 与真实标签进行对比  # zip: 元素个数与最短的列表一致
            if float(y_p) < 0.5 and int(y_t) == 0:
                correct += 1  # 负样本判断正确
            elif float(y_p) >= 0.5 and int(y_t) == 1:
                correct += 1  # 正样本判断正确
            else:
                wrong += 1
    print("正确预测个数: %d, 正确率: %f" % (correct, correct / (correct + wrong)))
    return correct / (correct + wrong)


def main():
    epoch_num = 10  # 训练轮次
    batch_size = 20  # 每次训练时样本个数
    train_sample = 1000  # 每轮训练总共训练的样本总数
    char_dim = 20  # 每个字的维度
    sentence_length = 6  # 样本文本长度
    vocab = build_vocab()  # 建立字表
    model = build_model(vocab, char_dim, sentence_length)  # 建立模型
    optim = torch.optim.Adam(model.parameters(), lr=0.005)  # 建立优化器
    log = []
    for epoch in range(epoch_num):  # 训练轮次
        model.train()  # 训练
        watch_loss = []  # 保存损失
        for batch in range(int(train_sample / batch_size)):  # 分批训练
            x, y = build_dataset(batch_size, vocab, sentence_length)
            optim.zero_grad()  # 梯度下降初始化为0
            loss = model(x, y)  # 根据线性结果输入模型中
            loss.backward()  # 计算梯度 反向传播?
            optim.step()  # 更新权重
            watch_loss.append(loss.item())  # 获得损失项
        print("===\n第%d轮平均loss:%f" % (epoch + 1, np.mean(watch_loss)))
        acc = evaluate(model, vocab, sentence_length)
        log.append([acc, np.mean(watch_loss)])
    plt.plot(range(len(log)), [l[0] for l in log])
    plt.plot(range(len(log)), [l[1] for l in log])
    plt.show()
    # 保存模型
    torch.save(model.state_dict(), "model.pth")
    # 保存词表
    writer = open("vocab.json", "w", encoding="utf8")
    writer.write(json.dumps(vocab, ensure_ascii=False, indent=2))
    writer.close()
    return


# 最终预测
def predict(model_path, vocab_path, input_strings):
    char_dim = 20  # 每个字的维度
    sentence_length = 6  # 样本文本长度
    vocab = json.load(open(vocab_path, "r", encoding="utf8"))
    model = build_model(vocab, char_dim, sentence_length)  # 建立模型
    model.load_state_dict(torch.load(model_path))  # 加载训练好的权重
    x = []
    for input_string in input_strings:
        x.append([vocab[char] for char in input_string])  # 将输入序列化
    model.eval()  # 测试模式，不使用dropout
    with torch.no_grad():  # 不计算梯度
        result = model.forward(torch.LongTensor(x))  # 模型预测
    for i, input_string in enumerate(input_strings):
        print(round(float(result[i])), input_string, result[i])  # 打印结果


if __name__ == "__main__":
    main()
    test_strings = ["abvxee", "casdfg", "rqweqg", "nlkdww"]
    predict("model.pth", "vocab.json", test_strings)
