import itertools
import keras.layers as kl


def tito_embedding(conf):
    """tito model used in talkingdata"""
    input_layer = list()
    embedded_layer = list()
    embed_conf, num_conf = conf
    for name, input_dim, output_dim in embed_conf:
        x_in = kl.Input(shape=(1,), dtype="int32")
        embedded = kl.Embedding(input_dim, output_dim, name=name, input_length=1)(x_in)
        embedded_layer.append(embedded)
        input_layer.append(x_in)

    embedded_layer = kl.concatenate(embedded_layer)
    embedded_layer = kl.SpatialDropout1D(embedded_layer)
    embedded_layer = [kl.Flatten(embedded_layer)]

    dense_layer = list()
    for name, setting in num_conf.items():
        x_in = kl.Input(shape=(setting[1],), name=name, dtype="float32")
        dense_layer.append(x_in)
        input_layer.append(x_in)

    if len(dense_layer) > 1:
        dense_num = kl.concatenate(dense_layer)
    else:
        dense_num = dense_layer[0]
    dense_num = kl.Dense(128)(dense_num)
    dense_num = kl.BatchNormalization()(dense_num)
    embedded_layer.append(dense_num)

    return input_layer, embedded_layer


def simple_embedding(conf):
    """make embedded and input layer
    """
    input_layer = list()
    embedded_layer = list()
    embed_conf, num_conf = conf
    for name, input_dim, output_dim in embed_conf:
        x_in = kl.Input(shape=(1,), dtype="int32")
        embedded = kl.Embedding(input_dim, output_dim, name=name, input_length=1)(x_in)
        embedded = kl.Flatten()(embedded)
        embedded = kl.Dense(output_dim)(embedded)
        embedded = kl.advanced_activations.PReLU()(embedded)
        embedded_layer.append(embedded)
        input_layer.append(x_in)

    dense_layer = list()
    for name, setting in num_conf.items():
        x_in = kl.Input(shape=(setting[1],), name=name, dtype="float32")
        dense_layer.append(x_in)
        input_layer.append(x_in)

    if len(dense_layer) > 1:
        dense_num = kl.concatenate(dense_layer)
    else:
        dense_num = dense_layer[0]
    dense_num = kl.Dense(10)(dense_num)
    dense_num = kl.advanced_activations.PReLU()(dense_num)
    embedded_layer.append(dense_num)

    return input_layer, embedded_layer


