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', (['dataDir'], {}), False, 'from os import listdir\n'), (27, 'numpy.amax', 'np.amax', (['targets[:, (0), :, :]'], {}), True, 'import numpy as np\n'), (16, 'numpy.load', 'np.load', (['(dataDir + file)'], {}), True, 'import numpy as np\n'), (23, 'numpy.sqrt', 'np.sqrt', (['(targets[:, (1), :, :] * targets[:, (1), :, :] + targets[:, (2), :, :] *\n    targets[:, (2), :, :])'], {}), True, 'import numpy as np\n'), (29, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'num': 'None', 'figsize': '(20, 10)', 'dpi': '(80)', 'facecolor': '"""w"""', 'edgecolor': '"""k"""'}), True, 'import matplotlib.pyplot as plt\n'), (32, 'matplotlib.pyplot.subplot', 'plt.subplot', (['(331)'], {}), True, 'import matplotlib.pyplot as plt\n'), (33, 'matplotlib.pyplot.title', 'plt.title', (['"""x vel"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (34, 'matplotlib.pyplot.imshow', 'plt.imshow', (['input[(0), :, :]'], {'cmap': '"""jet"""'}), 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', (['(332)'], {}), True, 'import matplotlib.pyplot as plt\n'), (37, 'matplotlib.pyplot.title', 'plt.title', (['"""y vel"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (38, 'matplotlib.pyplot.imshow', 'plt.imshow', (['input[(1), :, :]'], {'cmap': '"""jet"""'}), 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')]