# -*- coding: utf-8 -*-
"""
Created on Thu Mar 30 10:15:05 2017

@author: hazel.bain

    This module generates the PDFs for the
    Chen geoeffective magnetic cloud prediction Bayesian formulation. 
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
    
    The data is stored as a pickle file and should be read in as:
    
        events_frac = pickle.load(open("events_frac.p","rb"))
    
    The top level create_pdfs function generates all the input PDFs
    and returns the posterior PDF P((Bzm, tau) n e|Bzm', tau' ; f) along with
    some other diagnostic variables. 
    
        Pbzm_tau_e_bzmp_taup, norm_bzm_tau_e_bzmp_taup, P0, P1, P1_map = create_pdfs(events_frac, kernel_alg = 'scipy_stats')
    
    Due to a relatively small smaple of geoeffective events, kernel density estimation
    is used to smooth the data and generate a non parametric PDFs.


"""


from sklearn.neighbors import KernelDensity
from scipy import stats
import pickle as pickle
import scipy.integrate as integrate
   
#import kde_hb as hb

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
from mpl_toolkits.mplot3d import Axes3D


def create_pdfs(events_frac, kernel_alg = 'scipy_stats', \
                ranges = [-150, 150, -250, 250], nbins = [50j, 100j],\
                fracs = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0], ew = 2, nw = 0.5, plotting=[0,0,0,0,0], fname='' ):

    """
    Create the PDFs for the
    Chen geoeffective magnetic cloud prediction Bayesian formulation. 
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
       
    inputs:
        
    events_frac = dataframe
        contains output variables from a fit to solar wind magnetic field data
    kernel_alg = string
        choose between scikit learn and scipy stats python KDE algorithms
    ranges = 4 element array 
        defines the axis ranges for Bzm and tau [bmin, bmax, tmin, tmax]
    nbins = 2 elements array 
        defines the number of bins along bzm and tau [nbins_Bzm, nbins_tau]
    ew = float
        defines the kernel smoothing width for the geoeffective events
    nw = float
        defines the kernel smoothing width for the nongeoeffective events    
    plotting = int array
        indices indicate which PDFs to plot       
    
    """ 
    
    #width of smoothing box
    #ew = 2
    #nw = 0.5
    
    
    #range of Bzm and tau to define PDFs 
    #note - tmin is negative due to requirement to reflect raw data points in 
    #the tau = 0 axis to combat boundary effects when implementing kernel density
    #estimate smoothing
    #bmin = -150
    #bmax = 150
    #tmin = -250
    #tmax = 250
    
    #ranges = [bmin, bmax, tmin, tmax]
    
    #number of data bins in each dimension (dt takes into account the reflection)
    #db = 5j
    #dt = 10j
    
    #nbins = [db, dt]
    
    
    #create input PDFS
    Pe = P_e(events_frac)
    Pn = P_n(events_frac)
    
    #Add a perturbation to Bzmp to break symmtry for KDE to work
    #b_indices = np.mgrid[ranges[0]:ranges[1]:nbins[0]]
    #delta_b = b_indices[1] - b_indices[0]
    #np.random.seed(2)
    #events_frac.bzm_predicted = events_frac.bzm_predicted + np.random.uniform(-delta_b, delta_b, len(events_frac))
    
    
    Pbzm_tau_e, norm_bzm_tau_e = P_bzm_tau_e(events_frac, ranges=ranges,\
        nbins=nbins, kernel_width = ew, plotfig = plotting[0])
    
    Pbzmp_taup_e, norm_bzmp_taup_e = P_bzmp_taup_e(events_frac, ranges=ranges, \
        nbins=nbins, fracs=fracs, kernel_width = ew, plotfig = plotting[1])
    
    Pbzmp_taup_n, norm_bzmp_taup_n = P_bzmp_taup_n(events_frac, ranges=ranges,\
        nbins=nbins, fracs=fracs, kernel_width = nw, plotfig=plotting[2])
    
    Pbzmp_taup_bzm_tau_e, norm_bzmp_taup_bzm_tau_e, P0, tmpint, indices = P_bzmp_taup_bzm_tau_e(events_frac, \
        ranges=ranges, fracs=fracs, nbins=nbins, kernel_width = ew, plotfig=plotting[3])
    
    Pbzm_tau_e_bzmp_taup, norm_bzm_tau_e_bzmp_taup, P1, P1_map, axis_vals = P_bzm_tau_e_bzmp_taup(Pe, \
                                                    Pn,\
                                                    Pbzm_tau_e, \
                                                    Pbzmp_taup_e,\
                                                    Pbzmp_taup_n,\
                                                    Pbzmp_taup_bzm_tau_e, \
                                                    ranges = ranges, nbins = nbins,\
                                                    fracs=fracs, \
                                                    plotfig = plotting[4])
    
    
    ##plot all pdfs in output figure       
    
    for i in range(1, len(fracs)):
    
        gbzm = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm
        gtau = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau
        
        gbzmn = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm
        gtaun = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau
        
        gbzmp = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm_predicted
        gtaup = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau_predicted
        gbzmp.dropna(inplace=True)
        gtaup.dropna(inplace=True)

        gbzmp_n = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm_predicted
        gtaup_n = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau_predicted 
        gbzmp_n.dropna(inplace=True)
        gtaup_n.dropna(inplace=True)


        fig, ((ax1,ax2),(ax3,ax4),(ax5,ax6)) = plt.subplots(3,2,figsize = (8.5,11))
        fontP = FontProperties()
        fontP.set_size('x-small')
        
        c1 = ax1.imshow(np.rot90(Pbzm_tau_e), extent=(ranges[0],ranges[1],0,ranges[3]), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax1.plot(gbzm, gtau, 'k.', markersize=4, c='r', label = 'bzm, tau, g = 1')
        ax1.set_xlim([ranges[0], ranges[1]])
        ax1.set_ylim([0, ranges[3]])
        ax1.set_xlabel('Bzm')
        ax1.set_ylabel('Tau')
        ax1.set_title('P_bzm_tau_e, bandwidth = '+str(ew), fontsize = 'small')
        #cbaxes1 = ax1.add_axes([0.8, 0.1, 0.03, 0.8]) 
        fig.colorbar(c1, ax = ax1, fraction=0.025, format='%.2E')      #, cax = cbaxes1)
        leg1 = ax1.legend(loc='upper right', prop = fontP, fancybox=True)
        leg1.get_frame().set_alpha(0.5)
        
        c2 = ax2.imshow(np.rot90(Pbzmp_taup_e[:,:,i-1]), extent=(ranges[0],ranges[1],0,ranges[3]), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax2.plot(gbzmp, gtaup, 'k.', markersize=4, c='r', label = 'bzmp, taup, g = 1')
        ax2.set_xlim([ranges[0], ranges[1]])
        ax2.set_ylim([0, ranges[3]])
        ax2.set_xlabel('Bzm_p')
        ax2.set_ylabel('Tau_p')
        ax2.set_title('P_bzmp_taup_e, bandwidth = '+str(ew), fontsize = 'small')
        fig.colorbar(c2, ax = ax2, fraction=0.025, format='%.2E')
        leg2 = ax2.legend(loc='upper right', prop = fontP, fancybox=True)
        leg2.get_frame().set_alpha(0.5)
        
        c3 = ax3.imshow(np.rot90(Pbzmp_taup_n[:,:,i-1]), extent=(ranges[0],ranges[1],0,ranges[3]), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax3.plot(gbzmp_n, gtaup_n, 'k.', markersize=4, c='b', label = 'bzmp, taup, g = 0')
        ax3.set_xlim([ranges[0], ranges[1]])
        ax3.set_ylim([0, ranges[3]])
        ax3.set_xlabel('Bzm_p')
        ax3.set_ylabel('Tau_p')
        ax3.set_title('P_bzmp_taup_n, bandwidth = '+str(nw), fontsize = 'small')
        #cbaxes3 = ax1.add_axes([0.8, 0.1, 0.03, 0.8]) 
        fig.colorbar(c3, ax = ax5, fraction=0.025, format='%.2E')      #, cax = cbaxes3)
        leg3 = ax3.legend(loc='upper right', prop = fontP, fancybox=True)
        leg3.get_frame().set_alpha(0.5)
    
        #print(np.squeeze(indices).shape)
        predicted_duration = 15.0
        predicted_bzmax = -26.0
        indt = np.min(np.where(indices[3,:] > predicted_duration))
        indb = np.max(np.where(indices[2,:] < predicted_bzmax))
    
        c4 = ax4.imshow(np.rot90(Pbzmp_taup_bzm_tau_e[:,:,indb,indt,i-1]), extent=(ranges[0],ranges[1],0,ranges[3]), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax4.plot(gbzmp_n, gtaup_n, 'k.', markersize=4, c='b', label = 'bzmp, taup, g = 0')
        ax4.plot(gbzmp, gtaup, 'k.', markersize=4, c='r', label = 'bzmp, taup, g = 1')
        ax4.set_xlim([ranges[0], ranges[1]])
        ax4.set_ylim([0, ranges[3]])
        ax4.set_xlabel('Bzm_p')
        ax4.set_ylabel('Tau_p')
        ax4.set_title('P_bzm_tau_bzmp_taup_e, bandwidth = '+str(ew), fontsize = 'small')
        fig.colorbar(c4, ax = ax4, fraction=0.025, format='%.2E')
        leg4 = ax4.legend(loc='upper right', prop = fontP, fancybox=True)
        leg4.get_frame().set_alpha(0.5)
        
        c5 = ax5.imshow(np.rot90(Pbzm_tau_e_bzmp_taup[:,:,indb,indt,i-1]), extent=(ranges[0],ranges[1],0,ranges[3]), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax5.plot(gbzmn, gtaun, 'k.', markersize=4, c='b', label = 'bzm, tau, g = 1')
        ax5.plot(gbzm, gtau, 'k.', markersize=4, c='r', label = 'bzm, tau, g = 1')
        ax5.set_xlim([ranges[0], ranges[1]])
        ax5.set_ylim([0, ranges[3]])
        ax5.set_xlabel('Bzm_p')
        ax5.set_ylabel('Tau_p')
        ax5.set_title('P_bzm_tau_e_bzmp_taup, bandwidth = '+str(ew), fontsize = 'small')
        #cbaxes5 = ax5.add_axes([0.8, 0.1, 0.03, 0.8]) 
        #fig.colorbar(c5, ax = ax5, fraction=0.025, format='%2.2f')      #, cax = cbaxes5)
        leg5 = ax5.legend(loc='upper right', prop = fontP, fancybox=True)
        leg5.get_frame().set_alpha(0.5)
    
        c6 = ax6.imshow(np.rot90(P1_map[:,:,i-1]), extent=(ranges[0],ranges[1],0,ranges[3]), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax6.plot(gbzmn, gtaun, 'k.', markersize=4, c='b', label = 'bzm, tau, g = 0')
        ax6.plot(gbzm, gtau, 'k.', markersize=4, c='r', label = 'bzm, tau, g = 1')
        ax6.set_xlim([ranges[0], ranges[1]])
        ax6.set_ylim([0, ranges[3]])
        ax6.set_xlabel('Bzm_p')
        ax6.set_ylabel('Tau_p')
        ax6.set_title('P1_map, bandwidth = '+str(ew), fontsize = 'small')
        fig.colorbar(c6, ax = ax6, fraction=0.025, format='%2.2f')
        leg6 = ax6.legend(loc='upper right', prop = fontP, fancybox=True)
        leg6.get_frame().set_alpha(0.5)
    
    
        fig.savefig('C:/Users/hazel.bain/Documents/MC_predict/pyMCpredict/MCpredict/PDFs/plots/allpdfs_'+fname+'_f'+str(fracs[i])+'.pdf')
        plt.close()

    #create a dictionary to return PDFs etc
    P_dict = {}
    P_dict["P_e"] = Pe
    P_dict["P_n"] = Pn
    P_dict["P_bzm_tau_e"] = Pbzm_tau_e
    P_dict["norm_bzm_tau_e"] = norm_bzm_tau_e
    P_dict["P_bzmp_taup_e"] = Pbzmp_taup_e
    P_dict["norm_bzmp_taup_e"] = norm_bzmp_taup_e
    P_dict["P_bzmp_taup_n"] = Pbzmp_taup_n
    P_dict["norm_bzmp_taup_n"] = norm_bzmp_taup_n
    P_dict["P_bzmp_taup_bzm_tau_e"] = Pbzmp_taup_bzm_tau_e
    P_dict["norm_bzmp_taup_bzm_tau_e"] = norm_bzmp_taup_bzm_tau_e
    P_dict["P0"] = P0
    P_dict["tmpint"] =tmpint
    P_dict["indices"] = indices
    P_dict["P_bzm_tau_e_bzmp_taup"] = Pbzm_tau_e_bzmp_taup
    P_dict["norm_bzm_tau_e_bzmp_taup"] = norm_bzm_tau_e_bzmp_taup
    P_dict["axis_vals"] = axis_vals
    P_dict["P1"] = P1
    P_dict["P1_map"] = P1_map
    
    #save the input paramters as well
    P_dict["ew"] = ew
    P_dict["nw"] = nw
    P_dict["ranges"] = ranges
    P_dict["nbins"] = nbins
    P_dict["kernel_alg"] = kernel_alg
       
    #save a pickle file with P_dict
    pickle.dump(P_dict, open("C:/Users/hazel.bain/Documents/MC_predict/pyMCpredict/MCpredict/PDFs/Pdict_"+fname+".p", "wb"))
    
    return P_dict    


def load_pdfs(fname=''):
    
    pp = pickle.load(open("PDFs/Pdict_"+fname+".p","rb"))
    
    return pp['P_e'], pp['P_n'], pp['P_bzm_tau_e'], pp['P_bzmp_taup_e'], pp['P_bzmp_taup_n'], pp['P_bzmp_taup_bzm_tau_e'], pp['P_bzm_tau_e_bzmp_taup']

    
def P_e(events_frac):  
    
    """
    Determine the prior PDF P(e) for input into the following Chen geoeffective 
    magnetic cloud prediction Bayesian formulation
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where c_j can be e or n
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
    
    inputs:
        
    events_frac = dataframe
        contains output variables from a fit to solar wind magnetic field data
        as a function of the fraction of time through an event
            
    """
    
    #----P(e) - probability of a geoefective event 
    n_events = len(events_frac.drop_duplicates('evt_index'))
    n_geoeff_events = len(events_frac.drop_duplicates('evt_index').query('geoeff == 1.0'))
    
    #values from Chen paper
    #n_nongeoeff_events = 8600.
    #n_geoeff_events = 56.
    #n_events = n_nongeoeff_events + n_geoeff_events
    
    P_e = n_geoeff_events / n_events
    
    print('\n\n P_e: ' + str(P_e) + '\n\n')
    
    return P_e


def P_n(events_frac):  
    
    """
    Determine the prior PDF P(n) for input into the following Chen geoeffective 
    magnetic cloud prediction Bayesian formulation
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where c_j can be e or n
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
    
    inputs:
        
    events_frac = dataframe
        contains output variables from a fit to solar wind magnetic field data
        as a function of the fraction of time through an event
            
    """
    
    #----P(n) probability of nongeoeffective events
    n_events = len(events_frac.drop_duplicates('evt_index'))
    n_nongeoeff_events = len(events_frac.drop_duplicates('evt_index').query('geoeff == 0.0'))
    
    #values from Chen paper
    #n_nongeoeff_events = 8600.
    #n_geoeff_events = 56.
    #n_events = n_nongeoeff_events + n_geoeff_events
    
    P_n = n_nongeoeff_events / n_events

    print('\n\n P_n: ' + str(P_n) + '\n\n')

    return P_n


def P_bzm_tau_e(events_frac, kernel_alg = 'scipy_stats', \
                ranges = [-150, 150, -250, 250], nbins = [50j, 100j],\
                kernel_width = 2, plotfig = 0):  
    
    """
    Determine the prior PDF P(Bzm, tau|e), the probability of geoeffective event with observered bzm and tau
    for input into the following Chen geoeffective magnetic cloud prediction Bayesian formulation
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
    
    Due to a relatively small smaple of geoeffective events, kernel density estimation
    is used to smooth the data and generate a non parametric PDF. An artifact of 
    the KDE method is that the PDF will be smoothed to produce probabilities 
    extending to negative values of tau. To prevent this effect at the boundary 
    the raw data values are reflected in the tau = 0 axis and then the KDE is applied, 
    producing a symmetric density estimates about the axis of reflection - see 
    Silverman 1998 section 2.10. The required output PDF is obtained by selecting array 
    elements corresponding to postive values of tau 
    
    f'(x) = 2f(x) for x >= 0 and f'(x) = 0 for x < 0
    
    As a result of this reflection, the input range for tau extends to negative 
    values and nbins_tau is double nbins_bzm. The output PDF will be for 
    tau = [0, tmax]
    
    inputs:
        
    events_frac = dataframe
        contains output variables from a fit to solar wind magnetic field data
    kernel_alg = string
        choose between scikit learn and scipy stats python KDE algorithms        
    ranges = 4 element array 
        defines the axis ranges for Bzm and tau [bmin, bmax, tmin, tmax]
    nbins = 2 elements array 
        defines the number of bins along bzm and tau [nbins_Bzm, nbins_tau]
    kernel_widths = float
        defines the kernel smoothing width
    plotfit = int
        plot a figure of the distribution         
    
    """ 
    #range of Bzm and tau to define PDFs 
    bmin = ranges[0]
    bmax = ranges[1]
    tmin = ranges[2]
    tmax = ranges[3]
    
    #number of data bins in each dimension (dt takes into account the reflection)
    db = nbins[0]
    dt = nbins[1]
    
    #true boundary for tau and the corresponding index in the pdf array
    taumin = 0
    #dt0 = dt/2 
    
    #kernel smoothing width 
    ew = kernel_width
    
    #extract raw data points from dataframe of "actual" bzm and tau for geoeffective events
    gbzm = events_frac.drop_duplicates('evt_index').query('geoeff == 1.0 and bzm < 0.0').bzm
    gtau = events_frac.drop_duplicates('evt_index').query('geoeff == 1.0 and bzm < 0.0 ').tau
    
    #print("max bzm: "+str(gbzm.max())+", tau: "+ str(gtau.max()))
    
    
    #to handle boundary conditions and limit the density estimate to positve 
    #values of tau: reflect the data points along the tau axis, perform the 
    #density estimate and then set negative values of tau to 0 and double the 
    #density of positive values of tau
    gbzm_r = np.concatenate([gbzm, gbzm]) 
    gtau_r = np.concatenate([gtau, -gtau])
    
    #grid containing x, y positions 
    X_bzm, Y_tau = np.mgrid[bmin:bmax:db, tmin:tmax:dt]
    dt0 = int(len(Y_tau[1])/2.)
    Y_tau = Y_tau-((tmax-tmin)/len(Y_tau[1])/2.)      #to make sure the resulting PDF tau 
                                            #axis will start at 0. 

    #option to use scikit learn or scipy stats kernel algorithms
    if kernel_alg == 'sklearn':
        
        positions = np.vstack([X_bzm.ravel(), Y_tau.ravel()]).T
        values = np.vstack([gbzm_r, gtau_r]).T
        kernel_bzm_tau_e = KernelDensity(kernel='gaussian', bandwidth=ew).fit(values)
        Ptmp_bzm_tau_e = np.exp(np.reshape(kernel_bzm_tau_e.score_samples(positions).T, X_bzm.shape))
        
    elif kernel_alg == 'scipy_stats':
        
        positions = np.vstack([X_bzm.ravel(), Y_tau.ravel()])
        values = np.vstack([gbzm_r, gtau_r])
        kernel_bzm_tau_e = stats.gaussian_kde(values, bw_method = ew)
        Ptmp_bzm_tau_e = np.reshape(kernel_bzm_tau_e(positions).T, X_bzm.shape)

    #set the density estimate to 0 for negative tau, and x2 for positve tau 
    P_bzm_tau_e = Ptmp_bzm_tau_e[:,dt0::]*2

    #check the normalization, should normalize to 1
    norm_bzm_tau_e = integrate.simps(integrate.simps(P_bzm_tau_e,Y_tau[0,dt0::]), X_bzm[:,0])
    print('\n\n Normalization for P_bzm_tau_e: ' + str(norm_bzm_tau_e) + '\n\n')

    if plotfig == 1:
        
        fontP = FontProperties()                #legend
        #fontP.set_size('medium')1
        
        fig, ax = plt.subplots()
        c = ax.imshow(np.rot90(P_bzm_tau_e), extent=(bmin,bmax,taumin,tmax), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax.plot(gbzm, gtau, 'k.', markersize=4, label = 'bzm, tau, geoeff = 1')
        ax.set_xlim([bmin, bmax])
        ax.set_ylim([taumin, tmax])
        ax.set_xlabel('Bzm')
        ax.set_ylabel('Tau')
        ax.set_title('P_bzm_tau_e, bandwidth = '+str(ew))
        fig.colorbar(c)
        ax.legend(loc='upper right', prop = fontP, fancybox=True)

    #P_bzm_tau_e_rescale = P_bzm_tau_e/norm_bzm_tau_e
    #norm_bzm_tau_e_rescale = integrate.simps(integrate.simps(P_bzm_tau_e_rescale,Y_tau[0,dt0::]), X_bzm[:,0])

    #print("here00")
    #print(norm_bzm_tau_e_rescale)

    return P_bzm_tau_e, norm_bzm_tau_e


def P_bzmp_taup_e(events_frac, kernel_alg = 'scipy_stats', \
                ranges = [-150, 150, -250, 250], nbins = [50j, 100j],\
                fracs = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0], kernel_width = 2, plotfig = 0):  
    
    """
    Determine the PDF P(Bzm', tau'|e ; f), the probability of geoeffective event 
    with estimated values bzm' and tau' at fraction f throughout an event for 
    input into the following Chen geoeffective magnetic cloud prediction 
    Bayesian formulation.  This PDF contributes to the Bayesian "evidence".
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
    
    Due to a relatively small smaple of geoeffective events, kernel density estimation
    is used to smooth the data and generate a non parametric PDF. An artifact of 
    the KDE method is that the PDF will be smoothed to produce probabilities 
    extending to negative values of tau. To prevent this effect at the boundary 
    the raw data values are reflected in the tau' = 0 axis and then the KDE is applied, 
    producing a symmetric density estimates about the axis of reflection - see 
    Silverman 1998 section 2.10. The required output PDF is obtained by selecting array 
    elements corresponding to postive values of tau 
    
    f'(x) = 2f(x) for x >= 0 and f'(x) = 0 for x < 0
    
    As a result of this reflection, the input range for tau extends to negative 
    values and nbins_tau is double nbins_bzm. The output PDF will be for 
    tau' = [0, tmax]
    
    inputs:
        
    events_frac = dataframe
        contains output variables from a fit to solar wind magnetic field data
    kernel_alg = string
        choose between scikit learn and scipy stats python KDE algorithms
    ranges = 4 element array 
        defines the axis ranges for Bzm and tau [bmin, bmax, tmin, tmax]
    nbins = 2 elements array 
        defines the number of bins along bzm and tau [nbins_Bzm, nbins_tau]
    fracs = list
        fractions to split the events into for calculating evolution dependent pdfs
    kernel_width = float
        defines the kernel smoothing width
    plotfit = int
        plot a figure of the distribution 
        
    
    """
    #range of Bzm and tau to define PDFs 
    bmin = ranges[0]
    bmax = ranges[1]
    tmin = ranges[2]
    tmax = ranges[3]
    
    #number of data bins in each dimension (dt takes into account the reflection)
    db = nbins[0]
    dt = nbins[1]
    
    #true boundary for tau and the corresponding index in the pdf array
    taumin = 0
    #dt0 = dt/2 
    
    #kernel smoothing width
    ew = kernel_width
    
    #set grid containing x,y positions and use to get array size
    X_bzmp, Y_taup = np.mgrid[bmin:bmax:db, tmin:tmax:dt]
    db2 = int(len(X_bzmp[:,0]))
    dt2 = int(len(Y_taup[0,:]))
    
    #P_bzmp_taup_e is a function of the fraction of time f throughout an event 
    nfracs = len(fracs)-1
    Ptmp_bzmp_taup_e = np.zeros((db2,dt2,nfracs))
    for i in range(1, len(fracs)):
        
        #extract raw data points from dataframe of estimates bzm' and tau' for 
        #fraction f throughout geoeffective events 
        
        gbzmp = events_frac.query('geoeff == 1.0 and bzm < 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm_predicted
        gtaup = events_frac.query('geoeff == 1.0 and bzm < 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau_predicted
        
        #print("max bzmp: "+str(gbzmp.max())+", taup: "+ str(gtaup.max()))
        
        gbzmp.dropna(inplace=True)
        gtaup.dropna(inplace=True)
        
        if gtaup.max() == np.inf:
            inf_idx = gtaup[gtaup == np.inf].index.values[0]
            gtaup.drop([inf_idx], inplace = True)
            gbzmp.drop([inf_idx], inplace = True)
        
        
        #to handle boundary conditions and limit the density estimate to positve 
        #values of tau': reflect the data points along the tau' axis, perform the 
        #density estimate and then set negative values of tau' to 0 and double the 
        #density of positive values of tau'
        gbzmp_r = np.concatenate([gbzmp, gbzmp])
        gtaup_r = np.concatenate([gtaup, -gtaup])
    
        #make sure the resulting PDF tau axis will start at 0 hours. 
        dt0 = int(len(Y_taup[1])/2.)
        Y_taup = Y_taup-((tmax-tmin)/len(Y_taup[1])/2.)        
                                                            
        #option to use scikit learn or scipy stats kernel algorithms
        if kernel_alg == 'sklearn':
        
            positions = np.vstack([X_bzmp.ravel(), Y_taup.ravel()]).T
            values = np.vstack([gbzmp_r, gtaup_r]).T             
            kernel_bzmp_taup_e = KernelDensity(kernel='gaussian', bandwidth=ew).fit(values)
            
            Ptmp_bzmp_taup_e[:,:,i-1] = np.exp(np.reshape(kernel_bzmp_taup_e.score_samples(positions).T, X_bzmp.shape))

        elif kernel_alg == 'scipy_stats'  : 
            
            positions = np.vstack([X_bzmp.ravel(), Y_taup.ravel()])
            values = np.vstack([gbzmp_r, gtaup_r])
            kernel_bzmp_taup_e = stats.gaussian_kde(values, bw_method=ew)

            Ptmp_bzmp_taup_e[:,:,i-1] = np.reshape(kernel_bzmp_taup_e(positions).T, X_bzmp.shape)

       
    #set the density estimate to 0 for negative tau, and x2 for positve tau 
    P_bzmp_taup_e = Ptmp_bzmp_taup_e[:,dt0::,:]*2

    #check the normalization for the final fraction (usually at frac = 1.0)
    norm_bzmp_taup_e = integrate.simps(integrate.simps(P_bzmp_taup_e[:,:,-1],Y_taup[0,dt0::]), X_bzmp[:,0])
    print('\n\n Normalization for P_bzmp_taup_e: ' + str(norm_bzmp_taup_e) + '\n\n')

    if plotfig == 1:    

        fontP = FontProperties()                #legend
        #fontP.set_size('medium')
               
        fig, ax = plt.subplots()
        c = ax.imshow(np.rot90(P_bzmp_taup_e[:,:,-1]), extent=(bmin,bmax,taumin,tmax), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax.plot(gbzmp, gtaup, 'k.', markersize=4, label = 'bzm_p, tau_p, geoeff = 1')
        ax.set_xlim([bmin, bmax])
        ax.set_ylim([taumin, tmax])
        ax.set_xlabel('Bzm')
        ax.set_ylabel('Tau')
        ax.set_title('P_bzmp_taup_e, bandwidth = '+str(ew))
        fig.colorbar(c)
        ax.legend(loc='upper right', prop = fontP, fancybox=True)

    #P_bzmp_taup_e_rescale = np.zeros((db2,db2,nfracs))
    #for i in range(0, len(fracs)-1):
    #    P_bzmp_taup_e_rescale[:,:,i] = P_bzmp_taup_e[:,:,i]/norm_bzmp_taup_e
    #norm_bzmp_taup_e_rescale = integrate.simps(integrate.simps(P_bzmp_taup_e_rescale[:,:,-1],Y_taup[0,dt0::]), X_bzmp[:,0])

    #print("here00")
    #print(norm_bzmp_taup_e_rescale)


    #return P_bzmp_taup_e, norm_bzmp_taup_e

    return P_bzmp_taup_e, norm_bzmp_taup_e



def P_bzmp_taup_n(events_frac, kernel_alg = 'scipy_stats', \
                ranges = [-150, 150, -250, 250], nbins = [50j, 100j],\
                fracs = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0], kernel_width = 0.5, plotfig = 0):  
    
    """
    Determine the PDF P(Bzm', tau'|n ; f), the probability of nongeoeffective event 
    with estimated values bzm' and tau' at fraction f throughout an event for 
    input into the following Chen geoeffective magnetic cloud prediction 
    Bayesian formulation. This PDF contributes to the Bayesian "evidence".
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
    
    Due to a relatively small smaple of events, kernel density estimation
    is used to smooth the data and generate a non parametric PDF. An artifact of 
    the KDE method is that the PDF will be smoothed to produce probabilities 
    extending to negative values of tau. To prevent this effect at the boundary 
    the raw data values are reflected in the tau' = 0 axis and then the KDE is applied, 
    producing a symmetric density estimates about the axis of reflection - see 
    Silverman 1998 section 2.10. The required output PDF is obtained by selecting array 
    elements corresponding to postive values of tau 
    
    f'(x) = 2f(x) for x >= 0 and f'(x) = 0 for x < 0
    
    As a result of this reflection, the input range for tau extends to negative 
    values and nbins_tau is double nbins_bzm. The output PDF will be for 
    tau' = [0, tmax]
    
    inputs:
        
    events_frac = dataframe
        contains output variables from a fit to solar wind magnetic field data
    kernel_alg = string
        choose between scikit learn and scipy stats python KDE algorithms
    ranges = 4 element array 
        defines the axis ranges for Bzm and tau [bmin, bmax, tmin, tmax]
    nbins = 2 elements array 
        defines the number of bins along bzm and tau [nbins_Bzm, nbins_tau]
    fracs = list
        fractions to split the events into for calculating evolution dependent pdfs
    kernel_width = float
        defines the kernel smoothing width    
    plotfit = int
        plot a figure of the distribution 
        
    
    """ 
    

    #range of Bzm and tau to define PDFs 
    bmin = ranges[0]
    bmax = ranges[1]
    tmin = ranges[2]
    tmax = ranges[3]
    
    #number of data bins in each dimension (dt takes into account the reflection)
    db = nbins[0]
    dt = nbins[1]
    
    #true boundary for tau and the corresponding index in the pdf array
    taumin = 0
    #dt0 = dt/2 
    
    #kernel smoothing width
    nw = kernel_width
    
    #set grid containing x,y positions and use to get array size
    X_bzmp_n, Y_taup_n = np.mgrid[bmin:bmax:db, tmin:tmax:dt]
    db2 = int(len(X_bzmp_n[:,0]))
    dt2 = int(len(Y_taup_n[0,:]))
    
    #P_bzmp_taup_n is a function of the fraction of time f throughout an event
    nfracs = len(fracs)-1
    Ptmp_bzmp_taup_n = np.zeros((db2,dt2,nfracs))
    for i in range(1, len(fracs)):

        #extract raw data points from dataframe of estimates bzm' and tau' for 
        #fraction f throughout nongeoeffective events
        #note geoeff label can have values between 0-3. Anything that is not 1 is
        #considered nongeoeffective

        #gbzm_n = events_frac.query('geoeff == 0.0 and frac >= ' + str(fracs[i-1]) + ' and frac < ' + str(fracs[i])).bzm
        #gtau_n = events_frac.query('geoeff == 0.0 and frac >= ' + str(fracs[i-1]) + ' and frac < ' + str(fracs[i])).tau
        
        gbzmp_n = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm_predicted
        gtaup_n = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau_predicted
        
        #print("max bzmp_n: "+str(gbzmp_n.max())+", taup_n: "+ str(gtaup_n.max()))
        
        gbzmp_n.dropna(inplace=True)
        gtaup_n.dropna(inplace=True)

        if gtaup_n.max() == np.inf:

            inf_idx = gtaup_n[gtaup_n == np.inf].index.values
            gtaup_n.drop(inf_idx, inplace = True)
            gbzmp_n.drop(inf_idx, inplace = True)
 
        #to handle boundary conditions and limit the density estimate to positve 
        #values of tau': reflect the data points along the tau' axis, perform the 
        #density estimate and then set negative values of tau' to 0 and double the 
        #density of positive values of tau'
        gbzmp_nr = np.concatenate([gbzmp_n, gbzmp_n])
        gtaup_nr = np.concatenate([gtaup_n, -gtaup_n])
    
        #make sure the resulting PDF tau axis will start at 0 hours. 
        dt0 = int(len(Y_taup_n[1])/2.)
        Y_taup_n = Y_taup_n-((tmax-tmin)/len(Y_taup_n[1])/2.)    
        
        #option to use scikit learn or scipy stats kernel algorithms
        if kernel_alg == 'sklearn':
        
            positions = np.vstack([X_bzmp_n.ravel(), Y_taup_n.ravel()]).T
            values = np.vstack([gbzmp_nr, gtaup_nr]).T
            kernel_bzmp_taup_n = KernelDensity(kernel='gaussian', bandwidth=nw).fit(values)
            
            Ptmp_bzmp_taup_n[:,:,i-1] = np.exp(np.reshape(kernel_bzmp_taup_n.score_samples(positions).T, X_bzmp_n.shape))
            
        elif kernel_alg == 'scipy_stats'  :
            
            positions = np.vstack([X_bzmp_n.ravel(), Y_taup_n.ravel()])
            values = np.vstack([gbzmp_nr, gtaup_nr])  
            kernel_bzmp_taup_n = stats.gaussian_kde(values, bw_method=nw)

            Ptmp_bzmp_taup_n[:,:,i-1] = np.reshape(kernel_bzmp_taup_n(positions).T, X_bzmp_n.shape)

   
    #set the density estimate to 0 for negative tau, and x2 for positve tau 
    P_bzmp_taup_n = Ptmp_bzmp_taup_n[:,dt0::,:]*2

    #check the normalization
    norm_bzmp_taup_n = integrate.simps(integrate.simps(P_bzmp_taup_n[:,:,-1],Y_taup_n[0,dt0::]), X_bzmp_n[:,0])
    print('\n\n Normalization for P_bzmp_taup_n: ' + str(norm_bzmp_taup_n) + '\n\n')

    if plotfig == 1:
        
        fontP = FontProperties()                #legend
        #fontP.set_size('medium')1
        
        fig, ax = plt.subplots()
        c = ax.imshow(np.rot90(P_bzmp_taup_n[:,:,-1]), extent=(bmin,bmax,taumin,tmax), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax.plot(gbzmp_n, gtaup_n, 'k.', markersize=4, label = 'bzm_p, tau_p, geoeff = 0')
        ax.set_xlim([bmin, bmax])
        ax.set_ylim([taumin, tmax])
        ax.set_xlabel('Bzm')
        ax.set_ylabel('Tau')
        ax.set_title('P_bzmp_taup_n, bandwidth = '+str(nw))
        fig.colorbar(c)
        ax.legend(loc='upper right', prop = fontP, fancybox=True)

    return P_bzmp_taup_n, norm_bzmp_taup_n


def P_bzmp_taup_bzm_tau_e(events_frac, kernel_alg = 'scipy_stats', \
                ranges = [-150, 150, -250, 250], nbins = [50j, 100j],\
                fracs = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0], kernel_width = 2, plotfig = 0):  
    
    """
    Determine the prior PDF P(Bzm', tau'|(Bzm, tau) n e ; f), the probability 
    of a geoeffective event with estimates Bzm' and tau' for a MC with actual 
    values Bzm and tau, at fraction f throughout an event for input into the following 
    Chen geoeffective magnetic cloud prediction Bayesian formulation. This is 
    the bayesian likelihood PDF, relating the model to the data.  
    
    P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
    
    Due to a relatively small smaple of geoeffective events, kernel density estimation
    is used to smooth the data and generate a non parametric PDF. An artifact of 
    the KDE method is that the PDF will be smoothed to produce probabilities 
    extending to negative values of tau. To prevent this effect at the boundary 
    the raw data values are reflected in the tau' = 0 axis and then the KDE is applied, 
    producing a symmetric density estimates about the axis of reflection - see 
    Silverman 1998 section 2.10. The required output PDF is obtained by selecting array 
    elements corresponding to postive values of tau 
    
    f'(x) = 2f(x) for x >= 0 and f'(x) = 0 for x < 0
    
    As a result of this reflection, the input range for tau extends to negative 
    values and nbins_tau is double nbins_bzm. The output PDF will be for 
    tau' = [0, tmax]
    
    inputs:
        
    events_frac = dataframe
        contains output variables from a fit to solar wind magnetic field data
    kernel_alg = string
        choose between scikit learn and scipy stats python KDE algorithms
    ranges = 4 element array 
        defines the axis ranges for Bzm and tau [bmin, bmax, tmin, tmax]
    nbins = 2 elements array 
        defines the number of bins along bzm and tau [nbins_Bzm, nbins_tau]
    fracs = list
        fractions to split the events into for calculating evolution dependent pdfs
    kernel_width = float
        defines the kernel smoothing width
    plotfit = int
        plot a figure of the distribution 
        
    
    """ 
    #range of Bzm and tau to define PDFs 
    bmin = ranges[0]
    bmax = ranges[1]
    tmin = ranges[2]
    tmax = ranges[3]
    
    #number of data bins in each dimension (dt takes into account the reflection)
    db = nbins[0]
    dt = nbins[1]
    
    #true boundary for tau and the corresponding index in the pdf array
    taumin = 0
    #dt0 = dt/2 
    
    #kernel smoothing width
    ew = kernel_width
    
    #set grid containing x,y positions and use to get array size
    X_bzmp, Y_taup, XX_bzm, YY_tau = np.mgrid[bmin:bmax:db, tmin:tmax:dt, bmin:bmax:db, tmin:tmax:dt]
    db2 = int(len(X_bzmp[:,0,:,:]))
    dt2 = int(len(Y_taup[0,:,:,:]))
    
    #P_bzmp_taup_bzm_tau_e is a function of the fraction of time f throughout an event
    #currently the fit to the data considers every 5th of an event 
    nfracs = len(fracs)-1
    Ptmp_bzmp_taup_bzm_tau_e = np.zeros((db2,dt2,db2,dt2,nfracs))
    for i in range(1, len(fracs)):
    # TODO:
    #for i in [1.0]: 
        
        #extract raw data points from dataframe of estimates bzm and tau for 
        #fraction f throughout eoeffective events
        gbzm = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm
        gtau = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau
        
        #extract raw data points from dataframe of estimates bzm' and tau' for 
        #fraction f throughout eoeffective events
        gbzmp = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm_predicted
        gtaup = events_frac.query('geoeff == 1.0 and bzm < 0.0  and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau_predicted
        
        if gtaup.max() == np.inf:
            inf_idx = gtaup[gtaup == np.inf].index.values[0]
            gtaup.drop([inf_idx], inplace = True)
            gbzmp.drop([inf_idx], inplace = True)
            gtau.drop([inf_idx], inplace = True)
            gbzm.drop([inf_idx], inplace = True)
                
        gbzmp_n = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).bzm_predicted
        gtaup_n = events_frac.query('geoeff == 0.0 and frac_est >= ' + str(fracs[i-1]) + ' and frac_est < ' + str(fracs[i])).tau_predicted        

        if gtaup_n.max() == np.inf:
            inf_idx = gtaup_n[gtaup_n == np.inf].index.values[0]
            gtaup_n.drop([inf_idx], inplace = True)
            gbzmp_n.drop([inf_idx], inplace = True)


        #print("max bzm: "+str(gbzm.max())+", tau: "+ str(gtau.max()))
        #print("max bzmp: "+str(gbzmp.max())+", taup: "+ str(gtaup.max()))
        #print("max bzmp_n: "+str(gbzmp_n.max())+", taup_n: "+ str(gtaup_n.max()))
    
        #to handle boundary conditions and limit the density estimate to positve 
        #values of tau' and tau: reflect the data points along the tau' and tau
        #axes, perform the ensity estimate and then set negative values of 
        #tau' and tau to 0 and x4 the density of positive values of tau'
        gbzmp_r = np.concatenate([gbzmp, gbzmp, gbzmp, gbzmp])
        gtaup_r = np.concatenate([gtaup, -gtaup, gtaup, -gtaup])        
        gbzm_r = np.concatenate([gbzm, gbzm, gbzm, gbzm]) 
        gtau_r = np.concatenate([gtau, gtau, -gtau, -gtau])
        
        #make sure the resulting PDF tau axis will start at 0 hours. 
        dt0 = int(len(Y_taup[1])/2)
        Y_taup = Y_taup-((tmax-tmin)/len(Y_taup[1])/2.)
        YY_tau = YY_tau-((tmax-tmin)/len(YY_tau[1])/2.)
        
        #option to use scikit learn or scipy stats kernel algorithms
        if kernel_alg == 'sklearn':
            
            positions = np.vstack([X_bzmp.ravel(), Y_taup.ravel(), XX_bzm.ravel(), YY_tau.ravel()]).T
            values = np.vstack([gbzmp_r, gtaup_r, gbzm_r, gtau_r]).T        
            kernel_bzmp_taup_bzm_tau_e = KernelDensity(kernel='gaussian', bandwidth=ew).fit(values)
            
            Ptmp_bzmp_taup_bzm_tau_e[:,:,:,:,i-1] = np.exp(np.reshape(kernel_bzmp_taup_bzm_tau_e.score_samples(positions).T, X_bzmp.shape))
            
        elif kernel_alg == 'scipy_stats':    
            
            positions = np.vstack([X_bzmp.ravel(), Y_taup.ravel(), XX_bzm.ravel(), YY_tau.ravel()])
            values = np.vstack([gbzmp_r, gtaup_r, gbzm_r, gtau_r])  
            
            try:
                kernel_bzmp_taup_bzm_tau_e = stats.gaussian_kde(values, bw_method = ew)
                Ptmp_bzmp_taup_bzm_tau_e[:,:,:,:,i-1] = np.reshape(kernel_bzmp_taup_bzm_tau_e(positions).T, X_bzmp.shape)
            except Exception as e:
                print(e.__doc__)
                #print(e.message)
                break
            
 
    #set the density estimate to 0 for negative tau, and x4 for positve tau 
    P_bzmp_taup_bzm_tau_e = Ptmp_bzmp_taup_bzm_tau_e[:,dt0::,:,dt0::,:]*4       #*50*50*2    

    #check the normalization of the 4D space   
    bp = X_bzmp[:,0,0,0]
    tp = Y_taup[0,dt0::,0,0]
    b = XX_bzm[0,0,:,0]
    t = YY_tau[0,0,0,dt0::]
    
    indices = np.squeeze(np.array([[bp],[tp],[b],[t]]))
    
    norm_bzmp_taup_bzm_tau_e = integrate.simps(integrate.simps(integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e[:,:,:,:,-1],
                                tp),\
                                bp), \
                                t),\
                                b)
    
    #check the normalization of the 4D space - should be 1   
    predicted_duration = 15.0
    predicted_bzmax = -26.0
    indt = np.min(np.where(t > predicted_duration))
    indb = np.max(np.where(b < predicted_bzmax))
    
    
    P0 = integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e[:,:,indb,indt,-1],\
                                t),\
                                b)  
    
    print('\n\n Normalization for P_bzmp_taup_bzm_tau_e: ' + str(norm_bzmp_taup_bzm_tau_e) )
    print('\n Normalization for P_bzmp_taup_bzm_tau_e plane Bzm =-26nT, tau = 15 hrs, frac = 1.0: ' \
          + str(P0) + '\n\n')
                                
    if plotfig == 1:
        
        fontP = FontProperties()                #legend
        #fontP.set_size('medium')1
        
        fig, ax = plt.subplots()
        c = ax.imshow(np.rot90(P_bzmp_taup_bzm_tau_e[:,:,indb,indt,0]), extent=(bmin,bmax,taumin,tmax), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        ax.plot(gbzmp_n, gtaup_n, 'k.', markersize=4, c='b')
        ax.plot(gbzmp, gtaup, 'k.', markersize=4, c='r')
        ax.set_xlim([bmin, bmax])
        ax.set_ylim([taumin, tmax])
        ax.set_xlabel('Bzm')
        ax.set_ylabel('Tau')
        fig.colorbar(c)

    
    
    ############## FOLLOWING CODE IS MESSING ABOUT TO GET THE PLANES TO NORM TO 1 #############

    P0_2 = np.zeros(P_bzmp_taup_bzm_tau_e[:,:,indb,indt,:].shape)
    tmpint = np.zeros(P_bzmp_taup_bzm_tau_e[:,:,indb,indt,:].shape)
    invtmpint = np.zeros(P_bzmp_taup_bzm_tau_e[:,:,indb,indt,:].shape)
    P_bzmp_taup_bzm_tau_e2 = np.zeros(P_bzmp_taup_bzm_tau_e.shape)

    
    for i in range(nfracs):
    #for i in [5]:                          #note only calculating final dimension/time step in above code for quickness testing
        for j in range(len(P0_2[:,0])):
            for k in range(len(P0_2[0,:])):
                
                tmp = P_bzmp_taup_bzm_tau_e[:,:,j,k,i]
                tmpint[j,k,i] = integrate.simps(integrate.simps(tmp, t), b)  
                invtmpint[j,k,i] = 1/tmpint[j,k,i] 
                
                P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = tmp * (1/tmpint[j,k,i])
                
                #normalization in each plane - should be 1
                P0_2[j,k,i] = integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e2[:,:,j,k,i],\
                               t),\
                               b)   
    #total 4D normalization
    norm_bzmp_taup_bzm_tau_e2 = integrate.simps(integrate.simps(integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e2[:,:,:,:,-1],
                                tp),\
                                bp), \
                                t),\
                                b)


#==============================================================================
#     #normalize each plane
#     x = 0
#     
#     P_bzmp_taup_bzm_tau_e2 = np.zeros((50,50,50,50,6))
#     for i in np.arange(6):
#         for j in range(50):
#             for k in range(50):
#                 
#                 tmp = P_bzmp_taup_bzm_tau_e[:,:,j,k,i]
#                 
#                 #x += np.sum(P_bzmp_taup_bzm_tau_e[:,:,j,k,i])
#                 
#                 #tmpsum = integrate.simps(integrate.simps(tmp, t), b)   
#                 #P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = tmp/tmpsum/((b[1]-b[0]) * (t[1]-t[0]))
#                 #P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = tmp *(50*(b[1]-b[0]) * 50*(t[1]-t[0]))
#                 
#                 #tmpsum = np.sum( P_bzmp_taup_bzm_tau_e[:,:,j,k,i])
#                 
#                 #tmpsum = integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e[:,:,j,k,i], t), b)   
#                 
#                 #P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = tmp* (50*(b[1]-b[0]) * 50*(t[1]-t[0])) *tmpsum 
#                 
#                 P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = tmp *2*50*50
#                 
#                 #den = 0
#                 #for m in range(50):
#                 #    for n in range(50):
#                 #        #den +=  (b[1]-b[0]) * (t[1]-t[0]) * P_bzmp_taup_bzm_tau_e[m,n,j,k,i]
#                 #        den +=  P_bzmp_taup_bzm_tau_e[m,n,j,k,i]
#                         
#                 #tmpsum = np.sum( P_bzmp_taup_bzm_tau_e[:,:,j,k,i])
#                 
#                 #P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = tmp / tmpsum / (50*5.05*50*6.12)
# 
#                 
#                 #P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = tmp*50*50*2
#                 
#                 #if np.isnan(integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e2[:,:,j,k,i], t), b)):
#                     #print(j,k,i,integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e2[:,:,j,k,i],\
#                     #            t),\
#                     #            b) )
# 
#     #voxel = (49*(b[1]-b[0]) * 49*(t[1]-t[0]))
#     #voxel = 300*250
#     #P_bzmp_taup_bzm_tau_e2 = P_bzmp_taup_bzm_tau_e * voxel
#     
#     #P_bzmp_taup_bzm_tau_e2 = P_bzmp_taup_bzm_tau_e * (49*49)
#==============================================================================
    
#==============================================================================
#     P_bzmp_taup_bzm_tau_e2 = np.zeros((db2,db2,db2,db2,6))
#     for j in range(db2):
#         for k in range(db2):
#             for i in range(6):
#                 P_bzmp_taup_bzm_tau_e2[:,:,j,k,i] = P_bzmp_taup_bzm_tau_e[:,:,j,k,i] * 1/ P_bzmp_taup_bzm_tau_e[:,:,j,k,i]
#==============================================================================
    
    #P0_2 = integrate.simps(integrate.simps(P_bzmp_taup_bzm_tau_e2[:,:,20,3,5],\
    #                            t),\
    #                            b)          
    #print(P0_2)            
    
    ############## END MESSING ABOUT WITH CODE #############
    
    return P_bzmp_taup_bzm_tau_e2, norm_bzmp_taup_bzm_tau_e2, P0_2, tmpint, indices


def P_bzm_tau_e_bzmp_taup(P_e, P_n, P_bzm_tau_e, P_bzmp_taup_e, P_bzmp_taup_n, \
                P_bzmp_taup_bzm_tau_e, ranges = [-150, 150, -250, 250], \
                nbins = [50j, 100j], fracs = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0], plotfig = 0):  
    
    """
    Determine the posterior PDF P((Bzm, tau) n e |Bzm', tau' ; f), the probability 
    of a geoeffective event with parameters Bzm and tau given estimates  
    Bzm' and tau', at fraction f throughout an event for input into the following 
    Chen geoeffective magnetic cloud prediction Bayesian formulation. This is 
    the bayesian posterior PDF.
    
   P((Bzm, tau) n e|Bzm', tau' ; f) 
    = P(Bzm, tau|(Bzm, tau) n e ; f) * P(Bzm, tau|e) * P(e) / SUM_j( P(Bzm',tau' | c_j ; f) * P(c_j))
    
    where
    
    Bzm = actual value of Bz max for a magnetic cloud
    tau = actual duation of a magnetic cloud
    Bzm' = fitted/estimated value of Bz max at fraction (f) of an event
    tau' = fitted/estimated value of duration at fraction (f) of an event
    e = geoeffective event
    n = nongeoeffective event
    f = fraction of an event
       
    inputs:
    
    P_e = float
        Prior P(e)
    P_n = float
        Prior P(n)
    P_bzm_tau_e = data array
        Prior PDF P(bzm, tau|e)
    P_bzmp_taup_e = data array
        Evidence PDF P(bzm', tau'|e; f)
    P_bzmp_taup_n = data array
        Evidence PDF P(bzm', tau'|n; f)
    P_bzmp_taup_bzm_tau_e = data array
        Likelihood PDF P(Bzm', tau' | (Bzm, tau) n e ; f)

    ranges = 4 element array 
        defines the axis ranges for Bzm and tau [bmin, bmax, tmin, tmax]
    nbins = 2 elements array 
        defines the number of bins along bzm and tau [nbins_Bzm, nbins_tau]
    fracs = list
        fractions to split the events into for calculating evolution dependent pdfs
    plotfit = int
        plot a figure of the distribution     

    
    output:
        
    P_bzm_tau_e_bzmp_taup = data array
        Posterior PDF P((Bzm, tau) n e | Bzm', tau' ; f)
    
        
    """ 
    #range of Bzm and tau to define PDFs 
    bmin = ranges[0]
    bmax = ranges[1]
    tmin = ranges[2]
    tmax = ranges[3]
    
    #number of data bins in each dimension (dt takes into account the reflection)
    db = nbins[0]
    dt = nbins[1]
    
    #set grid containing x,y positions and use to get array size
    X_bzmp, Y_taup, XX_bzm, YY_tau = np.mgrid[bmin:bmax:db, tmin:tmax:dt, bmin:bmax:db, tmin:tmax:dt] 
    dt0 = int(len(Y_taup[1])/2.)
    db2 = int(len(X_bzmp[:,0,:,:]))
    dt2 = int(len(Y_taup[0,:,:,:])/2)
    
    #true boundary for tau and the corresponding index in the pdf array
    taumin = 0
    #dt0 = dt/2 
    
    #P_bzm_tau_e_bzmp_taup_e is a function of the fraction of time f throughout an event
    nfracs = len(fracs)-1    
    P_bzm_tau_e_bzmp_taup = np.zeros((db2,dt2,dt2,db2,nfracs))
    for i in range(1, len(fracs)):
    #for i in np.arange(6)*0.2:

        num = np.multiply(P_bzmp_taup_bzm_tau_e[:,:,:,:,i-1], P_bzm_tau_e) * P_e
        denom = (P_bzmp_taup_e[:,:,i-1] * P_e) + (P_bzmp_taup_n[:,:,i-1] * P_n)
        
        P_bzm_tau_e_bzmp_taup[:,:,:,:,i-1] = np.divide(num, denom)
        
        #num = np.multiply(Ptmp_bzmp_taup_bzm_tau_e[:,:,:,:,int(i*5)], Ptmp_bzm_tau_e) * P_e
        #denom = (Ptmp_bzmp_taup_e[:,:,int(i*5)] * P_e) + (Ptmp_bzmp_taup_n[:,:,int(i*5)] * P_n)
        #Ptmp_bzm_tau_e_bzmp_taup[:,:,:,:,int(i*5)] = np.divide(num, denom)

    #check the normalization of the 4D space 
    bp = X_bzmp[:,0,0,0]
    tp = Y_taup[0,dt0::,0,0]
    b = XX_bzm[0,0,:,0]
    t = YY_tau[0,0,0,dt0::]
    
    axis_vals = np.array([bp,tp,b,t])
    

    norm_bzm_tau_e_bzmp_taup = integrate.simps(integrate.simps(integrate.simps(integrate.simps(P_bzm_tau_e_bzmp_taup[:,:,:,:,-1],\
                                tp),\
                                bp), \
                                t),\
                                b)
    #check the normalization of the 4D space - should be 1   
    # corresponds to P1 in Chen 97 paper
    b = XX_bzm[0,0,:,0]
    t = YY_tau[0,0,0,dt0::]
    
    
    predicted_duration = 15.0
    predicted_bzmax = -26.0
    indt = np.min(np.where(t > predicted_duration))
    indb = np.max(np.where(b < predicted_bzmax))
    
    #ranget1 =  np.min(np.where(t > 80.0))
    #rangeb0 =  np.max(np.where(b < -40.0))
    #rangeb1 =  np.min(np.where(b > 40.0))
    
    P1 = integrate.simps(integrate.simps(P_bzm_tau_e_bzmp_taup[:,:,indb,indt,-1],\
                                t),\
                                b)

    #map of P1 for all planes
    P1_map = np.zeros((db2,db2,nfracs))

    for i in range(nfracs):
        for j in range(db2):
            for k in range(db2):
                P1_map[j,k,i] = integrate.simps(integrate.simps(P_bzm_tau_e_bzmp_taup[:,:,j,k,i],\
                                    t),\
                                    b)
    
    print('\n\n Normalization for P_bzm_tau_e_bzmp_taup: ' + str(norm_bzm_tau_e_bzmp_taup) )
    print('\n Normalization for P_bzm_tau_e_bzmp_taup plane Bzmp =-26nT, taup = 15 hrs, frac = 1.0: ' \
          + str(P1) + '\n\n')
    print('\n\n Max P1_map: ' + str(P1_map.max()) )
    
    if plotfig == 1: 
        
        fontP = FontProperties()                #legend
        #fontP.set_size('medium')1
        
        fig, ax = plt.subplots()
        c = ax.imshow(np.rot90(P_bzm_tau_e_bzmp_taup[:,:,indb,indt,i-1]), extent=(bmin,bmax,taumin,tmax), cmap=plt.cm.gist_earth_r, interpolation = 'none')
        #ax.imshow(np.rot90(num[:,:,20,3]), extent=(bmin,bmax,taumin,tmax), cmap=plt.cm.gist_earth_r)
        ax.set_xlim([bmin, bmax])
        ax.set_ylim([taumin, tmax])
        ax.set_xlabel('Bzm')
        ax.set_ylabel('Tau')
        fig.colorbar(c)


    return P_bzm_tau_e_bzmp_taup, norm_bzm_tau_e_bzmp_taup, P1, P1_map, axis_vals   