import os
from time import time

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from icecream import ic
from tensorflow import keras


# os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

global_random_seed = 3884
np.random.seed(global_random_seed)


def gene_test_data(num_points: int, low_val=-1, high_val=1):
    return np.linspace(low_val, high_val, num_points)


def get_network(input_dim, output_dim, width, depth):
    # Random weight
    initializer = keras.initializers.RandomUniform(
        minval=-1., maxval=1., seed=global_random_seed)
    # initializer = keras.initializers.Constant(1)

    # different activation functions
    activation_func = keras.activations.tanh
    # activation_func = tf.math.sin

    hidden_layers = [keras.layers.InputLayer(input_shape=[input_dim, ])]
    hidden_layers += [keras.layers.Dense(width, activation=activation_func,
                                         kernel_initializer=initializer, bias_initializer=initializer) for _ in range(depth-1)]
    hidden_layers.append(keras.layers.Dense(
        output_dim, kernel_initializer=initializer))
    model = keras.Sequential(hidden_layers)
    return model


def predict(input):
    result = []
    for model in all_models:
        temp = model.predict(input)
        result.append(temp)
    return result


model_5 = get_network(1, 1, 5, 15)
model_10 = get_network(1, 1, 10, 15)
model_20 = get_network(1, 1, 20, 15)
model_40 = get_network(1, 1, 40, 15)
all_models = [model_5, model_10, model_20, model_40]

if __name__ == '__main__':
    st_time = time()
    test_data = gene_test_data(100)

    result = predict(test_data)
    ic(np.shape(result))
    print(time()-st_time)
    # result = list(zip(*result))
    ic(np.shape(result))
    result_fft = []
    for i in range(len(result)):
        result_fft.append(list(np.fft.fft(result[i], axis=0)))

    ic(np.shape(result_fft))

    fig, ((p1, p2, p3, p4), (p5, p6, p7, p8)) = plt.subplots(2, 4)
    pics = [p1, p2, p3, p4, p5, p6, p7, p8]
    for i in range(len(result)):
        pics[i].scatter(range(len(result[i])), list(result[i]), s=1.5, c='r')
        pics[i+4].scatter(range(len(result_fft[i])),
                          list(result_fft[i]), s=0.5, c='b')
        pics[i].set(adjustable='box', aspect=1.0/pics[i].get_data_ratio())
        pics[i+4].set(adjustable='box', aspect=1.0/pics[i+4].get_data_ratio())
    plt.show()