import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
import bisect
import copy
import random
import shelve
import shutil
import numpy as np
from tqdm import tqdm
from keras.models import Sequential, load_model
from keras.layers import Dense, LSTM
from keras.utils import Sequence
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.preprocessing import StandardScaler, MaxAbsScaler

from keras.backend.tensorflow_backend import set_session
import tensorflow as tf
config = tf.ConfigProto(intra_op_parallelism_threads=1, inter_op_parallelism_threads=1)  # to allow a3c running more workers
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))


class InputSequenceSharedData:
    def __init__(self, *, frac_train):
        print('InputSequenceSharedData: reading data..')
        with np.load('data/features.npz') as data:
            self.data_user_ids = data['user_id']
            self.num_rows = int(data['num_rows'])
            self.num_users = int(data['num_users'])
        with np.load('data/data.npz') as data:
            self.data_x, self.data_y = data['x'], data['y']

        print('InputSequenceSharedData: scaling features...')
        self.scalar = MaxAbsScaler().fit(self.data_x)
        self.data_x = self.scalar.transform(self.data_x)

        print('InputSequenceSharedData: building data structures...')
        self.num_train = int(self.num_rows * frac_train)
        self.train_user_rows = [[] for _ in range(self.num_users)]
        self.test_user_rows = [[] for _ in range(self.num_users)]
        for row in range(self.num_train):
            user_id = self.data_user_ids[row]
            self.train_user_rows[user_id].append(row)
        for row in range(self.num_train, self.num_rows):
            user_id = self.data_user_ids[row]
            self.test_user_rows[user_id].append(row)

        self.train_user_offsets, self.test_user_offsets = [0], [0]
        for user_id in range(self.num_users):
            self.train_user_offsets.append(self.train_user_offsets[-1] + len(self.train_user_rows[user_id]))
            self.test_user_offsets.append(self.test_user_offsets[-1] + len(self.test_user_rows[user_id]))
        self.train_user_offsets = self.train_user_offsets[1:]
        self.test_user_offsets = self.test_user_offsets[1:]

        print('InputSequenceSharedData: init done')


class InputSequence(Sequence):
    def __init__(self, data, *, is_train, batch_size, lookback):
        self._data = data
        self._is_train = is_train
        self._batch_size = batch_size
        self._lookback = lookback
        self._total_rows = self._data.num_train if self._is_train else self._data.num_rows - self._data.num_train
        self._random_state = np.random.RandomState(1234)
        self.on_epoch_end()

    @property
    def dim_features(self):
        return self._data.data_x.shape[1] + 1

    def on_epoch_end(self):
        self._perm = self._random_state.permutation(self._total_rows)  # shuffle

    def __len__(self):
        return int(np.ceil(self._total_rows / self._batch_size))

    def __getitem__(self, batch_idx):
        start_idx = batch_idx * self._batch_size
        end_idx = min((batch_idx+1) * self._batch_size, self._total_rows)
        ret_rows = end_idx - start_idx
        batch_x = np.zeros((ret_rows, self._lookback+1, self.dim_features), dtype=np.float32)
        batch_y = np.empty(ret_rows, dtype=np.int8)
        user_offsets = self._data.train_user_offsets if self._is_train else self._data.test_user_offsets
        user_rows = self._data.train_user_rows if self._is_train else self._data.test_user_rows

        def get_row(user_row_offset, t):
            if user_row_offset-t >= 0:
                row = user_rows[user_id][user_row_offset-t]
                return row
            if self._is_train:
                train_rows = self._data.train_user_rows[user_id]
                row_offset = len(train_rows) + (user_row_offset - t)
                if row_offset >= 0:
                    row = train_rows[row_offset]
                    return row
            return None  # otherwise (the very begining of a user), just fill 0

        for i, idx in enumerate(self._perm[start_idx:end_idx]):
            user_id = bisect.bisect_right(user_offsets, idx)
            user_row_offset = idx - user_offsets[user_id-1] if user_id != 0 else idx
            row = get_row(user_row_offset, self._lookback+1)
            for t in range(self._lookback, -1, -1):
                # label of previous time step
                if row is not None:
                    batch_x[i, self._lookback-t, -1] = self._data.data_y[row]
                # features of current time step
                row = get_row(user_row_offset, t)
                if row is not None:
                    batch_x[i, self._lookback-t, :-1] = self._data.data_x[row]
            # label of current time step
            batch_y[i] = self._data.data_y[row]
        return batch_x, batch_y


def run_train_test():
    batch_size = 256
    lookback = 5

    print('preparing data...')
    data = InputSequenceSharedData(frac_train=0.8)

    print('fetching test set...')
    train_input_seq = InputSequence(data, is_train=True, batch_size=batch_size, lookback=lookback)
    test_input_seq = InputSequence(data, is_train=False, batch_size=batch_size, lookback=lookback)
    x_test, y_test = [], []
    for i in range(len(test_input_seq)):
        batch_x, batch_y = test_input_seq[i]
        x_test.append(batch_x)
        y_test.append(batch_y)
    x_test = np.concatenate(x_test)
    y_test = np.concatenate(y_test)
    print(x_test.shape, y_test.shape)

    # keras: build graph
    print('building graph...')
    model = Sequential()
    model.add(LSTM(10, input_shape=(lookback+1, train_input_seq.dim_features)))
    model.add(Dense(1, activation='sigmoid'))
    model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
    model.summary()

    # keras: train model
    print('training...')
    best_auc = 0
    for epoch in range(1, 10):
        model.fit_generator(train_input_seq,
            epochs=1, verbose=1,
            use_multiprocessing=True, workers=2,
            shuffle=False,  # InputSequence handles shuffling
        )
        filename = 'data/rnn/model-epoch-{}.h5'.format(epoch)
        model.save(filename)

        # validate
        y_pred_proba = model.predict(x_test).squeeze()
        y_pred = y_pred_proba.copy()
        y_pred[y_pred > 0.5] = 1
        y_pred[y_pred <= 0.5] = 0
        acc = accuracy_score(y_test, y_pred)
        auc = roc_auc_score(y_test, y_pred_proba)
        print('epoch', epoch, 'acc:', acc)
        print('epoch', epoch, 'auc:', auc)
        if auc > best_auc:
            best_auc = auc
            shutil.copy(filename, 'data/rnn/model.h5')

    # validate
    print('validating...')
    model = load_model('data/rnn/model.h5')
    y_pred_proba = model.predict(x_test).squeeze()
    y_pred = y_pred_proba.copy()
    y_pred[y_pred > 0.5] = 1
    y_pred[y_pred <= 0.5] = 0
    print('acc:', accuracy_score(y_test, y_pred))
    print('auc:', roc_auc_score(y_test, y_pred_proba))


if __name__ == '__main__':
    run_train_test()

