import numpy as np
import pandas as pd

from gensim.models import KeyedVectors
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.utils import to_categorical


class AgNewsData:
    MAX_NB_WORDS = 200000
    EMBEDDING_DIM = 300
    MAX_SEQUENCE_LENGTH = 1000

    @classmethod
    def load(cls):
        word2vec = KeyedVectors.load_word2vec_format('dataset/embedding/GoogleNews-vectors-negative300.bin', binary=True)

        column_names = ['category_id', 'title', 'description']
        with open("dataset/ag_news_csv/train.csv", "r") as file:
            df_train = pd.read_csv(file, names=column_names, header=None)
        with open("dataset/ag_news_csv/test.csv", "r") as file:
            df_test = pd.read_csv(file, names=column_names, header=None)

        X_texts = list(df_train.title + df_train.description)
        y_labels = list(df_train.category_id)


        val_X_texts = list(df_test.title + df_test.description)
        val_y_labels = list(df_test.category_id)


        tokenizer = Tokenizer(num_words=cls.MAX_NB_WORDS)
        tokenizer.fit_on_texts(X_texts + val_X_texts)
        X_sequences = tokenizer.texts_to_sequences(X_texts)
        val_X_sequences = tokenizer.texts_to_sequences(val_X_texts)

        word_index = tokenizer.word_index

        X = pad_sequences(X_sequences, maxlen=cls.MAX_SEQUENCE_LENGTH)
        val_X = pad_sequences(val_X_sequences, maxlen=cls.MAX_SEQUENCE_LENGTH)

        y_labels = to_categorical(np.asarray(y_labels))
        val_y_labels = to_categorical(np.asarray(val_y_labels))


import numpy as np
import pandas as pd

from gensim.models import KeyedVectors
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.utils import to_categorical

import logging
logger = logging.getLogger(__name__)


WORD2VEC_PATH = 'dataset/embedding/GoogleNews-vectors-negative300.bin'
AG_NEWS_TRAIN_PATH = 'dataset/ag_news_csv/train.csv'
AG_NEWS_TEST_PATH = 'dataset/ag_news_csv/test.csv'

class DataSet:
    @classmethod
    def load_word2vec(cls):
        try:
            word2vec = KeyedVectors.load_word2vec_format(WORD2VEC_PATH, binary=True)
        except Exception as e:
            raise e
        return word2vec


class AgNews(DataSet):
    MAX_NB_WORDS = 200000
    EMBEDDING_DIM = 300
    MAX_SEQUENCE_LENGTH = 1000

    def __init__(self):
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None
        self.df_train = None
        self.df_test = None
        self.word_index = None

    def load(self):
        column_names = ['category_id', 'title', 'description']
        with open(AG_NEWS_TRAIN_PATH, "r") as file:
            self.df_train = pd.read_csv(file, names=column_names, header=None)
        with open(AG_NEWS_TEST_PATH, "r") as file:
            self.df_test = pd.read_csv(file, names=column_names, header=None)

    def create_dataset(self):
        X_train_texts = list(self.df_train.title + self.df_train.description)
        y_train_labels = list(self.df_train.category_id)
        logger.info('train text size is '.format(len(X_train_texts)))
        logger.info('train label size is'.format(len(y_train_labels)))

        X_test_texts = list(self.df_test.title + self.df_test.description)
        y_test_labels = list(self.df_test.category_id)

        tokenizer = Tokenizer(num_words=AgNews.MAX_NB_WORDS)
        tokenizer.fit_on_texts(X_train_texts + X_test_texts)
        X_train_sequences = tokenizer.texts_to_sequences(X_train_texts)
        X_test_sequences = tokenizer.texts_to_sequences(X_test_texts)

        self.word_index = tokenizer.word_index

        self.X_train = pad_sequences(X_train_sequences, maxlen=AgNews.MAX_SEQUENCE_LENGTH)
        self.X_test = pad_sequences(X_test_sequences, maxlen=AgNews.MAX_SEQUENCE_LENGTH)

        self.y_train = to_categorical(np.asarray(y_train_labels))
        self.y_test = to_categorical(np.asarray(y_test_labels))

format data helper