# -*- coding: utf-8 -*-
"""
Created on Tue Jun 16 08:36:09 2020

@author: leclaire
"""

# Script for DA Challenges output writing (example of the First DA challenge dataset)

import numpy as np
import h5py

# Mesh extent and number of points in each direction

Xlin = np.linspace(-50,50,251)
Ylin = np.linspace(-25,25,126)
Zlin = np.linspace(0.01,30.01,76)

# 3D mesh: first index describes variation in X, second in Y, third in Z

Y,X,Z = np.meshgrid(Ylin,Xlin,Zlin)

VX = np.zeros_like(X)
VY = np.zeros_like(X)
VZ = np.zeros_like(X)
dVXdX = np.zeros_like(X)
dVXdY = np.zeros_like(X)
dVXdZ = np.zeros_like(X)
dVYdX = np.zeros_like(X)
dVYdY = np.zeros_like(X)
dVYdZ = np.zeros_like(X)
dVZdX = np.zeros_like(X)
dVZdY = np.zeros_like(X)
dVZdZ = np.zeros_like(X)
p = np.zeros_like(X)

"""
If not in that format already, write/interpolate your results into the 3D numpy arrays for velocity, velocity gradient and pressure, etc.
VX = ...
VY = ...
...
p = ...
"""

# Pressure normalization: in the First DA challenge, this is with respect to (X,Y,Z)=(0,0.2,0.01)

p = p-p[125,63,0]

# Output file naming

algoStr = "myAlgo"
# In case you upload two results obtained with two different algorithms, then use different team names, e.g. myAlgo_1 and myAlgo_2
myPpp = 0.16

resultFile = '{}_OutputDA_ppp_0_{:3d}.h5'.format(algoStr,int(np.around(myPpp*1000)))

"""
At file writing, convert the 3D arrays to a 1D linear array with ordering as requested in the instructions.

Note that VXout = VX.flatten(order='F') is equivalent to e.g.

ii=0
for k in range(0,len(Zlin):
    for j in range(0,len(Ylin)):
        for i in range(0,len(Xlin)):
            VXout[ii] = VX[i,j,k]
            ii = ii+1
"""


with h5py.File(resultFile,'w') as f:
    f.create_dataset('VX',data = VX.flatten(order='F'))
    f.create_dataset('VY',data = VY.flatten(order='F'))
    f.create_dataset('VZ',data = VZ.flatten(order='F'))
    f.create_dataset('dVXdX',data = dVXdX.flatten(order='F'))
    f.create_dataset('dVYdX',data = dVYdX.flatten(order='F'))
    f.create_dataset('dVZdX',data = dVZdX.flatten(order='F'))
    f.create_dataset('dVXdY',data = dVXdY.flatten(order='F'))
    f.create_dataset('dVYdY',data = dVYdY.flatten(order='F'))
    f.create_dataset('dVZdY',data = dVZdY.flatten(order='F'))
    f.create_dataset('dVXdZ',data = dVXdZ.flatten(order='F'))
    f.create_dataset('dVYdZ',data = dVYdZ.flatten(order='F'))
    f.create_dataset('dVZdZ',data = dVZdZ.flatten(order='F'))
    f.create_dataset('p',data = p.flatten(order='F'))
    

# Display contents of created file
with h5py.File(resultFile,'r') as f:
    for variable in list(f.keys()):
        print(variable)
        data = f[variable]
        print(data.shape)
        print(data.dtype)