import tensorflow.keras as tf
import numpy as np
from sklearn.preprocessing import StandardScaler

verbose = 0


def predict_series(values, r1_nodes=5, r2_nodes=0, fc1_nodes=0, steps=20, use_lstm=True, *args, **kwargs):
    
    train = np.array(values)
    
    train_last_value = train[-1]
    train = train[1:] - train[:-1]
    sc = StandardScaler()
    train = sc.fit_transform(train.reshape(-1, 1))

    T = 25
    X = []
    Y = []
    for t in range(len(train) - T):
        x = train[t:t + T]
        X.append(x)
        Y.append(train[t + T])

    X = np.array(X).reshape(-1, T, 1)
    Y = np.array(Y)

    nb_stats = 0
    """
    X_temp = np.zeros(X.size + nb_stats * len(X)).reshape(-1, T + nb_stats)

    step_size = 1 / (len(X) + steps)


    def update_stats(row):
        new_stat = row[T:]
        new_stat[0] += step_size  # number of sample

        minimum = min(row[:T])  # minimum value, and when it occurred
        if minimum < row[T + 1]:
            new_stat[1], new_stat[2] = minimum, new_stat[0]

        maximum = max(row[:T])  # maximum value, and when it occurred
        if maximum > row[T + 3]:
            new_stat[3], new_stat[4] = maximum, new_stat[0]
        
        new_stat[5] = (row[T + 5] * row[T] + row[T - 1]) / (new_stat[0])  # rolling average

        difference10 = row[T - 1] - row[T - 11]  # the biggest difference within 10 items
        if difference10 > row[T + 6]:
            new_stat[6], new_stat[7] = difference10, new_stat[0]
        if difference10 < row[T + 8]:
            new_stat[8], new_stat[9] = difference10, new_stat[0]

        abs_difference10 = abs(difference10)  # the biggest absolute difference within 10 items
        if abs_difference10 > row[T + 10]:
            new_stat[10], new_stat[11] = abs_difference10, new_stat[0]
        if abs_difference10 < row[T + 12]:
            new_stat[12], new_stat[13] = abs_difference10, new_stat[0]
        
        return new_stat

    X_temp[0] = X[0]  #np.append(X[0])#, [0, np.inf, 0, -np.inf, 0])  #, 0, -np.inf, 0, +np.inf, 0, 0, 0, np.inf, 0])
    for i in range(1, len(X)):
        X_temp[i] = np.append(X[i][:T], X_temp[i - 1][T:])
        X_temp[i][T:] = update_stats(X_temp[i])
    """
    #X = X_temp[1:].reshape(-1, T + nb_stats, 1)
    #Y = Y[1:]

    i = tf.layers.Input(shape=(T + nb_stats, 1))
    
    if use_lstm:
        rnn_layer = tf.layers.LSTM
    else:
        rnn_layer = tf.layers.GRU

    if r2_nodes:
        x = rnn_layer(r1_nodes, return_sequences=True)(i)
        x = rnn_layer(r2_nodes)(x)
    else:
        x = rnn_layer(r1_nodes)(i)
    if fc1_nodes:
        x = tf.layers.Dense(fc1_nodes, activation='relu')(x)
    x = tf.layers.Dense(1)(x)
    model = tf.models.Model(i, x)
    
    
    """lr_schedule = tf.optimizers.schedules.ExponentialDecay(
        initial_learning_rate=0.2,
        decay_steps=10,
        decay_rate=0.8)
    optimizer = tf.optimizers.Ftrl(learning_rate=0.001, learning_rate_power=-0.1)"""
    #for i in range(0, 500, 10):
        #print('{}: {}'.format(i, lr_schedule(i)))
    
    
    model.compile(
        loss='mse', #tf.losses.LogCosh(),
        optimizer=tf.optimizers.Adamax(lr=0.1) #LogCosh()'sgd'
    )

    callbacks = [tf.callbacks.EarlyStopping(patience=150, monitor='loss', restore_best_weights=True)]

    r = model.fit(
        X, Y,
        epochs=500,
        callbacks=callbacks,
        verbose=verbose,
        validation_split=0.0
    )
    pred = np.array([])
    last_x = X[-1]


    for _ in range(steps):
        p = model.predict(last_x.reshape(1, -1, 1))[0, 0]
        pred = np.append(pred, p)
        #last_x[:T] = np.roll(last_x[:T], -1)
        #last_x[T - 1] = p
        #last_x[T:] = update_stats(last_x)
        last_x = np.roll(last_x, -1)
        last_x[-1] = p

    pred = sc.inverse_transform(pred.reshape(-1, 1))
    # pred = np.array(pred).astype('float64')
    # pred = list(pred)
    # logging.info(pred)
    
    pred.reshape(-1)
    pred[0] = train_last_value + pred[0]
    
    for i in range(1, len(pred)):
        pred[i] += pred[i-1]
    

    result = {'result': list(pred.reshape(-1)), 'epochs': r.epoch[-1] + 1, 'loss': min(r.history['loss']), 'loss_last': r.history['loss'][-1]}
    return result


if __name__ == "__main__":
    from time import time
    t1 = time()
    verbose = 2
    data = np.sin(np.arange(0.0, 28.0, 0.35)*2)
    result = predict_series(data, steps=66, r1_nodes=14, r2_nodes=14, fc1_nodes=20)
    print('exec time: {:8.3f}'.format(time()-t1))
    #print(result['result'][:2])
    print(print(result['epochs'], result['loss']))
    import seaborn as sns
    sns.lineplot(x=range(30), y=data[-30:], color='r')
    sns.lineplot(x=range(30, 30+len(result['result'])), y=result['result'], color='b')