import numpy as np
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
from os import listdir
from tensorflow.keras.callbacks import ModelCheckpoint

dataDir = "./data/trainSmallFA/"
files = listdir(dataDir)
files.sort()
totalLength = len(files)
inputs = np.empty((len(files), 3, 64, 64))
targets = np.empty((len(files), 3, 64, 64))

for i, file in enumerate(files):
    npfile = np.load(dataDir + file)
    d = npfile['a']
    inputs[i] = d[0:3]  # inx, iny, mask
    targets[i] = d[3:6]  # p, velx, vely

# print("inputs shape = ", inputs.shape)
print(np.shape(targets[:, 1, :, :].flatten()))
maxvel = np.amax(np.sqrt(targets[:, 1, :, :]* targets[:, 1, :, :]
                         + targets[:, 2, :, :]* targets[:, 2, :, :]))
print(maxvel)
targets[:, 1:3, :, :] /= maxvel
targets[:, 0, :, :] /= np.amax(targets[:, 0, :, :])
for input in inputs:
    plt.figure(num=None, figsize=(20, 10), dpi=80, facecolor='w', edgecolor='k')

    # predicted data
    plt.subplot(331)
    plt.title('x vel')
    plt.imshow(input[0, :, :], cmap='jet')  # vmin=-100,vmax=100, cmap='jet')
    plt.colorbar()
    plt.subplot(332)
    plt.title('y vel')
    plt.imshow(input[1, :, :], cmap='jet')
    plt.colorbar()

    plt.show()
[(9, 'os.listdir', 'listdir', False, 'from os import listdir\n'), (27, 'numpy.amax', 'np.amax', True, 'import numpy as np\n'), (16, 'numpy.load', 'np.load', True, 'import numpy as np\n'), (23, 'numpy.sqrt', 'np.sqrt', True, 'import numpy as np\n'), (29, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (32, 'matplotlib.pyplot.subplot', 'plt.subplot', True, 'import matplotlib.pyplot as plt\n'), (33, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (34, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (35, 'matplotlib.pyplot.colorbar', 'plt.colorbar', True, 'import matplotlib.pyplot as plt\n'), (36, 'matplotlib.pyplot.subplot', 'plt.subplot', True, 'import matplotlib.pyplot as plt\n'), (37, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (38, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (39, 'matplotlib.pyplot.colorbar', 'plt.colorbar', True, 'import matplotlib.pyplot as plt\n'), (41, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n')]