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


program_description = '''\
Read all DNS data files and collect c and ddx(c) at selected x coordinates.
Compute conditional mean of ddx(c) given condition c by data binning
'''

# Commandline argument parser
parser = argparse.ArgumentParser(description=program_description)
parser.add_argument("-c", "--case", help="target case name, required")
parser.add_argument("-l", "--list", help="list all case names", action='store_true')
parser.add_argument("-t", "--test", help="test run, use only 10 timesteps", action='store_true')
args = parser.parse_args()
params = vars(args)

cases = dnstool.case_library()

if params["list"] :
    for c in cases:
        print(c, cases[c].case_root)
    sys.exit(0)
else:
    if params["case"] is None :
        sys.exit("--case CASE not given")

casename = params["case"]
case = cases[casename]
nx, ny, nz = case.shape

n_timestep = len(case.data_files)
if params["test"] :
    n_timestep = 10

y_idx = 4

x_indices_all = {
    'IC1': np.array([272, 285, 294, 302, 309, 315, 321, 329, 339]),
    'IC2': np.array([251, 273, 286, 298, 309, 320, 330, 342, 359]),
    'IC3': np.array([246, 270, 285, 299, 311, 322, 333, 346, 363]),
    'IC4': np.array([262, 279, 291, 300, 309, 317, 326, 337, 349]),
    'IC8': np.array([254, 277, 291, 302, 312, 321, 330, 340, 354]),
    'IC10': np.array([285, 294, 300, 304, 308, 312, 316, 320, 326]),
    'IC11': np.array([264, 281, 292, 301, 309, 317, 325, 334, 347]),
    'IC12': np.array([253, 276, 290, 301, 311, 320, 330, 340, 355]),
    'IC13': np.array([243, 269, 287, 300, 312, 323, 334, 347, 365]),
    'VIC1': np.array([286, 294, 299, 304, 308, 312, 315, 320, 326]),
    'VIC2': np.array([264, 280, 291, 300, 309, 317, 325, 335, 348]),
    'VIC4': np.array([273, 287, 295, 302, 309, 315, 321, 328, 338]),
    'VIC10': np.array([292, 298, 302, 305, 308, 310, 313, 316, 320]),
    'VIC13': np.array([253, 276, 291, 302, 312, 321, 331, 341, 355]),}

x_indices = x_indices_all[casename]

print("Calculating Conditional Mean of ddx(c) given c")
print(" *    Case:", casename)
print(" * x index:", x_indices)


class ConditionalMeanBinning:
    def __init__ (self, name='', nbins=100, boundary=0.1):
        self.eps = np.finfo(np.float).eps

        self.name = name

        # Point grid containing interval boundaries
        self.grid = np.hstack(
                (np.linspace(0, boundary, nbins),
                 np.linspace(boundary, 1-boundary, nbins)[1:],
                 np.linspace(1-boundary, 1, nbins)[1:],
                )
                )
        self.grid[0] = -self.eps # to include c = 0 samples to the first bin

        # interval center points
        self.cstar = np.zeros(self.grid.size + 1)
        self.cstar[1:-1] = (self.grid[1:]+self.grid[:-1])/2.
        self.cstar[-1] = 1.0

        self.v_sums = np.zeros(self.grid.size)
        self.v_counts = np.zeros(self.grid.size, dtype=int)

    def feed_samples (self, c, v):
        ''' Feed arrays c, v of samples and conditions '''

        bin_indices = np.digitize(c, self.grid, right=True)
        # assign bin indices. bin intervals are right inclusive
        # i-th bin: c_bin[i-1] < c <= c_bin[i]

        for i in range(1, self.grid.size):
            binned = v[bin_indices == i]
            self.v_sums[i] += binned.sum()
            self.v_counts[i] += binned.size

    def average (self):
        v_avg = np.zeros(self.v_sums.size+1)
        v_avg[:-1] = self.v_sums / (self.v_counts + self.eps)
        # add eps to prevent nan values from divide by zero

        d = self.sum_count()
        d["bin_cstar"] = self.cstar
        d["bin_avg"] = v_avg
        return d

    def sum_count (self):
        return {
            "bin_grid" : self.grid, 
            "bin_sum" : self.v_sums, 
            "bin_count" : self.v_counts
        }

cs = CompactScheme(nx, ny, nz, False, True, True, 4, 2, 2)

# c mean at sampling x coordinates (verification purpose)
cmean_verify = np.zeros(len(x_indices))

avg_objs = [ConditionalMeanBinning(name=str(xi)) for xi in x_indices]

for i, fname in enumerate(case.data_files[:n_timestep]):
    print(datetime.datetime.now(), '{:5d}/{}'.format(i+1, n_timestep), fname)

    fpath = os.path.join(case.case_root, fname)
    time, y = case.read_single_field(fpath, y_idx)
    c = (1. - y)
    ddxc = cs.ddx(c)

    for j, xi in enumerate(x_indices):
        avg_objs[j].feed_samples(c[:,:,xi], ddxc[:,:,xi])
        cmean_verify[j] += np.sum(c[:,:,xi])
print(datetime.datetime.now(), 'Finished')

print()
print("Verify mean(c) at sampling x coordinates")
cmean_verify /= (n_timestep * ny * nz)
for xi_cmean_pair in zip(x_indices, cmean_verify):
    print (xi_cmean_pair)

print()
print("Save the result")
for j, xi in enumerate(x_indices):
    filename = "cavg_ddxc_c_{}_{}".format(casename, avg_objs[j].name)
    np.savez(filename, **avg_objs[j].average())
    print("saved ", filename)