# Copyright 2020 The ElasticDL Authors. All rights reserved.
# 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 tensorflow as tf
from tensorflow import feature_column as fc

from elasticdl.python.elasticdl.callbacks import LearningRateScheduler
from elasticdl_preprocessing.feature_column import feature_column as edl_fc
from model_zoo.census_model_sqlflow.wide_and_deep.feature_configs import (
    FEATURE_TRANSFORM_INFO_EXECUTE_ARRAY,
    INPUT_SCHEMAS,
    TRANSFORM_OUTPUTS,
    age_bucketize,
    capital_gain_bucketize,
    capital_loss_bucketize,
    education_hash,
    group1_embedding_deep,
    group1_embedding_wide,
    group2_embedding_deep,
    group2_embedding_wide,
    group3_embedding_deep,
    hours_per_week_bucketize,
    marital_status_lookup,
    native_country_hash,
    occupation_hash,
    race_lookup,
    relationship_lookup,
    sex_lookup,
    workclass_lookup,
)
from model_zoo.census_model_sqlflow.wide_and_deep.transform_ops import (
    TransformOpType,
)


# The model definition from model zoo. It's functional style.
# Input Params:
#   input_layers: The input layers dict of feature inputs
#   wide_embedding: A tensor. Embedding for the wide part.
#   deep_embedding: A tensor. Embedding for the deep part.
def wide_and_deep_classifier(input_layers, wide_embedding, deep_embedding):
    # Wide Part
    wide = wide_embedding  # shape = (None, 3)

    # Deep Part
    dnn_input = deep_embedding
    for i in [16, 8, 4]:
        dnn_input = tf.keras.layers.Dense(i)(dnn_input)

    # Output Part
    concat_input = tf.concat([wide, dnn_input], 1)

    logits = tf.reduce_sum(concat_input, 1, keepdims=True)
    probs = tf.reshape(tf.sigmoid(logits), shape=(-1,))

    return tf.keras.Model(
        inputs=input_layers,
        outputs={"logits": logits, "probs": probs},
        name="wide_deep",
    )


# Build the input layers from the schema of the input features
def get_input_layers(input_schemas):
    input_layers = {}

    for schema_info in input_schemas:
        input_layers[schema_info.name] = tf.keras.layers.Input(
            name=schema_info.name, shape=(1,), dtype=schema_info.dtype
        )

    return input_layers


# Build the transform logic from the metadata in feature_configs.py.
def transform(inputs):
    feature_column_dict = {}

    for feature_transform_info in FEATURE_TRANSFORM_INFO_EXECUTE_ARRAY:
        if feature_transform_info.op_type == TransformOpType.HASH:
            feature_column_dict[
                feature_transform_info.output
            ] = tf.feature_column.categorical_column_with_hash_bucket(
                feature_transform_info.input,
                hash_bucket_size=feature_transform_info.hash_bucket_size,
            )
        elif feature_transform_info.op_type == TransformOpType.BUCKETIZE:
            feature_column_dict[
                feature_transform_info.output
            ] = tf.feature_column.bucketized_column(
                fc.numeric_column(feature_transform_info.input),
                boundaries=feature_transform_info.boundaries,
            )
        elif feature_transform_info.op_type == TransformOpType.LOOKUP:
            feature_column_dict[
                feature_transform_info.output
            ] = tf.feature_column.categorical_column_with_vocabulary_list(
                feature_transform_info.input,
                vocabulary_list=workclass_lookup.vocabulary_list,
            )
        elif feature_transform_info.op_type == TransformOpType.CONCAT:
            concat_inputs = [
                feature_column_dict[name]
                for name in feature_transform_info.input
            ]
            concat_column = edl_fc.concatenated_categorical_column(
                concat_inputs
            )
            feature_column_dict[feature_transform_info.output] = concat_column
        elif feature_transform_info.op_type == TransformOpType.EMBEDDING:
            feature_column_dict[
                feature_transform_info.output
            ] = tf.feature_column.embedding_column(
                feature_column_dict[feature_transform_info.input],
                dimension=feature_transform_info.output_dim,
            )
        elif feature_transform_info.op_type == TransformOpType.ARRAY:
            feature_column_dict[feature_transform_info.output] = [
                feature_column_dict[name]
                for name in feature_transform_info.input
            ]

    return tuple(
        [
            tf.keras.layers.DenseFeatures(feature_column_dict[name])(inputs)
            for name in TRANSFORM_OUTPUTS
        ]
    )


