"""
## 1. Requirements
"""
!pip install torchtext
"""
## 2. Load Dataset
"""
class DataLoader(object):
    def __init__(self,
                 max_vocab=9999,
                 min_freq=1,
                 init_token='<bos>',
                 eos_token='<eos>'):
        super().__init__()

        self.max_vocab = max_vocab
        self.min_freq = min_freq
        self.label = data.Field(sequential=False,
                                unk_token=None)
        self.text = data.Field(init_token=init_token,
                               eos_token=eos_token,
                               batch_first=True,
                               )

    def get_loaders(self, config):
        train, valid = data.TabularDataset(
            path=config.file_path,
            format='tsv',
            fields=[
                ('label', self.label),
                ('text', self.text),
            ],
        ).split(split_ratio=config.train_ratio)

        self.train_loader, self.valid_loader = data.BucketIterator.splits(
            (train, valid),
            batch_size=config.batch_size,
            device= 'cuda:{}'.format(config.gpu_id) if config.gpu_id >= 0 else 'cpu',
            shuffle=True,
            sort_key = lambda x: len(x.text),
            sort_within_batch = True
        )

        self.label.build_vocab(train)
        self.text.build_vocab(train, max_size=self.max_vocab, min_freq=self.min_freq)

        return self.train_loader, self.valid_loader
"""
## 3. Modeling
"""
class RNNclassifier(nn.Module):
    def __init__(self, input_size, emb_dim, hidden_size, n_layers, n_classes, dropout):
        self.input_size = input_size
        self.emb_dim = emb_dim
        self.hidden_size = hidden_size
        self.n_layers = n_layers
        self.dropout = dropout
        self.n_classes = n_classes

        super().__init__()

        # x = (bs, length)
        self.emb = nn.Embedding(input_size, emb_dim)
        # x = (bs, length, dim)
        self.rnn = nn.LSTM(
            input_size=emb_dim,
            hidden_size=hidden_size,
            dropout=dropout,
            num_layers=n_layers,
            batch_first=True,
            bidirectional=True
        )
        # x = (bs, length, hidden * 2)
        self.generator = nn.Linear(in_features=hidden_size * 2,
                                   out_features=n_classes)
        self.activation = nn.LogSoftmax(dim=-1)

    def forward(self, x):
        # x = (bs, length)
        x1 = self.emb(x)
        # x = (bs, length, dim)
        x2, _ = self.rnn(x1)
        # x = (bs, length, hidden * 2)
        y = self.activation(self.generator(x2[:, -1]))

        return y

class CNNclassifier(nn.Module):
    def __init__(self,
                 input_size,
                 emb_dim,
                 window_sizes,
                 n_filters,
                 use_padding,
                 use_batchnorm,
                 dropout,
                 n_classes):
        # emb
        self.input_size = input_size
        self.emb_dim = emb_dim
        # cnn module
        self.window_sizes = window_sizes
        self.n_filters = n_filters
        self.use_padding = use_padding
        self.use_batchnorm = use_batchnorm
        self.dropout = dropout
        # generator
        self.n_classes = n_classes
        super().__init__()

        # x = (bs, length)
        self.emb = nn.Embedding(input_size, emb_dim)
        # x = (bs, length, emb)
        self.cnnmodule = nn.ModuleList()

        for window_size in window_sizes:
            self.cnnmodule.append(nn.Sequential(
                    nn.Conv2d(in_channels=1,
                              out_channels=n_filters,
                              kernel_size=(window_size, emb_dim)),
                              # padding=(1, 1) if use_padding else 0),
                    nn.ReLU(),
                    nn.BatchNorm2d(n_filters) if use_batchnorm else nn.Dropout(dropout),
                )
            )
        # x = (bs, n_filters, length -  window_size + 1, 1)
        self.generator = nn.Linear(in_features=n_filters * len(window_sizes),
                                   out_features=n_classes)
        self.activation = nn.LogSoftmax(dim=-1)


    def forward(self, x):
        # x = (bs, length)
        x = self.emb(x)
        # x = (bs, length, emb)
        min_length = max(self.window_sizes)
        if min_length > x.size(1):
            pad = x.new(x.size(0), min_length - x.size(1), self.emb_dim).zero_()
            x = torch.cat([x, pad], dim=1)

        x = x.unsqueeze(dim=1)
        # x = (bs, 1, length, emb)

        cnn_outs = []
        for block in self.cnnmodule:
            cnn_out = block(x)
            # x = (bs, n_filters, length - window_size + 1, 1)
            cnn_out = nn.functional.max_pool1d(input=cnn_out.squeeze(dim=-1),
                                               kernel_size=cnn_out.size(-2)).squeeze(dim=-1)
            # x = (bs, n_filters)
            cnn_outs += [cnn_out]

        x = torch.cat(cnn_outs, dim=-1)
        # x = (bs, sum(n_filters))
        x = self.generator(x)
        y = self.activation(x)

        return y
