import jieba

text = "深度学习(英语: deep learning)是机器学习的分支， 是一种以人工神经网络为架构， 对数据进行表征学习的算法。"

cuted = jieba.lcut(text)

print(cuted)

print([cuted[i:i+2] for i in range(len(cuted) - 1)])

torch.nn.Embedding(num_embeddings, embedding_dim)

embedding = nn.Embedding(vocab_size, 300) #实例化
input_embeded = embedding(input) # 进行embedding的操作

def tokenize(content):
    # filters = '!"#$%^&'
    # $, . 符号是个特殊字符，需要转义
    filters = ['\$','\t', '\n', '\x97', '\x96', '#',  '%', '&',  '\.', "<.*?>", '"']
    # | 在正则里面表示或
    content = re.sub("|".join(filters), " ", content)
    content = re.sub("！", " ！", content)
    tokens = [i.strip().lower for i in content.split()]
    return tokens



class ImdbDataset(Dataset):
    def __init__(self, train=True):
        self.train_data_path = r"./datasets/aclImdb/train"
        self.test_data_path = r"./datasets/aclImdb/test"

        data_path = self.train_data_path if train else self.test_data_path

        temp_data_path = [os.path.join(data_path, "pos"), os.path.join(data_path, "neg")]
        self.total_file_path = []
        for path in temp_data_path:
            file_path_list = os.listdir(path)
            file_path_list = [os.path.join(path, i) for i in file_path_list]
            self.total_file_path.extend(file_path_list)


    def __getitem__(self, index):
        file_path = self.total_file_path[index]

        # 获取label
        label_str = file_path.split("/")[-2]
        label = 0 if label_str == "neg" else 1

        # 获取内容
        content = open(file_path).read()
        tokens = tokenize(content)
        return tokens, label


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


def collate_fn(batch):
    """
    param batch: ([tokens, label], [tokens, label])
    """
    # print(batch)
    content, label = list(zip(*batch))
    return content, label


def get_dataloader(train=True):
    imdb_dataset = ImdbDataset(train)
    data_loader = DataLoader(
        imdb_dataset,
        batch_size=2,
        shuffle=True,
        collate_fn=collate_fn)

    return data_loader

class Word2Sequence():
    UNK_TAG = "UNK"
    PAD_TAG = "PAD"

    UNK = 0
    PAD = 1

    def __init__(self):
        self.dict = {
            self.UNK_TAG: self.UNK,
            self.PAD_TAG: self.PAD
        }

        self.count = {}

    def fit(self, sentence):
        """
        把单个句子保存到dict中
            sentence： [word1, word2, word3, ...]
        """
        for word in sentence:
            self.count[word] = self.count.get(word, 0) + 1

    def build_vocab(self, min=5, max=None, max_features=None):
        """
        生成词典
            min: 最小出现的次数
            max: 最大的次数
            max_features： 一共保留多少个词语
        """
        # 删除count中词频小于min的word
        if min is not None:
            self.count = {word:value for word, value in self.count.items() if value>min}

        # 删除词语大于max的值
        if max is not None:
            self.count = {word:value for word, value in self.count.items() if value<max}

        # 限制保留的词语数
        # items： [(key1, value1), (key2, value2), ...]
        # key = lambda x:x[-1]， 表示按照value进行排序
        if max_features is not None:
            temp = sorted(self.count.items(), key =lambda x: x[-1], reverse=True)[:max_features]
            self.count = dict(temp)

        # 遍历self.count里的key
        for word in self.count:
            # 每添加一个次到dict中后， 就用现在dict的长度来表示这个词
            self.dict[word] = len(self.dict)

        # 得到一个翻转的dict字典

        self.inverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def transform(self, sentence, max_len=max_len):
        """
        把句子转化为序列
            sentence: [word1, word2, ...]
            max_len: int， 对句子进行填充或者裁剪
        """

        if max_len is not None:
            if max_len > len(sentence):
                # 填充
                sentence = sentence + [self.PAD_TAG] * (max_len - len(sentence))
                print(sentence)
            if max_len < len(sentence):
                # 裁剪
                sentence = sentence[:max_len]
            
        return [self.dict.get(word, self.UNK) for word in sentence]


    def inverse_transform(self, indices):
        """
        把序列转化为句子
            indices: [1, 2, 3, 4, ...]
        """
        return [self.inverse_dict.get(idx) for idx in indices]

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

import torch
import torch.nn as nn
import torch.nn.functional as F

batch_size = 10
seq_len = 20
vocab_size = 100
embedding_dim = 30

hidden_size = 18
num_layer = 1

input = torch.randint(low=0, high=100, size=[batch_size, seq_len])

embedding = nn.Embedding(vocab_size, embedding_dim)
input_embeded = embedding(input)

lstm = nn.LSTM(input_size=embedding_dim, hidden_size=hidden_size, num_layers=num_layer, batch_first=True)

output, (h_n, c_n) = lstm(input_embeded)

print(output.size())
print("*" * 100)
print(h_n.size())
print("*" * 100)
print(c_n.size())

class IMDBLSTMModel(nn.Module):
    def __init__(self):
        super(IMDBLSTMModel, self).__init__()
        self.embedding = nn.Embedding(len(ws), 100)

        self.lstm = nn.LSTM(input_size=100, hidden_size=lib.hidden_size, num_layers=lib.num_layers, batch_first=True, bidirectional=lib.bidirectional, dropout=lib.dropout)

        self.fc = nn.Linear(lib.hidden_size * 2, 2)


    def forward(self, x):
        # 进行embedding操作， 形状： [batch_szie, max_len, 100]
        x = self.embedding(x)
        # x: [batch_size, max_len, 2 * hidden_size]， h_n, c_n: [2 * 2, batch_size, hidden_size]
        x, (h_n, c_n) = self.lstm(x)
        # 获取两个方向最后一次的output进行concat操作
        output_fw = h_n[-2, :, :] # 正向最后一次输出
        output_bw = h_n[-1, :, :] # 反向最后一次输出
        output = torch.cat([output_fw, output_bw], dim=-1) # [batch_size, hidden_size * 2]

        out = self.fc(output)
        return F.log_softmax(out, dim=-1)



model = IMDBLSTMModel().to(lib.device)
optimizer = torch.optim.Adam(model.parameters(), 0.001)

if os.path.exists("./saved_models/model.pkl"):
    model.load_state_dict(torch.load("./saved_models/model.pkl"))
    optimizer.load_state_dict(torch.load("./saved_models/optimizer.pkl"))

def train(epoch):
    for idx, (input, target) in enumerate(get_dataloader(train=True)):
        input = input.to(lib.device)
        target = target.to(lib.device)

        output= model(input)

        loss = F.nll_loss(output, target)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(epoch, idx, loss.item())

        if idx % 100 == 0:
            torch.save(model.state_dict(), "./saved_models/model.pkl")
            torch.save(optimizer.state_dict(), "./saved_models/optimizer.pkl")


def eval():
    loss_list = []
    acc_list = []
    for input, target in tqdm(get_dataloader(train=False), total=len(get_dataloader(train=False)), ascii=True, desc="test: "):
        input = input.to(lib.device)
        target = target.to(lib.device)
        with torch.no_grad():
            output = model(input)
            cur_loss = F.nll_loss(output, target)
            loss_list.append(cur_loss.cpu().item())

            pred = output.max(dim=-1)[-1]
            cur_acc = pred.eq(target).float().mean()
            acc_list.append(cur_acc.cpu().item())

    print("total loss, acc:  ", np.mean(loss_list), np.mean(acc_list))

if __name__ == "__main__":
    # for i in range(10):
    #     train(i)

    eval()