# The following code has the same logic with the `transform` function above.
# It can be generated from the parsed meta in feature_configs using code_gen.
def transform_from_code_gen(source_inputs):
    education_hash_fc = fc.categorical_column_with_hash_bucket(
        "education", hash_bucket_size=education_hash.hash_bucket_size
    )

    occupation_hash_fc = fc.categorical_column_with_hash_bucket(
        "occupation", hash_bucket_size=occupation_hash.hash_bucket_size
    )

    native_country_hash_fc = fc.categorical_column_with_hash_bucket(
        "native_country", hash_bucket_size=native_country_hash.hash_bucket_size
    )

    workclass_lookup_fc = fc.categorical_column_with_vocabulary_list(
        "workclass", vocabulary_list=workclass_lookup.vocabulary_list
    )

    marital_status_lookup_fc = fc.categorical_column_with_vocabulary_list(
        "marital_status", vocabulary_list=marital_status_lookup.vocabulary_list
    )

    relationship_lookup_fc = fc.categorical_column_with_vocabulary_list(
        "relationship", vocabulary_list=relationship_lookup.vocabulary_list
    )

    race_lookup_fc = fc.categorical_column_with_vocabulary_list(
        "race", vocabulary_list=race_lookup.vocabulary_list
    )

    sex_lookup_fc = fc.categorical_column_with_vocabulary_list(
        "sex", vocabulary_list=sex_lookup.vocabulary_list
    )

    age_bucketize_fc = fc.bucketized_column(
        fc.numeric_column("age"), boundaries=age_bucketize.boundaries
    )

    capital_gain_bucketize_fc = fc.bucketized_column(
        fc.numeric_column("capital_gain"),
        boundaries=capital_gain_bucketize.boundaries,
    )

    capital_loss_bucketize_fc = fc.bucketized_column(
        fc.numeric_column("capital_loss"),
        boundaries=capital_loss_bucketize.boundaries,
    )

    hours_per_week_bucketize_fc = fc.bucketized_column(
        fc.numeric_column("hours_per_week"),
        boundaries=hours_per_week_bucketize.boundaries,
    )

    group1_fc = edl_fc.concatenated_categorical_column(
        categorical_columns=[
            workclass_lookup_fc,
            hours_per_week_bucketize_fc,
            capital_gain_bucketize_fc,
            capital_loss_bucketize_fc,
        ]
    )

    group2_fc = edl_fc.concatenated_categorical_column(
        categorical_columns=[
            education_hash_fc,
            marital_status_lookup_fc,
            relationship_lookup_fc,
            occupation_hash_fc,
        ]
    )

    group3_fc = edl_fc.concatenated_categorical_column(
        categorical_columns=[
            age_bucketize_fc,
            sex_lookup_fc,
            race_lookup_fc,
            native_country_hash_fc,
        ]
    )

    group1_wide_embedding_fc = fc.embedding_column(
        group1_fc, dimension=group1_embedding_wide.output_dim,
    )

    group2_wide_embedding_fc = fc.embedding_column(
        group2_fc, dimension=group2_embedding_wide.output_dim,
    )

    group1_deep_embedding_fc = fc.embedding_column(
        group1_fc, dimension=group1_embedding_deep.output_dim,
    )

    group2_deep_embedding_fc = fc.embedding_column(
        group2_fc, dimension=group2_embedding_deep.output_dim,
    )

    group3_deep_embedding_fc = fc.embedding_column(
        group3_fc, dimension=group3_embedding_deep.output_dim,
    )

    wide_feature_columns = [
        group1_wide_embedding_fc,
        group2_wide_embedding_fc,
    ]

    deep_feature_columns = [
        group1_deep_embedding_fc,
        group2_deep_embedding_fc,
        group3_deep_embedding_fc,
    ]

    return (
        tf.keras.layers.DenseFeatures(wide_feature_columns)(source_inputs),
        tf.keras.layers.DenseFeatures(deep_feature_columns)(source_inputs),
    )


# The entry point of the submitter program
def custom_model():
    input_layers = get_input_layers(input_schemas=INPUT_SCHEMAS)
    wide_embedding, deep_embedding = transform_from_code_gen(input_layers)

    return wide_and_deep_classifier(
        input_layers, wide_embedding, deep_embedding
    )


def loss(labels, predictions):
    logits = predictions["logits"]
    return tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(
            labels=tf.cast(tf.reshape(labels, (-1, 1)), tf.float32),
            logits=logits,
        )
    )


def optimizer(lr=0.001):
    return tf.keras.optimizers.Adam(learning_rate=lr)


def eval_metrics_fn():
    return {
        "logits": {
            "accuracy": lambda labels, predictions: tf.equal(
                tf.cast(tf.reshape(predictions, [-1]) > 0.5, tf.int32),
                tf.cast(tf.reshape(labels, [-1]), tf.int32),
            )
        },
        "probs": {"auc": tf.keras.metrics.AUC()},
    }


def callbacks():
    def _schedule(model_version):
        if model_version < 5000:
            return 0.0003
        elif model_version < 12000:
            return 0.0002
        else:
            return 0.0001

    return [LearningRateScheduler(_schedule)]


if __name__ == "__main__":
    model = custom_model()
    print(model.summary())

    output = model.call(
        {
            "education": tf.constant([["Bachelors"]], tf.string),
            "occupation": tf.constant([["Tech-support"]], tf.string),
            "native_country": tf.constant([["United-States"]], tf.string),
            "workclass": tf.constant([["Private"]], tf.string),
            "marital_status": tf.constant([["Separated"]], tf.string),
            "relationship": tf.constant([["Husband"]], tf.string),
            "race": tf.constant([["White"]], tf.string),
            "sex": tf.constant([["Female"]], tf.string),
            "age": tf.constant([[18]], tf.float32),
            "capital_gain": tf.constant([[100.0]], tf.float32),
            "capital_loss": tf.constant([[1.0]], tf.float32),
            "hours_per_week": tf.constant([[40]], tf.float32),
        }
    )

    print(output)