def tosh_embedding(conf):
    """make embedded and input layer tosh version
    """
    input_layer = list()
    embedded_layer = list()
    embed_conf, num_conf = conf

    # embedding
    for name, input_dim, output_dim in embed_conf:
        x_in = kl.Input(shape=(1,), name=name, dtype="int32")
        embedded = kl.Embedding(input_dim, output_dim, input_length=1)(x_in)
        if output_dim > 10:
            embedded = kl.Dropout(0.1)(embedded)
        embedded = kl.Flatten()(embedded)
        embedded_layer.append(embedded)
        input_layer.append(x_in)

    x = kl.concatenate(embedded_layer)
    x = kl.Dense(64)(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    emb_out = kl.Dropout(0.2)(x)

    table_layers = [emb_out]
    for tab_name, setting in num_conf.items():
        x_in = kl.Input(shape=(setting[1],), name=tab_name, dtype="float32")

        dense_num = kl.Dense(setting[2])(x_in)
        dense_num = kl.advanced_activations.PReLU()(dense_num)
        dense_num = kl.BatchNormalization()(dense_num)
        dense_num = kl.Dropout(0.2)(dense_num)

        input_layer.append(x_in)
        table_layers.append(dense_num)

    return input_layer, table_layers


def embedding(conf, embedding_trainable=True):
    """make embedded and input layer
    """
    input_layer = list()
    out_layer = list()
    embedded_layer = list()
    embed_conf, num_conf = conf
    print(f"embedding train: {embedding_trainable}")
    for name, input_dim, output_dim in embed_conf:
        x_in = kl.Input(shape=(1,), name=name, dtype="int32")
        embedded = kl.Embedding(
            input_dim,
            output_dim,
            input_length=1,
            name=f"{name}_embedding",
            trainable=embedding_trainable
            # embeddings_initializer='zeros'
            # embeddings_regularizer=l2(0.001)
        )(x_in)
        if output_dim > 10 and embedding_trainable:
            embedded = kl.Dropout(0.1)(embedded)
        embedded = kl.Flatten()(embedded)
        embedded_layer.append(embedded)
        input_layer.append(x_in)

    x = kl.concatenate(embedded_layer)
    x = kl.Dense(128, init="he_normal")(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    # x = kl.Dropout(0.1)(x)
    emb_out = x

    dense_layer = list()
    for name, setting in num_conf.items():
        x_in = kl.Input(shape=(setting[1],), name=name, dtype="float32")
        dense_num = kl.Dense(setting[2], init="he_normal")(x_in)
        dense_num = kl.advanced_activations.PReLU()(dense_num)
        dense_num = kl.BatchNormalization()(dense_num)
        input_layer.append(x_in)
        dense_layer.append(dense_num)

    out_layer = dense_layer + [emb_out]

    return input_layer, out_layer


def embedding_for_ae(conf, out_size):
    """make embedded and input layer
    """
    input_layer = list()
    out_layer = list()
    embedded_layer = list()
    embed_conf, num_conf = conf
    for name, input_dim, output_dim in embed_conf:
        x_in = kl.Input(shape=(1,), name=name, dtype="int32")
        embedded = kl.Embedding(input_dim, output_dim, input_length=1)(x_in)
        embedded = kl.Flatten()(embedded)
        embedded_layer.append(embedded)
        input_layer.append(x_in)

    x = kl.concatenate(embedded_layer)
    x = kl.Dense(out_size // 2, init="he_normal")(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    # x = kl.Dropout(0.1)(x)
    emb_out = x

    dense_layer = list()
    for name, setting in num_conf.items():
        x_in = kl.Input(shape=(setting[1],), name=name, dtype="float32")
        dense_num = kl.Dense(out_size // 2, init="he_normal")(x_in)
        dense_num = kl.advanced_activations.PReLU()(dense_num)
        dense_num = kl.BatchNormalization()(dense_num)
        input_layer.append(x_in)
        dense_layer.append(dense_num)

    out_layer = dense_layer + [emb_out]

    return input_layer, out_layer


def embedding_v2(conf):
    """make embedded and input layer
    """
    input_layer = list()
    out_layer = list()
    embedded_layer = list()
    embed_conf, num_conf = conf
    for name, input_dim, output_dim in embed_conf:
        x_in = kl.Input(shape=(1,), name=name, dtype="int32")
        embedded = kl.Embedding(input_dim, output_dim, input_length=1)(x_in)
        embedded = kl.Flatten()(embedded)
        if name == "SmartScreen":
            ss_layer = embedded
            ss_size = output_dim
        else:
            embedded_layer.append(embedded)

        input_layer.append(x_in)

    fm_layer = list()
    for l in embedded_layer:
        x = kl.Dense(ss_size)(l)
        x = kl.merge.dot([x, ss_layer], axes=1)
        fm_layer.append(x)

    x = kl.concatenate(fm_layer)
    x = kl.Dense(128)(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    fm_out = x

    x = kl.concatenate(embedded_layer)
    x = kl.Dense(128)(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    emb_out = x

    dense_layer = list()
    for name, setting in num_conf.items():
        x_in = kl.Input(shape=(setting[1],), name=name, dtype="float32")
        dense_num = kl.Dense(setting[2])(x_in)
        dense_num = kl.advanced_activations.PReLU()(dense_num)
        dense_num = kl.BatchNormalization()(dense_num)
        input_layer.append(x_in)
        dense_layer.append(dense_num)

    out_layer = dense_layer + [emb_out, fm_out]

    return input_layer, out_layer


def fm_embedding(conf):
    """fm embedding model.
    embedding size must be equal.
    """
    input_layer = list()
    embedded_layer = list()
    embed_conf, num_conf = conf
    for name, input_dim, _ in embed_conf:
        x_in = kl.Input(shape=(1,), dtype="int32")
        embedded = kl.Embedding(input_dim, 6, name=name, input_length=1)(x_in)
        embedded_layer.append(kl.Flatten(embedded))
        input_layer.append(x_in)

    fm_layer = list()
    for emb1, emb2 in itertools.combinations(embedded_layer, 2):
        dot_layer = kl.merge([emb1, emb2], mode="dot", dot_axes=1)
        fm_layer.append(dot_layer)

    for i, (name, input_dim, _) in enumerate(embed_conf):
        linear = kl.Embedding(input_dim, 1, name=name, input_length=1)(input_layer[i])
        fm_layer.append(linear)

    dense_layer = list()
    for name, setting in num_conf.items():
        x_in = kl.Input(shape=(setting[1],), name=name, dtype="float32")
        dense_layer.append(x_in)
        input_layer.append(x_in)

    if len(dense_layer) > 1:
        dense_num = kl.concatenate(dense_layer)
    else:
        dense_num = dense_layer[0]
    dense_num = kl.Dense(128)(dense_num)
    dense_num = kl.BatchNormalization()(dense_num)
    fm_layer.append(dense_num)

    return input_layer, fm_layer
