import argparse
import datetime
import numpy as np
import dnstool
from pycompact import CompactScheme


program_description = '''\
Calculate product of conditional fluctuations of u and ddx(c).
'''

# Commandline argument parser
parser = argparse.ArgumentParser(description=program_description)
parser.add_argument("-c", "--case", help="target case name", required=True)
args = parser.parse_args()
params = vars(args)

# Case Parameters
cases = dnstool.case_library()
casename = params["case"]
case = cases[casename]
nx, ny, nz = case.shape

# Initialize Compact Scheme object
cs = CompactScheme(nx, ny, nz, False, True, True, 4, 2, 2)

# Open Files
cfile = "c.dat"
ufile = "u.dat"
ddxcfile = "ddxc.dat"
resfile = "uddxc.dat"

c = np.memmap(cfile, mode="r", dtype=np.double).reshape((nx,-1,ny,nz))
u = np.memmap(ufile, mode="r", dtype=np.double).reshape((nx,-1,ny,nz))
ddxc = np.memmap(ddxcfile, mode="r", dtype=np.double).reshape((nx,-1,ny,nz))
storage = np.memmap(resfile, mode='w+', dtype=np.double, shape=c.shape)

# Read pre-calculated Conditional Means
batch_cmean_ddxc_c = np.load('./cmean_ddxc_given_c.npz')
cstar_ddxc_c = batch_cmean_ddxc_c['cstar']
cmean_ddxc_c = batch_cmean_ddxc_c['cmean']

batch_cmean_u_c = np.load('./cmean_u_given_c.npz')
cstar_u_c = batch_cmean_u_c['cstar']
cmean_u_c = batch_cmean_u_c['cmean']

# Compute product of conditional fluctuations at each x index
for xidx in range(nx):

    print (datetime.datetime.now(), xidx) # Print wall time and x index

    # Calculate Conditional u Fluctuation
    u_mean_at_x = np.interp (c[xidx], cstar_u_c[xidx], cmean_u_c[xidx])
    u_at_x = u[xidx]
    u_flux_at_x = u_at_x - u_mean_at_x

    # Calculate Conditional d/dx(c) Fluctuation
    ddxc_mean_at_x = np.interp (c[xidx], cstar_ddxc_c[xidx], cmean_ddxc_c[xidx])
    ddxc_at_x = u[xidx]
    ddxc_flux_at_x = ddxc_at_x - ddxc_mean_at_x

    # Save u' * (d/dx(c))'
    storage[xidx,:,:,:] = u_flux_at_x * ddxc_flux_at_x

# Flush all write to disk
storage.flush()