"""
## 4. Train loop
"""
SILNET = 0
BATCHWISE = 1
EPOCHWISE = 2


class KerasLikeEngine(Engine) :
    def __init__(self, func, model, optimizer, loss, config) :
        self.model = model
        self.optimizer = optimizer
        self.loss = loss
        self.config = config

        super().__init__(func)

        self.best_loss = np.inf
        self.best_model = None

        self.device = next(model.parameters()).device

    @staticmethod
    def train(engine, minibatch) :
        engine.model.train()

        x, y = minibatch.text, minibatch.label
        x, y = x.to(engine.device), y.to(engine.device)

        y_hat = engine.model(x)
        loss_i = engine.loss(y_hat.squeeze(), y)

        ### ---- train only---- ###
        engine.optimizer.zero_grad()
        loss_i.backward()
        engine.optimizer.step()
        ### ---- train only---- ###

        if isinstance(y, torch.LongTensor) or isinstance(y, torch.cuda.LongTensor) :
            accuracy = (torch.argmax(y_hat, dim=-1) == y).sum() / float(y.size(0))
        else :
            accuracy = 0

        return {
            'loss' : float(loss_i),
            'accuracy' : float(accuracy),
        }

    @staticmethod
    def valid(engine, minibatch) :
        engine.model.eval()

        with torch.no_grad() :
            x, y = minibatch.text, minibatch.label
            x, y = x.to(engine.device), y.to(engine.device)

            y_hat = engine.model(x)
            loss_i = engine.loss(y_hat.squeeze(), y)

            if isinstance(y, torch.LongTensor) or isinstance(y, torch.cuda.LongTensor) :
                accuracy = (torch.argmax(y_hat, dim=-1) == y).sum() / float(y.size(0))
            else :
                accuracy = 0

            return {
                'loss' : float(loss_i),
                'accuracy' : float(accuracy),
            }

    @staticmethod
    def attach(train_engine, val_engine, verbose=EPOCHWISE) :
        def attach_running_average(engine, metric_name) :
            RunningAverage(output_transform=lambda x : x[metric_name]).attach(
                engine,
                metric_name,
            )

        training_metric_names = ['loss', 'accuracy']

        for metric_name in training_metric_names :
            attach_running_average(train_engine, metric_name)

        if verbose == BATCHWISE :
            pbar = ProgressBar(bar_format=None, ncols=120)
            pbar.attach(train_engine, ['loss', 'accuracy'])

        if verbose == EPOCHWISE :
            pbar = ProgressBar(bar_format=None, ncols=120)
            pbar.attach(train_engine, ['loss', 'accuracy'])

            @train_engine.on(Events.EPOCH_COMPLETED)
            def print_logs(engine) :
                print('Epoch {} Train - Accuracy: {:.4f} Loss: {:.4f}'.format(
                    engine.state.epoch,
                    engine.state.metrics['accuracy'],
                    engine.state.metrics['loss'],
                ))

        validation_metric_names = ['loss', 'accuracy']

        for metric_name in validation_metric_names :
            attach_running_average(val_engine, metric_name)

        if verbose == BATCHWISE :
            pbar = ProgressBar(bar_format=None, ncols=120)
            pbar.attach(val_engine, ['loss', 'accuracy'])

        if verbose == EPOCHWISE :
            pbar = ProgressBar(bar_format=None, ncols=120)
            pbar.attach(val_engine, ['loss', 'accuracy'])

            # @val_engine.on(Events.EPOCH_COMPLETED)
            # def print_logs(engine) :
            #     print('        Valid - Accuracy: {:.4f} Loss: {:.4f} Lowest_loss={:.4f}'.format(
            #         engine.state.metrics['accuracy'],
            #         engine.state.metrics['loss'],
            #         engine.best_loss,
            #     ))

    @staticmethod
    def check_best(engine) :
        loss = float(engine.state.metrics['loss'])
        if loss <= engine.best_loss :
            engine.best_loss = loss
            engine.best_model = deepcopy(engine.model.state_dict())

    @staticmethod
    def save_model(engine, config, **kwargs) :
        torch.save({
            'model' : engine.best_model,
            'config' : config,
            **kwargs
        }, config.model_fn)

    @staticmethod
    def print_logs(engine) :
        print('        Valid - Accuracy: {:.4f} Loss: {:.4f} Lowest_loss={:.4f}'.format(
            engine.state.metrics['accuracy'],
            engine.state.metrics['loss'],
            engine.best_loss,
        ))


