import torch
from torch import nn, optim
import torch.nn.functional as F
from torchtext import data
from torch.utils.data import DataLoader
from torchtext.data import Iterator
from gensim.models import KeyedVectors


class BucketIteratorWrapper(DataLoader):
    __initialized__ = False

    def __init__(self, iterator: Iterator):
        self.batch_size = iterator.batch_size
        self.num_workers = 1
        self.collate_fn = None
        self.pin_memory = False
        self.drop_last = False
        self.timeout = 0
        self.worker_init_fn = None
        self.sampler = iterator
        self.batch_sampler = iterator
        self.__initialized__ = True

    def __iter__(self):
        return map(lambda batch: {
            'features': batch.TEXT,
            'targets': batch.LABEL,
        }, self.batch_sampler.__iter__())

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


class CNN(nn.Module):

    def __init__(self, output_dim, kernel_num, kernel_sizes=[3, 4, 5], dropout=0.5, static=False):
        super(CNN, self).__init__()

        model = KeyedVectors.load_word2vec_format('ch07/GoogleNews-vectors-negative300.bin', binary=True)
        weights = torch.FloatTensor(model.vectors)
        self.embed = nn.Embedding.from_pretrained(weights)
        self.convs1 = nn.ModuleList([nn.Conv2d(1, kernel_num, (k, self.embed.weight.shape[1])) for k in kernel_sizes])
        self.dropout = nn.Dropout(dropout)
        self.fc1 = nn.Linear(len(kernel_sizes) * kernel_num, output_dim)
        self.static = static

    def conv_and_pool(self, x, conv):
        x = F.relu(conv(x)).squeeze(3)
        x = F.max_pool1d(x, x.size(2)).squeeze(2)
        return x

    def forward(self, x):
        x = self.embed(x)

        if self.static:
            x = x.detach()

        x = x.unsqueeze(1)
        x = x.float()
        x = [F.relu(conv(x)).squeeze(3) for conv in self.convs1]

        x = [F.max_pool1d(i, i.size(2)).squeeze(2) for i in x]

        x = torch.cat(x, 1)
        x = self.dropout(x)
        logit = self.fc1(x)
        return logit


TEXT = data.Field(sequential=True, lower=True, batch_first=True)
LABELS = data.Field(sequential=False, batch_first=True, use_vocab=False)

train, val, test = data.TabularDataset.splits(
    path='ch06', train='train2.txt',
    validation='valid2.txt', test='test2.txt', format='tsv',
    fields=[('TEXT', TEXT), ('LABEL', LABELS)])

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_iter, val_iter, test_iter = data.BucketIterator.splits(
    (train, val, test), batch_sizes=(64, 64, 64), device=device, repeat=False, sort=False)

train_loader = BucketIteratorWrapper(train_iter)
valid_loader = BucketIteratorWrapper(val_iter)
loaders = {"train": train_loader, "valid": valid_loader}

TEXT.build_vocab(train, min_freq=2)
LABELS.build_vocab(train)
model = CNN(output_dim=4, kernel_num=3, kernel_sizes=[3, 4, 5], dropout=0.2)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

for epoch in range(1):
    model.train()
    for batch in train_iter:
        x, y = batch.TEXT, batch.LABEL
        y_pred = model(x)
        print(y_pred)
        print(y_pred.shape)
pythonを用いて、ID番号で表現された単語列$$\boldsymbol{x} = (x_1, x_2, \dots, x_T)$$がある。ただし、$$T$$は単語列の長さ、$$x_t \in \mathbb{R}^{V}$$は単語のID番号のone-hot表記である（$$V$$は単語の総数である）。畳み込みニューラルネットワーク（CNN: Convolutional Neural Network）を用い、単語列$$\boldsymbol{x}$$からカテゴリ$$y$$を予測するモデルを実装せよ。

ただし、畳み込みニューラルネットワークの構成は以下の通りとする。

+ 単語埋め込みの次元数: $$d_w$$
+ 畳み込みのフィルターのサイズ: 3 トークン
+ 畳み込みのストライド: 1 トークン
+ 畳み込みのパディング: あり
+ 畳み込み演算後の各時刻のベクトルの次元数: $$d_h$$
+ 畳み込み演算後に最大値プーリング（max pooling）を適用し、入力文を$$d_h$$次元の隠れベクトルで表現

すなわち、時刻$$t$$の特徴ベクトル$$p_t \in \mathbb{R}^{d_h}$$は次式で表される。

$$
p_t = g(W^{(px)} [\mathrm{emb}(x_{t-1}); \mathrm{emb}(x_t); \mathrm{emb}(x_{t+1})] + b^{(p)})
$$

ただし、$$W^{(px)} \in \mathbb{R}^{d_h \times 3d_w}, b^{(p)} \in \mathbb{R}^{d_h}$$はCNNのパラメータ、$$g$$は活性化関数（例えば$$\tanh$$やReLUなど）、$$[a; b; c]$$はベクトル$$a, b, c$$の連結である。なお、行列$$W^{(px)}$$の列数が$$3d_w$$になるのは、3個のトークンの単語埋め込みを連結したものに対して、線形変換を行うためである。

最大値プーリングでは、特徴ベクトルの次元毎に全時刻における最大値を取り、入力文書の特徴ベクトル$$c \in \mathbb{R}^{d_h}$$を求める。$$c[i]$$でベクトル$$c$$の$$i$$番目の次元の値を表すことにすると、最大値プーリングは次式で表される。

$$
c[i] = \max_{1 \leq t \leq T} p_t[i]
$$

最後に、入力文書の特徴ベクトル$$c$$に行列$$W^{(yc)} \in \mathbb{R}^{L \times d_h}$$とバイアス項$$b^{(y)} \in \mathbb{R}^{L}$$による線形変換とソフトマックス関数を適用し、カテゴリ$$y$$を予測する。

$$
y = {\rm softmax}(W^{(yc)} c + b^{(y)})
$$

なお、この問題ではモデルの学習を行わず、ランダムに初期化された重み行列で$$y$$を計算するだけでよい。