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

@author: leclaire
"""

# Script for 1st LPT and DA Challenge output writing, DA cases

import numpy as np
import h5py

# Mesh extent and number of points in each direction
# Example of a domain covering [-50,50]x[-25,25]x[0.01,30.01] (X Y Z), with 251x126x76 points

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)
T = 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
X = ...
...
VX = ...
...
p = ...
T = ...
"""

# Pressure normalization: with respect to a specified point, in the present example point (X,Y,Z)=(0,0.2,0.01)
p = p-p[125,63,0]

# Output file naming
Team = "ONERA"
# In case you upload two results obtained with two different algorithms, then use different team names, e.g. ONERA_XXXX and ONERA_YYYY
Case = 1
ImageStrategy = "TR"
Ppp = 0.16
InstantNumber = 25

if Case==2:
    Frequency = 1 # only for case 2
    InertialPart = 0 # only for case 2
    resultFileTemplate = '{}_DA_CASE{:02d}_{}_ppp_0_{:3d}_FR_{}_I{}_PartFile_{:04d}.h5'
    resultFile = resultFileTemplate.format(Team,Case,ImageStrategy,int(np.around(Ppp*1000)),Frequency,InertialPart,InstantNumber)
else:
    resultFileTemplate = '{}_DA_CASE{:02d}_{}_ppp_0_{:3d}_PartFile_{:04d}.h5'
    resultFile = resultFileTemplate.format(Team,Case,ImageStrategy,int(np.around(Ppp*1000)),InstantNumber)

"""
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('X',data = X.flatten(order='F'))
    f.create_dataset('Y',data = Y.flatten(order='F'))
    f.create_dataset('Z',data = Z.flatten(order='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)