# Copyright 2019 IBM Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from spacy.lang.en import English
import numpy as np
from lale.operators import make_operator
import lale.helpers
from sklearn.utils.validation import _is_arraylike
import urllib.request
import os
import zipfile

logging.basicConfig(level=logging.INFO)


class GloveEmbeddingEncoderImpl(object):
    """
    GloveEmbeddingEncoderImpl is a module allows simple generation of sentence embeddings using
    glove word embeddings

    Parameters
    ----------
    dim: int, (default=300), dimension of word embeddings to use
    combiner: string, (default=mean), apply mean pooling or max pooling on word embeddings to generate
        sentence embedding

    References
    ----------
    R. JeffreyPennington and C. Manning. Glove: Global vectors for word representation. 2014

    """
    def __init__(self, dim=300, combiner='mean'):

        self.combiner = combiner
        self.dim = dim

        if self.combiner not in ["mean", "max"]:
            raise ValueError("Combiner must be either mean or max")

        if self.dim not in [50, 100, 200, 300]:
            raise ValueError("dim must be in 50, 100, 200, or 300")

        glove_path = os.path.join(os.path.dirname(__file__),
                                  "resources",
                                  "glove.6B.{}d.txt".format(dim))

        if not os.path.exists(glove_path):
            download_url = 'http://nlp.stanford.edu/data/glove.6B.zip'
            download_dir = os.path.join(os.path.dirname(__file__), 'resources', 'glove.6B.zip')
            if not os.path.exists(os.path.join(os.path.dirname(__file__), "resources")):
                os.mkdir(os.path.join(os.path.dirname(__file__), "resources"))
            urllib.request.urlretrieve(download_url, download_dir)
            with zipfile.ZipFile(download_dir, 'r') as zip_ref:
                zip_ref.extractall(os.path.join(os.path.dirname(__file__), "resources"))
            os.remove(download_dir)

        self.glove_dict = self._load_glove(glove_path)
        nlp = English()
        self.tokenizer = nlp.Defaults.create_tokenizer(nlp)
        print('finish initialization')

    def fit(self, X, y=None):
        return self

    def _load_glove(self, glove_path):
        glove_dict = dict()
        with open(glove_path, 'r', encoding='utf-8') as file_handle:
            for line in file_handle:
                line = line.split()
                glove_dict[line[0]] = np.array(line[1:], dtype=np.float32)
        return glove_dict

    def transform(self, X):
        if not _is_arraylike(X):
            raise TypeError("X is not iterable")

        transformed_X = list()
        for text in X:
            temp_vec = list()
            for token in self.tokenizer(text.lower()):
                temp_vec.append(self.glove_dict.get(token, self.glove_dict['unk']))

            if self.combiner == 'mean':
                sentence_vec = np.mean(temp_vec, axis=0)
            else:
                sentence_vec = np.amax(temp_vec, axis=0)

            transformed_X.append(sentence_vec)
        return transformed_X

_input_schema_fit = {
    '$schema': 'http://json-schema.org/draft-04/schema#',
    'description': 'Input data schema for training.',
    'type': 'object',
    'required': ['X'],
    'additionalProperties': False,
    'properties': {
        'X': {
            'description': 'Input Text',
            'anyOf': [
                {'type': 'array',
                 'items': {'type': 'string'}},
                {'type': 'array',
                 'items': {
                     'type': 'array', 'minItems': 1, 'maxItems': 1,
                     'items': {'type': 'string'}}}]}
    }}

_input_schema_predict = {
    '$schema': 'http://json-schema.org/draft-04/schema#',
    'description': 'Input data schema for training.',
    'type': 'object',
    'required': ['X'],
    'additionalProperties': False,
    'properties': {
        'X': {
            'description': 'Input Text',
            'anyOf': [
                {'type': 'array',
                 'items': {'type': 'string'}},
                {'type': 'array',
                 'items': {
                     'type': 'array', 'minItems': 1, 'maxItems': 1,
                     'items': {'type': 'string'}}}]}
    }}

_output_schema = {
    '$schema': 'http://json-schema.org/draft-04/schema#',
    'description': 'Output data schema for transformed data.',
    'type': 'array',
    'items': {'type': 'array', 'items': {'type': 'number'}}}

_hyperparams_schema = {
    '$schema': 'http://json-schema.org/draft-04/schema#',
    'description': 'Hyperparameter schema.',
    'allOf': [
        {'description':
             'This first sub-object lists all constructor arguments with their '
             'types, one at a time, omitting cross-argument constraints.',
         'type': 'object',
         'additionalProperties': False,
         'relevantToOptimizer': [],
         'properties': {
             'combiner': {
                'enum': ['mean', 'max'],
                'default': 'mean'},
                'dim':{'enum': [50, 100, 200, 300],
                       'default': 300}
         }}]}

_combined_schemas = {
    '$schema': 'http://json-schema.org/draft-04/schema#',
    'description': 'Combined schema for expected data and hyperparameters for a transformer for'
                   ' a text data transformer based on pre-trained glove embedding',
    'type': 'object',
    'tags': {
        'pre': ['text'],
        'op': ['transformer', '~interpretable'],
        'post': ['embedding']},
    'properties': {
        'input_fit': _input_schema_fit,
        'input_predict': _input_schema_predict,
        'output': _output_schema,
        'hyperparams': _hyperparams_schema}}


if __name__ == "__main__":
    lale.helpers.validate_is_schema(_combined_schemas)

GloveEmbeddingEncoder = make_operator(GloveEmbeddingEncoderImpl(), _combined_schemas)