class Trainer :
    def __init__(self, config) :
        self.config = config

    def train(self, model, optimizer, loss, train_loader, val_loader) :
        # Make Engine object
        train_engine = KerasLikeEngine(KerasLikeEngine.train, model, optimizer, loss, self.config)
        val_engine = KerasLikeEngine(KerasLikeEngine.valid, model, optimizer, loss, self.config)

        # Attach Metrics
        KerasLikeEngine.attach(train_engine, val_engine, verbose=self.config.verbose)

        # Event Handling
        def run_validation(engine, validation_engine, valid_loader) :
            validation_engine.run(valid_loader, max_epochs=1)

        train_engine.add_event_handler(Events.EPOCH_COMPLETED,
                                       run_validation, val_engine, val_loader)
        val_engine.add_event_handler(Events.EPOCH_COMPLETED,
                                     KerasLikeEngine.check_best)
        val_engine.add_event_handler(Events.EPOCH_COMPLETED,
                                     KerasLikeEngine.save_model, train_engine, self.config)
        val_engine.add_event_handler(Events.EPOCH_COMPLETED,
                                     KerasLikeEngine.print_logs, val_engine)


        # running
        train_engine.run(train_loader, max_epochs=self.config.n_epochs)

        model.load_state_dict(val_engine.best_model)

        return model
"""
## 5. Run on CLI
"""
import argparse

import torch
import torch.nn as nn
import torch.optim as optim

from model import RNNclassifier, CNNclassifier
from trainer import Trainer
from data_loader import DataLoader


def define_argparse():

    p.add_argument('--model_fn', required=True)
    p.add_argument('--model_name', type=str, required=True)
    p.add_argument('--file_path', type=str, required=True)

    p.add_argument('--gpu_id', type=int, default= 0 if torch.cuda.is_available() else -1)
    p.add_argument('--batch_size', type=int, default=64)
    p.add_argument('--n_epochs', type=int, default=5)
    p.add_argument('--train_ratio', type=float, default=.8)
    p.add_argument('--verbose', type=int, default=2)

    p.add_argument('--emb_dim', type=int, default=32)
    p.add_argument('--dropout', type=float, default=.2)
    # rnn
    p.add_argument('--hidden_size', type=int, default=32)
    p.add_argument('--n_layers', type=int, default=3)
    # cnn
    p.add_argument('--window_sizes', type=list, default=[2, 3, 4])
    p.add_argument('--n_filters', type=int, default=32)
    p.add_argument('--use_padding', type=bool, default=True)
    p.add_argument('--use_batchnorm', type=bool, default=True)

    config = p.parse_args()

    return config


def main(config) :
    dataloader = DataLoader()
    train_loader, valid_loader = dataloader.get_loaders(config)

    print('Train size : {} Valid size : {}'.format(
        len(train_loader.dataset),
        len(valid_loader.dataset),
    ))

    input_size = len(dataloader.text.vocab)
    n_classes = len(dataloader.label.vocab)

    if config.model_name == 'rnn':
        model = RNNclassifier(input_size, config.emb_dim, config.hidden_size, config.n_layers, n_classes, config.dropout)
    elif config.model_name == 'cnn':
        model = CNNclassifier(input_size, config.emb_dim, config.window_sizes, config.n_filters, config.use_padding, config.use_batchnorm,
                              config.dropout, n_classes)
    loss = nn.NLLLoss() # expect 1d tensor
    optimizer = optim.Adam(model.parameters())

    trainer = Trainer(config)
    trainer.train(model, optimizer, loss, train_loader, valid_loader)

if __name__ == '__main__':
    config = define_argparse()
    main(config)