import pymc3 as pm
import seaborn as sns
import scipy
import scipy.stats
import numpy as np
import spacepy.toolbox as tb
import matplotlib.pyplot as plt
from theano import theano, tensor as tt

sns.set(font_scale=1.5)

"""
Explanation: Demonstrate how to use DensityDist
Utilize a scipy.stats.rayleigh distribution to demonstrate
End of explanation
"""


from pymc3.distributions.dist_math import bound

def pot1(z):
    z = z.T
    return .5*((z.norm(2, axis=0)-2.)/.4)**2 - tt.log(tt.exp(-.5*((z[0]-2.)/.6)**2) +
                                                      tt.exp(-.5*((z[0]+2.)/.6)**2))


def cust_logp(z):
    # this works to cut things off, useful or not but 0 is a nature place a lot of times
    return bound(-pot1(z), z>-5, z<5) 
#     return -pot1(z)

with pm.Model() as pot1m:
    pm.DensityDist('pot1', logp=cust_logp, shape=(2,))
    trace = pm.sample(1000)

pm.traceplot(trace);

pm.plot_posterior(trace);

"""
Explanation: Given example
https://pymc3.readthedocs.io/en/latest/notebooks/normalizing_flows_overview.html?highlight=densitydist
End of explanation
"""


dist = scipy.stats.rayleigh(loc=2, scale=20)
rvs = dist.rvs(10000)
sns.distplot(rvs)

# need to make a callable function of the histogram of this. 
h, b = np.histogram(rvs, 30)
b = tb.bin_edges_to_center(b)
plt.plot(b, h)

f = scipy.interpolate.interp1d(b,h, bounds_error=False, fill_value=0)
XX = np.linspace(0, 100, 1000)
plt.plot(XX, f(XX))


from theano.tensor import _shared

# @theano.as_op(itypes=[theano.tensor.dscalar],otypes=[theano.tensor.dscalar])
def dens_fn(x):
    return tt.log(f(x) )
# dens_fn(_shared(np.linspace(-1, 5, 10)))

with pm.Model() as model:
    dens = pm.DensityDist('density', logp=dens_fn)
    trace = pm.sample(5000)
    

pm.traceplot(trace);

pm.plot_posterior(trace);

"""
Explanation: Now can we do this from a histogram?
End of explanation
"""


