import torch
from torchtext.vocab import Vectors
import torch.nn.functional as F
from tqdm import tqdm
from pprint import pprint

from .model import Config, TextCNN, LSTMClassifier
from .utils.learning import adjust_learning_rate
from .utils.dataset import get_dataset, get_iterator, get_vectors, get_vocab
from .utils.evaluate import get_score
from .utils.log import logger
from .utils.pad import pad_sequnce
from .config import DEVICE, DEFAULT_CONFIG

seed = 2019
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

class SA(object):
    """
    """
    def __init__(self):
        self._model = None
        self._word_vocab = None
        self._label_vocab = None
    
    def train(self, train_path, save_path=DEFAULT_CONFIG['save_path'], dev_path=None, vectors_path=None, **kwargs):
        train_dataset = get_dataset(train_path)
        if dev_path:
            dev_dataset = get_dataset(dev_path)
            word_vocab, label_vocab = get_vocab(train_dataset, dev_dataset)
        else:
            word_vocab, label_vocab = get_vocab(train_dataset)
        self._word_vocab = word_vocab
        self._label_vocab = label_vocab
        print(self._word_vocab.itos[:5])
        print(self._label_vocab.stoi)
        train_iter = get_iterator(train_dataset, batch_size=DEFAULT_CONFIG['batch_size'])
        config = Config(word_vocab, label_vocab, save_path=save_path, **kwargs)
        textcnn = TextCNN(config)
        print(textcnn)
        self._model = textcnn
        optim = torch.optim.Adam(textcnn.parameters(), lr=config.lr)
        for epoch in range(config.epoch):
            textcnn.train()
            acc_loss = 0
            for fuck in tqdm(train_iter):
                optim.zero_grad()
                # print('text', fuck.text)
                # print('text shape', fuck.text.size())
                logits = self._model(fuck.text.t())
                # print('predict', logits)
                # print('true', fuck.label)
                # break
                item_loss = F.cross_entropy(logits, fuck.label)
                acc_loss += item_loss.item()
                item_loss.backward()
                optim.step()
            logger.info('epoch: {}, acc_loss: {}'.format(epoch, acc_loss))
            if dev_path:
                # for dev_score in self._validate(dev_dataset):
                #     logger.info('dev score:{}'.format(dev_score))
                dev_score = self._validate(dev_dataset)
                logger.info('dev score:{}'.format(dev_score))

            adjust_learning_rate(optim, config.lr / (1 + (epoch + 1) * config.lr_decay))
        config.save()
        textcnn.save()

    
    def predict(self, text):
        text = pad_sequnce(text, 5)
        vec_text = torch.tensor([self._word_vocab.stoi[x] for x in text])
        vec_text = vec_text.reshape(1, -1).to(DEVICE)
        # print('text', vec_text, 'shape', vec_text.size())
        vec_predict = self._model(vec_text)[0]
        # print('predict', vec_predict)
        soft_predict = torch.softmax(vec_predict, dim=0)
        predict_prob, predict_index = torch.max(soft_predict.cpu().data, dim=0)
        predict_class = self._label_vocab.itos[predict_index]
        predict_prob = predict_prob.item()
        return predict_prob, predict_class

    
    def load(self, save_path=DEFAULT_CONFIG['save_path']):
        config = Config.load(save_path)
        textcnn = TextCNN(config)
        textcnn.load()
        self._model = textcnn
        self._word_vocab = config.word_vocab
        self._label_vocab = config.label_vocab
    
    def test(self, test_path):
        test_dataset = get_dataset(test_path)
        test_score = self._validate(test_dataset)
        logger.info('test score:{}'.format(test_score))
    
    def _validate(self, dev_dataset, batch_size=10):
        self._model.eval()
        dev_predict_list = []
        dev_true_list = []
        for dev_item in tqdm(dev_dataset):
            _, predict_class = self.predict(dev_item.text)
            dev_true_list.append(dev_item.label)
            dev_predict_list.append(predict_class)
        return get_score(dev_predict_list, dev_true_list)
        