from logging import getLogger

import numpy as np
import tensorflow as tf
import tensorflow.keras.layers as layers

logger = getLogger(__name__)


class FeatureAggregationSimilarityDataset(tf.keras.utils.Sequence):
    def __init__(self,
                 x_item_indices: np.ndarray,
                 y_item_indices: np.ndarray,
                 x_item_features: np.ndarray,
                 y_item_features: np.ndarray,
                 scores: np.ndarray,
                 batch_size: int = 2**10) -> None:
        self.x_item_indices = x_item_indices
        self.y_item_indices = y_item_indices
        self.x_item_features = x_item_features
        self.y_item_features = y_item_features
        self.scores = scores
        self.batch_size = batch_size

    def get(self, size: int, batch_size: int):
        idx = np.arange(self.scores.shape[0])
        np.random.shuffle(idx)
        idx = idx[:size]
        return FeatureAggregationSimilarityDataset(x_item_indices=self.x_item_indices.copy()[idx],
                                                   y_item_indices=self.y_item_indices.copy()[idx],
                                                   x_item_features=self.x_item_features.copy()[idx],
                                                   y_item_features=self.y_item_features.copy()[idx],
                                                   scores=self.scores.copy()[idx],
                                                   batch_size=batch_size)

    def train_test_split(self, test_size_rate: float):
        data_size = self.scores.shape[0]
        test_size = int(data_size * test_size_rate)
        train = FeatureAggregationSimilarityDataset(x_item_indices=self.x_item_indices.copy()[test_size:],
                                                    y_item_indices=self.y_item_indices.copy()[test_size:],
                                                    x_item_features=self.x_item_features.copy()[test_size:],
                                                    y_item_features=self.y_item_features.copy()[test_size:],
                                                    scores=self.scores.copy()[test_size:],
                                                    batch_size=self.batch_size)
        test = FeatureAggregationSimilarityDataset(x_item_indices=self.x_item_indices.copy()[:test_size],
                                                   y_item_indices=self.y_item_indices.copy()[:test_size],
                                                   x_item_features=self.x_item_features.copy()[:test_size],
                                                   y_item_features=self.y_item_features.copy()[:test_size],
                                                   scores=self.scores.copy()[:test_size],
                                                   batch_size=self.batch_size)
        return train, test

    def __getitem__(self, index):
        start_idx = index * self.batch_size
        last_idx = start_idx + self.batch_size
        x_item_indices = self.x_item_indices[start_idx:last_idx]
        y_item_indices = self.y_item_indices[start_idx:last_idx]
        x_item_features = self.x_item_features[start_idx:last_idx]
        y_item_features = self.y_item_features[start_idx:last_idx]
        scores = self.scores[start_idx:last_idx]
        return (x_item_indices, y_item_indices, x_item_features, y_item_features), scores

    def __len__(self):
        return int(np.floor(len(self.x_item_indices) / self.batch_size))


