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

import sys
import argparse
import numpy as np
from fastkde import fastKDE


# Commandline argument parser
parser = argparse.ArgumentParser()
parser.add_argument("-x", "--xindex", help="sampling x index", type=int, required=True)
parser.add_argument("-p", "--point-sigma", help="number of pdf grid points per sigma", default=20, )
args = parser.parse_args()
params = vars(args)


xidx = int(params["xindex"])
numPointsPerSigma = int(params["point_sigma"])

sigwidth = 1.

print ("Computing pdf at {}".format(xidx))


def sigmoid (y, a=1.0):
    eps = np.finfo(np.float).eps
    return 1. / (1. + np.exp(-y/a))

def isigmoid(x, a=1.):
    eps = np.finfo(np.float).eps
    return - a * np.log((1.+2*eps)/(x+eps) - 1.)

def disigmoid(x, a=1.):
    eps = np.finfo(np.float).eps
    return - a * (1./((1.+2*eps)/(x+eps) - 1.)) * (- 1. / ((x+eps)**2.))


sc = np.memmap("c.dat", mode="r", dtype=np.double).reshape((512,-1,256,256))
sabsk = np.memmap("u.dat", mode="r", dtype=np.double).reshape((512,-1,256,256))

c_absk_pdf, (ax1, ax2) = fastKDE.conditional(
        sabsk[xidx].ravel(), 
        isigmoid(sc[xidx].ravel(), a=sigwidth), 
        peakFrac=0.001,
        numPointsPerSigma=numPointsPerSigma
        )

finite_ax1 = sigmoid(ax1, a=sigwidth)

arr_dict = {}
arr_dict["cpdf"] = np.asarray(c_absk_pdf)
arr_dict["ax1"] = np.asarray(ax1)
arr_dict["ax2"] = np.asarray(ax2)

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