import keras.layers as kl
from keras.models import Model
from keras.optimizers import Adam
from keras.losses import binary_crossentropy
from keras import backend as K
from .loss import Dense_with_AMsoftmax_loss
from .lovasz_loss import lovasz_loss
from .embedding import embedding, simple_embedding
from .embedding import tosh_embedding, fm_embedding


def embedded_mlp(embed_conf, embedding_trainable=True, am_softmax=False, lovasz=False):
    """Multi-layer Perceptron for embedded data
    """

    input_layer, embedded_layer = embedding(embed_conf, embedding_trainable)
    x = kl.concatenate(embedded_layer)
    x = kl.Dropout(0.2)(x)

    x = kl.Dense(64, init="he_normal")(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    x = kl.Dropout(0.2)(x)

    x = kl.Dense(16, init="he_normal")(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    x = kl.Dropout(0.2)(x)

    opt = Adam(amsgrad=True)
    if am_softmax:
        AM_softmax_logits = Dense_with_AMsoftmax_loss(2, m=0.2, scale=1)
        outputs = AM_softmax_logits(x)
        model = Model(input_layer, outputs, name="MLP")
        model.compile(optimizer=opt, loss=AM_softmax_logits.loss, metrics=["accuracy"])
    elif lovasz:
        outputs = kl.Dense(1, name="linear")(x)
        logit = kl.Dense(1, activation="sigmoid", name="logit")(x)
        model = Model(input_layer, [outputs, logit], name="MLP")
        import ipdb
        with ipdb.launch_ipdb_on_exception():
            model.compile(
                optimizer=opt,
                loss={"linear": lovasz_loss, 
                      "logit": "binary_crossentropy"},
                metrics=["accuracy", "binary_crossentropy"]
            )
    else:
        outputs = kl.Dense(1, activation="sigmoid", init="he_normal")(x)
        model = Model(input_layer, outputs, name="MLP")
        model.compile(optimizer=opt, loss="binary_crossentropy", metrics=["accuracy"])
    return model


def embedded_mlp_multi(embed_conf, out_size):
    """Multi-layer Perceptron for embedded data
    """

    input_layer, embedded_layer = embedding(embed_conf)
    x = kl.concatenate(embedded_layer)
    x = kl.Dropout(0.2)(x)

    x = kl.Dense(64, init="he_normal")(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    x = kl.Dropout(0.2)(x)

    x = kl.Dense(16, init="he_normal")(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    x = kl.Dropout(0.2)(x)

    outputs = kl.Dense(out_size, activation="softmax", init="he_normal")(x)
    model = Model(input_layer, outputs, name="MLP")
    opt = Adam(amsgrad=True)
    model.compile(optimizer=opt, loss="categorical_crossentropy", metrics=["accuracy"])
    return model


def deep_fm(fm_conf):
    """Multi-layer Perceptron for embedded data
    """

    input_layer, fm_layer = fm_embedding(fm_conf)
    x = kl.concatenate(fm_layer)
    x = kl.Dropout(0.2)(x)

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

    x = kl.Dense(16)(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    x = kl.Dropout(0.2)(x)

    outputs = kl.Dense(1, activation="sigmoid")(x)
    model = Model(input_layer, outputs, name="FM")
    opt = Adam(amsgrad=True)
    model.compile(optimizer=opt, loss="binary_crossentropy", metrics=["accuracy"])
    return model


def simple_embedded_mlp(embed_conf):
    """Multi-layer Perceptron for embedded data
    """

    input_layer, embedded_layer = simple_embedding(embed_conf)
    x = kl.concatenate(embedded_layer)
    x = kl.Dropout(0.5)(x)

    x = kl.Dense(256)(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    x = kl.Dropout(0.5)(x)

    x = kl.Dense(64)(x)
    x = kl.advanced_activations.PReLU()(x)
    x = kl.BatchNormalization()(x)
    x = kl.Dropout(0.5)(x)

    outputs = kl.Dense(1, activation="sigmoid")(x)
    model = Model(input_layer, outputs, name="MLP")
    model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
    return model


def tosh_nn(embed_conf, am_softmax=False, lovasz=False):
    """MLP based on tosh's NN."""
    input_layer, table_layer = tosh_embedding(embed_conf)

    # tab = kl.Add()(fm_layer)
    tab = kl.concatenate(table_layer)
    fc1 = kl.Dense(128)(tab)
    fc1 = kl.advanced_activations.PReLU()(fc1)
    fc1 = kl.BatchNormalization()(fc1)
    fc1 = kl.Dropout(0.25)(fc1)

    fc2 = kl.Dense(32)(fc1)
    fc2 = kl.advanced_activations.PReLU()(fc2)
    fc2 = kl.BatchNormalization()(fc2)
    fc2 = kl.Dropout(0.25)(fc2)

    opt = Adam(amsgrad=True)
    if am_softmax:
        AM_softmax_logits = Dense_with_AMsoftmax_loss(2, m=0.2, scale=1)
        outputs = AM_softmax_logits(fc2)
        model = Model(input_layer, outputs, name="MLP")
        model.compile(optimizer=opt, loss=AM_softmax_logits.loss, metrics=["accuracy"])
    elif lovasz:
        outputs = kl.Dense(1, name="linear")(fc2)
        logit = kl.Dense(1, activation="sigmoid", name="logit")(fc2)
        model = Model(input_layer, [outputs, logit], name="MLP")
        import ipdb
        with ipdb.launch_ipdb_on_exception():
            model.compile(
                optimizer=opt,
                loss={"linear": lovasz_loss, 
                      "logit": "binary_crossentropy"},
                metrics=["accuracy", "binary_crossentropy"]
            )
    else:
        outputs = kl.Dense(1, activation="sigmoid", init="he_normal")(fc2)
        model = Model(input_layer, outputs, name="MLP")
        model.compile(optimizer=opt, loss="binary_crossentropy", metrics=["accuracy"])
    return model


def embedding_multi_out(embed_conf):
    """ Multi-output embedding network
    """
    output = list()
    x_in, out_ls = embedding(embed_conf)
    for layer in out_ls:
        x = kl.Dense(16)(layer)
        x = kl.advanced_activations.PReLU()(x)
        x = kl.BatchNormalization()(x)
        x = kl.Dropout(0.2)(x)

        x = kl.Dense(1, activation="sigmoid")(x)
        output.append(x)

    all_out = kl.concatenate(out_ls)

    all_out = kl.Dense(256)(all_out)
    all_out = kl.advanced_activations.PReLU()(all_out)
    all_out = kl.BatchNormalization()(all_out)
    all_out = kl.Dropout(0.5)(all_out)

    all_out = kl.Dense(64)(all_out)
    all_out = kl.advanced_activations.PReLU()(all_out)
    all_out = kl.BatchNormalization()(all_out)
    all_out = kl.Dropout(0.5)(all_out)

    all_out = kl.Dense(1, activation="sigmoid")(all_out)
    output.append(all_out)

    model = Model(x_in, output, name="Multi-in Multi-out")
    model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
    return model


def sparse_mlp(input_shape):
    """sparse MLP
    """
    x_in = kl.Input(shape=(input_shape,))
    x = kl.Dense(1024, activation="relu")(x_in)
    x = kl.Dropout(0.5)(x)
    x = kl.Dense(512, activation="relu")(x)
    x = kl.Dropout(0.5)(x)
    x = kl.Dense(256, activation="relu")(x)
    x = kl.Dropout(0.5)(x)
    x = kl.Dense(64, activation="relu")(x)
    outputs = kl.Dense(1, activation="sigmoid")(x)
    model = Model(x_in, outputs, name="Sparse_MLP")
    model.compile(optimizer="rmsprop", loss="binary_crossentropy", metrics=["accuracy"])
    return model
