#!/usr/bin/env python
# coding: utf-8

import sys
import os
import argparse
import numpy as np


program_description = '''\
Compute conditional mean of v given condition c by data binning
'''

# Commandline argument parser
parser = argparse.ArgumentParser(description=program_description)
parser.add_argument("-x", "--xindex", help="sampling x index", type=int, required=True)
parser.add_argument("-v", "--variable", help="file containing variable to calculate conditional mean", required=True)
parser.add_argument("-n", "--nbins", help="number of bins per condtion c interval", default=45, )
parser.add_argument("-b", "--boundary", help="width of boundary region for fine mesh", default=0.06, )
args = parser.parse_args()
params = vars(args)


# Parameters 
xidx = int(params["xindex"])
vfilename = params["variable"]
vname = os.path.splitext(os.path.basename(vfilename))[0]

nbins = int(params["nbins"])
boundary = float(params["boundary"])

cfile = "c.dat"
ufile = vfilename # "u.dat"



print ("Computing conditional mean at {}".format(xidx))



def cmean_binning (c, v, nbins=100, boundary=0.1, quiet=False):
    '''
    Compute conditional mean of v given condition c by data binning
    '''
    eps = np.finfo(np.float).eps

    c_bins = np.hstack(
        (np.linspace(0, boundary, nbins),
         np.linspace(boundary, 1-boundary, nbins)[1:],
         np.linspace(1-boundary, 1, nbins)[1:],
        )
    )
    c_bins[0] = -eps # to include c = 0 samples to the first bin

    cstar = np.zeros(c_bins.size + 1)
    cstar[1:-1] = (c_bins[1:]+c_bins[:-1])/2.
    cstar[-1] = 1.0

    v_sums = np.zeros(c_bins.size)
    v_counts = np.zeros(c_bins.size, dtype=int)
    
    ntimes = c.shape[0]
    
    progress_view = 0
    indicator_length = 50
    
    if not quiet: print (indicator_length * "|")

    for tidx, cplane in enumerate(c):
        
        progress = indicator_length * (tidx / ntimes)
        if np.floor(progress) > progress_view:
            for j in range(int(np.floor(progress) - progress_view)):
                if not quiet: print("=", end='')
            progress_view = np.floor(progress)

        v_plane = v[tidx]

        bin_indices = np.digitize(cplane, c_bins, 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, c_bins.size):
            binned = v_plane[bin_indices == i]
            v_sums[i] += binned.sum()
            v_counts[i] += binned.size

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

    # Boundary treatment

    return cstar, v_avg, c_bins, v_counts


# Load data

sc = np.memmap(cfile, mode="r", dtype=np.double).reshape((512,-1,256,256))
sv = np.memmap(ufile, mode="r", dtype=np.double).reshape((512,-1,256,256))

c_at_x = sc[xidx]
v_at_x = sv[xidx]

print ("Completed loading data".format(xidx))


# Compute conditional mean

cstar, cmean, bin_edges, count = cmean_binning(c_at_x, v_at_x, nbins=nbins, boundary=boundary, quiet=True)


# Save result
arr_dict = {}
arr_dict["cstar"] = np.asarray(cstar)
arr_dict["cmean"] = np.asarray(cmean)
arr_dict["bins"] = np.asarray(bin_edges)
arr_dict["count"] = np.asarray(count)

np.savez("cmean_{}_given_c_{:03d}".format(vname, xidx), **arr_dict)