import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class SkipGram(nn.Module):
    def __init__(self, vocab_size, embd_size):
        super(SkipGram, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embd_size)

    def forward(self, focus, context):
        embed_focus = self.embeddings(focus)
        embed_ctx = self.embeddings(context)
        # score = torch.mm(embed_focus, torch.t(embed_ctx))
        score = torch.mul(embed_focus, embed_ctx).sum(dim=1)
        log_probs = score #F.logsigmoid(score)

        return log_probs

    def loss(self, log_probs, target):
        loss_fn = nn.BCEWithLogitsLoss()
        # loss_fn = nn.NLLLoss()
        loss = loss_fn(log_probs, target)
        return loss


class CBOW(nn.Module):
    def __init__(self, vocab_size, embd_size, context_size, hidden_size):
        super(CBOW, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embd_size)
        self.linear1 = nn.Linear(2 * context_size * embd_size, hidden_size)
        self.linear2 = nn.Linear(hidden_size, vocab_size)

    def forward(self, inputs):
        embedded = self.embeddings(inputs).view((1, -1))
        hid = F.relu(self.linear1(embedded))
        out = self.linear2(hid)
        log_probs = F.log_softmax(out)
        return log_probs

import random
import re

import torch
import torch.optim as optim
from tqdm import tqdm
from pytorch_word2vec_model import SkipGram

epochs = 50
negative_sampling = 4
window = 2
vocab_size = 1
embd_size = 300


def batch_data(x, batch_size=128):
    in_w = []
    out_w = []
    target = []
    for text in x:
        for i in range(window, len(text) - window):
            word_set = set()
            in_w.append(text[i])
            in_w.append(text[i])
            in_w.append(text[i])
            in_w.append(text[i])

            out_w.append(text[i - 2])
            out_w.append(text[i - 1])
            out_w.append(text[i + 1])
            out_w.append(text[i + 2])

            target.append(1)
            target.append(1)
            target.append(1)
            target.append(1)
            # negative sampling
            count = 0
            while count < negative_sampling:
                rand_id = random.randint(0, vocab_size-1)
                if not rand_id in word_set:
                    in_w.append(text[i])
                    out_w.append(rand_id)
                    target.append(0)
                    count += 1

            if len(out_w) >= batch_size:
                yield [in_w, out_w, target]
                in_w = []
                out_w = []
                target = []
    if out_w:
        yield [in_w, out_w, target]


def train(train_text_id, model,opt):
    model.train()  # 启用dropout和batch normalization
    ave_loss = 0
    pbar = tqdm()
    cnt=0
    for x_batch in batch_data(train_text_id):
        in_w, out_w, target = x_batch
        in_w_var = torch.tensor(in_w)
        out_w_var = torch.tensor(out_w)
        target_var = torch.tensor(target,dtype=torch.float)

        model.zero_grad()
        log_probs = model(in_w_var, out_w_var)
        loss = model.loss(log_probs, target_var)
        loss.backward()
        opt.step()
        ave_loss += loss.item()
        pbar.update(1)
        cnt += 1
        pbar.set_description('< loss: %.5f >' % (ave_loss / cnt))
    pbar.close()
text_id = []
vocab_dict = {}

with open('corpus.txt',encoding='utf-8') as fp:
    for line in fp:
        lines = re.sub("[^A-Za-z0-9']+", ' ', line).lower().split()
        line_id = []
        for s in lines:
            if not s:
                continue
            if s not in vocab_dict:
                vocab_dict[s] = len(vocab_dict)
            id = vocab_dict[s]
            line_id.append(id)
            if id==11500:
                print(id,s)
        text_id.append(line_id)
vocab_size = len(vocab_dict)
print('vocab_size', vocab_size)
model = SkipGram(vocab_size, embd_size)

for epoch in range(epochs):
    print('epoch', epoch)
    opt = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),
    lr=0.001, weight_decay=0)
    train(text_id, model,opt)