import numpy as np

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import backend as K


from disttf import GaussianRealisation,BiasedRealisation,BoxRealisation
from disttf import MixtureLayer, MotioLayer, ScaleLayer, SplitLayer
from disttf import NonLinearityLayer, SeperateLayer, RecombineLayer


def gen_model(inputs, splits=10, realisation="gauss", mixture=0, nonlin=False):

    i=keras.layers.Input(shape=(inputs,))
    q=i
    v=tf.constant(1.0,dtype=tf.float32)

    if splits>1:
        q,v=SplitLayer(splits)([q,v])
    q,v=MotioLayer()([q,v])
    q,v=ScaleLayer()([q,v])
    for j in range(mixture):
        q=MixtureLayer()(q)
        if nonlin:
            q=NonLinearityLayer()(q)
    if realisation=="gauss":
        q=GaussianRealisation()([q,v])
    elif realisation=="biased":
        q=BiasedRealisation()([q,v])
    elif realisation=="box":
        q=BoxRealisation()([q,v])
    else:
        raise ValueError("Unknown realisation type: "+realisation)

    model=keras.Model(inputs=i,outputs=q)

    loss=K.log(K.abs(q)+1e-6)
    loss=-K.mean(loss)
    model.add_loss(loss)

    return model