class Average(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(Average, self).__init__(**kwargs)

    def build(self, input_shape):
        super(Average, self).build(input_shape)

    def call(self, inputs, mask, **kwargs):
        mask = tf.cast(mask, tf.float32)
        return tf.div_no_nan(tf.keras.backend.batch_dot(inputs, mask, axes=1), tf.reduce_sum(mask, axis=1, keepdims=True))

    def compute_mask(self, inputs, mask=None):
        return None

    def get_config(self):
        base_config = super(Average, self).get_config()
        return base_config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


class Clip(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(Clip, self).__init__(**kwargs)

    def build(self, input_shape):
        super(Clip, self).build(input_shape)

    def call(self, inputs, **kwargs):
        return tf.keras.backend.clip(inputs, -1.0, 1.0)

    def compute_mask(self, inputs, mask=None):
        return None

    def get_config(self):
        base_config = super(Clip, self).get_config()
        return base_config

    @classmethod
    def from_config(cls, config):
        return cls(**config)


class FeatureAggregationSimilarityGraph(object):
    def __init__(self,
                 feature_size: int,
                 embedding_size: int,
                 item_size: int,
                 max_feature_index: int,
                 embeddings_initializer=None,
                 bias_embeddings_initializer=None,
                 embeddings_regularizer=None):
        embeddings_initializer = embeddings_initializer or tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.005)
        embeddings_regularizer = embeddings_regularizer or tf.keras.regularizers.l2(0.0001)
        bias_embeddings_initializer = bias_embeddings_initializer or tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.005)
        self.input_x_index = layers.Input(shape=(1, ), name='input_x_index')
        self.input_y_index = layers.Input(shape=(1, ), name='input_y_index')
        self.input_x_feature = layers.Input(shape=(feature_size, ), name='input_x_feature')
        self.input_y_feature = layers.Input(shape=(feature_size, ), name='input_y_feature')

        self.embedding = layers.Embedding(
            max_feature_index + 1,
            embedding_size,
            mask_zero=True,
            embeddings_initializer=embeddings_initializer,
            embeddings_regularizer=embeddings_regularizer,
        )
        self.bias_embedding = tf.keras.layers.Embedding(
            item_size + 1,
            1,
            mask_zero=True,
            embeddings_initializer=bias_embeddings_initializer,
        )

        self.embedding_x = self.average(self.embedding(self.input_x_feature))
        self.embedding_y = self.average(self.embedding(self.input_y_feature))
        self.bias_x = self.average(self.bias_embedding(self.input_x_index))
        self.bias_y = self.average(self.bias_embedding(self.input_y_index))

        self.inner_prod = tf.keras.layers.dot([self.embedding_x, self.embedding_y], axes=1, normalize=True)
        self.similarity = tf.keras.layers.add([self.inner_prod, self.bias_x, self.bias_y])
        self.similarity = self.clip(self.similarity)

    @staticmethod
    def average(x):
        return Average()(x)

    @staticmethod
    def clip(x):
        return Clip()(x)


class FeatureAggregationSimilarityModel(object):
    def __init__(
        self,
        embedding_size: int,
        learning_rate: float,
        feature_size: int,
        item_size: int,
        max_feature_index: int,
    ) -> None:
        self.feature_size = feature_size
        graph = FeatureAggregationSimilarityGraph(feature_size=feature_size,
                                                  embedding_size=embedding_size,
                                                  item_size=item_size,
                                                  max_feature_index=max_feature_index)
        self.model = tf.keras.models.Model(inputs=[graph.input_x_index, graph.input_y_index, graph.input_x_feature, graph.input_y_feature],
                                           outputs=graph.similarity)
        self.embeddings = tf.keras.models.Model(inputs=[graph.input_x_feature], outputs=graph.embedding_x)
        self.model.compile(optimizer=tf.train.AdamOptimizer(learning_rate), loss=tf.keras.losses.mse, metrics=[tf.keras.metrics.mse])

    def __getstate__(self):
        return self.feature_size, self.model.to_json(), self.model.get_weights(), self.embeddings.to_json(), self.embeddings.get_weights()

    def __setstate__(self, state):
        feature_size, json_config, weights, embedding_json_config, embedding_weights = state
        self.feature_size = feature_size
        self.model = tf.keras.models.model_from_json(json_config, custom_objects={'Clip': Clip, 'Average': Average})
        self.model.set_weights(weights)
        self.embeddings = tf.keras.models.model_from_json(embedding_json_config, custom_objects={'Clip': Clip, 'Average': Average})
        self.embeddings.set_weights(embedding_weights)

    def fit(self, dataset: FeatureAggregationSimilarityDataset, epoch_size: int, test_size_rate: float = 0.05, early_stopping_patience: int = 2):
        logger.info('prepare data...')
        callbacks = [tf.keras.callbacks.EarlyStopping(patience=early_stopping_patience)]
        logger.info('start to fit...')
        data, validation_data = self._make_dataset(dataset=dataset, test_size_rate=test_size_rate)
        self.model.fit_generator(data, epochs=epoch_size, callbacks=callbacks, validation_data=validation_data)

    def calculate_similarity(self, x_item_indices, y_item_indices, x_item_features, y_item_features, batch_size=2**14):
        return self.model.predict(x=(x_item_indices, y_item_indices, x_item_features, y_item_features), batch_size=batch_size).reshape(-1)

    def calculate_embeddings(self, item_features, batch_size=2**14):
        return self.embeddings.predict(x=(item_features, ), batch_size=batch_size)

    def _make_dataset(self, dataset: FeatureAggregationSimilarityDataset, test_size_rate: float):
        data, validation_data = dataset.train_test_split(test_size_rate=test_size_rate)
        return data, validation_data
