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()
[((224, 240), 'os.listdir', 'listdir', (['dataDir'], {}), '(dataDir)\n', (231, 240), False, 'from os import listdir\n'), ((833, 863), 'numpy.amax', 'np.amax', (['targets[:, (0), :, :]'], {}), '(targets[:, (0), :, :])\n', (840, 863), True, 'import numpy as np\n'), ((413, 436), 'numpy.load', 'np.load', (['(dataDir + file)'], {}), '(dataDir + file)\n', (420, 436), True, 'import numpy as np\n'), ((645, 751), 'numpy.sqrt', 'np.sqrt', (['(targets[:, (1), :, :] * targets[:, (1), :, :] + targets[:, (2), :, :] *\n    targets[:, (2), :, :])'], {}), '(targets[:, (1), :, :] * targets[:, (1), :, :] + targets[:, (2), :,\n    :] * targets[:, (2), :, :])\n', (652, 751), True, 'import numpy as np\n'), ((887, 963), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'num': 'None', 'figsize': '(20, 10)', 'dpi': '(80)', 'facecolor': '"""w"""', 'edgecolor': '"""k"""'}), "(num=None, figsize=(20, 10), dpi=80, facecolor='w', edgecolor='k')\n", (897, 963), True, 'import matplotlib.pyplot as plt\n'), ((990, 1006), 'matplotlib.pyplot.subplot', 'plt.subplot', (['(331)'], {}), '(331)\n', (1001, 1006), True, 'import matplotlib.pyplot as plt\n'), ((1011, 1029), 'matplotlib.pyplot.title', 'plt.title', (['"""x vel"""'], {}), "('x vel')\n", (1020, 1029), True, 'import matplotlib.pyplot as plt\n'), ((1034, 1074), 'matplotlib.pyplot.imshow', 'plt.imshow', (['input[(0), :, :]'], {'cmap': '"""jet"""'}), "(input[(0), :, :], cmap='jet')\n", (1044, 1074), True, 'import matplotlib.pyplot as plt\n'), ((1112, 1126), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ([], {}), '()\n', (1124, 1126), True, 'import matplotlib.pyplot as plt\n'), ((1131, 1147), 'matplotlib.pyplot.subplot', 'plt.subplot', (['(332)'], {}), '(332)\n', (1142, 1147), True, 'import matplotlib.pyplot as plt\n'), ((1152, 1170), 'matplotlib.pyplot.title', 'plt.title', (['"""y vel"""'], {}), "('y vel')\n", (1161, 1170), True, 'import matplotlib.pyplot as plt\n'), ((1175, 1215), 'matplotlib.pyplot.imshow', 'plt.imshow', (['input[(1), :, :]'], {'cmap': '"""jet"""'}), "(input[(1), :, :], cmap='jet')\n", (1185, 1215), True, 'import matplotlib.pyplot as plt\n'), ((1218, 1232), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ([], {}), '()\n', (1230, 1232), True, 'import matplotlib.pyplot as plt\n'), ((1238, 1248), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (1246, 1248), True, 'import matplotlib.pyplot as plt\n')]