import numpy as np
import random
from scipy.stats import rv_continuous
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit


'''
Mass distribution:

assume m can go from 1 solar mass to infinite solar Mass

Salpeter distribution:
epsilon(m)dm = epsilon_0*(m/M_solar)^(-2.35)*(dm/M_solar),
epsilon_0 is a constant relating to the local stellar density

set m = m/M_solar, then p(m)dm = C*m^(-2.35)*dm
Normalize going from m = 1 to infinity, get C = 1.35

So the pdf = 1.35*m^(-2.35)  (check this)

how can I make this more efficient?
generating same random number each time
'''
'''
num_events=10000
y_range = 100
grad=3
xbins = np.linspace(1,y_range,grad*(y_range))


class Mass_pdf(rv_continuous):
    def _pdf(self, x):
        return 1.35*x**(-2.35)

#random distribution
mass_pdf = Mass_pdf(name="Mass distribution", a=1)
x=np.array([np.random.rand(num_events)])
y_rand=mass_pdf.ppf(x)[0]

#actual distribution
y_act= 1.35*np.power(xbins,np.full((1,grad*(y_range)),-2.35))
y_act=y_act[0]

#FIGURE 1 - LINEAR X AXIS
fig1 = plt.figure(1)
ax1 = fig1.add_subplot(111)
plt.hist(y_rand,normed=True,bins=xbins,label="probability")
plt.plot(xbins,y_act,label="expected function")
fig1.suptitle("Expected Mass Distribution for BH's")
plt.ylabel("p")
plt.xlabel("mass (M_solar)")

#FIGURE 2 - LOG Y AXIS
fig5 = plt.figure(5)
ax5 = fig5.add_subplot(111)
plt.hist(y_rand,normed=True,bins=xbins,label="probability")
plt.plot(xbins,y_act,label="expected function")
fig5.suptitle("Expected Mass Distribution for BH's")
plt.ylabel("p")
plt.yscale('log', nonposy='clip')
plt.xlabel("mass (M_solar)")

#FIGURE 3 - LOG ALL AXIS
fig6 = plt.figure(6)
ax6= fig6.add_subplot(111)
logbins = 10 ** np.linspace(np.log10(1), np.log10(y_range), grad*y_range)
plt.hist(y_rand,normed=True,bins=logbins,label="probability")
plt.plot(xbins,y_act,label="expected function")
fig6.suptitle("Expected Mass Distribution for BH's")
plt.ylabel("p")
plt.gca().set_xscale("log")
plt.yscale('log', nonposy='clip')
plt.xlabel("mass (M_solar)")

# curve fitting
n, bins = np.histogram(y_rand,normed=True,bins=xbins)
bin_centers = bins[:-1] + 0.5 * (bins[1:] - bins[:-1])
def f(x, a, b):
    return a*x**b
popt, pcov = curve_fit(f, bin_centers, n, p0 = [1, -3])
print popt

text = "a*x^b a="+str(popt[0])+",b="+str(popt[1])+" vs. a=1.35,b=-2.35"

plt.figure(1)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax1.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.figure(5)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax5.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.figure(6)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax6.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.show()
'''

'''
Radial distribution:

Go from r=0 to infinity

Is this a poisson process??? Show him Exam 1

Else I would say pdf(r) = rho*4*pi*r^2, where rho is the density of stars

Need to normalize out to a horizen distance:

integral_0^horizon_dist C*pdf(r)dr = 1

Normalizing, C= (3/4)(1/pi)(1/rho)(1/horizon_dist)^3
'''
'''

#1 gpc look up how many within a GPC

num_events=10000
grad=100
density = 50000000

horizon_dist = 1


class Distance_pdf(rv_continuous):
    def _pdf(self, x):
        return 3.0*(x**2)*(1.0/horizon_dist)**3

xbins = np.linspace(0,horizon_dist,grad*(horizon_dist))

#random distribution
distance_pdf = Distance_pdf(name="Distance distribution", a=0,b=horizon_dist)
x=np.array([np.random.rand(num_events)])
y_rand=distance_pdf.ppf(x)[0]

#actual distribution
y_act = 3.0*(1.0/horizon_dist)**3*np.power(xbins,np.full((1,grad*(horizon_dist)),2))
y_act=y_act[0]

#FIGURE 2 - LINEAR X AXIS
fig2 = plt.figure(2)
ax2 = fig2.add_subplot(111)
plt.hist(y_rand,normed=True,bins=xbins,label="probability")
plt.plot(xbins,y_act,label="expected function")
fig2.suptitle("Expected Radial Distribution for BH's,normed")
plt.ylabel("p")
plt.xlabel("r (Gpc)")

#FIGURE 3 - LOG Y AXIS
fig3 = plt.figure(3)
ax3 = fig3.add_subplot(111)
plt.hist(y_rand,normed=True,bins=xbins,label="probability")
plt.plot(xbins,y_act,label="expected function")
fig3.suptitle("Expected Radial Distribution for BH's,normed")
plt.ylabel("p")
plt.yscale('log', nonposy='clip')
plt.xlabel("r (Gpc)")

#FIGURE 4 - LOG ALL AXES
fig4 = plt.figure(4)
ax4 = fig4.add_subplot(111)
logbins = 10 ** np.linspace(np.log10(1), np.log10(horizon_dist), grad*horizon_dist)
plt.hist(y_rand,normed=True,bins=xbins,label="probability")
plt.plot(xbins,y_act,label="expected function")
fig4.suptitle("Expected Radial Distribution for BH's,normed")
plt.ylabel("p")
plt.gca().set_xscale("log")
plt.yscale('log', nonposy='clip')
plt.xlabel("r (Gpc)")

# curve fitting
n, bins = np.histogram(y_rand,normed=True,bins=xbins)
bin_centers = bins[:-1] + 0.5 * (bins[1:] - bins[:-1])
def f(x, a, b):
    return a*x**b
popt, pcov = curve_fit(f, bin_centers, n, p0 = [1,1])
print popt

text = "a*x^b a="+str(popt[0])+",b="+str(popt[1])+" vs. a="+str(3.0*(1.0/horizon_dist)**3)+",b=2"

plt.figure(2)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax2.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.figure(3)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax3.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.figure(4)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax4.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.show()

'''




num_events=10000
grad=100
density = 50000000

horizon_dist = 1


class Distance_pdf(rv_continuous):
    def _pdf(self, x):
        return 3.0*(x**2)*(1.0/horizon_dist)**3

xbins = np.linspace(0,horizon_dist,grad*(horizon_dist))


#random distribution
distance_pdf = Distance_pdf(name="Distance distribution", a=0,b=horizon_dist)
x=np.array([np.random.rand(num_events)])
y_rand=distance_pdf.ppf(x)[0]
normed_value = (4.0/3)*np.pi*density*horizon_dist**3
hist, bins = np.histogram(y_rand, bins=xbins, density=True)
widths = np.diff(bins)
hist *= normed_value


#actual distribution
y_act = (4.0)*np.pi*density*np.power(xbins,np.full((1,grad*(horizon_dist)),2))
y_act=y_act[0]

#FIGURE 7 - LINEAR X AXIS
fig7 = plt.figure(7)
ax7 = fig7.add_subplot(111)
plt.bar(bins[:-1], hist, widths,label="probability",color="b")
plt.plot(xbins,y_act,label="expected function")
fig7.suptitle("Expected Radial Distribution for BH's")
plt.ylabel("p")
plt.xlabel("r (Gpc)")

#FIGURE 8 - LOG Y AXIS
fig8 = plt.figure(8)
ax8 = fig8.add_subplot(111)
plt.bar(bins[:-1], hist, widths,label="probability",color="b")
plt.plot(xbins,y_act,label="expected function")
fig8.suptitle("Expected Radial Distribution for BH's")
plt.ylabel("p")
plt.yscale('log', nonposy='clip')
plt.xlabel("r (Gpc)")

#FIGURE 9 - LOG ALL AXES
fig9 = plt.figure(9)
ax9 = fig9.add_subplot(111)
logbins = 10 ** np.linspace(np.log10(1), np.log10(horizon_dist), grad*horizon_dist)
plt.bar(logbins[:-1], hist, widths,label="probability",color="b")
plt.plot(xbins,y_act,label="expected function")
fig9.suptitle("Expected Radial Distribution for BH's")
plt.ylabel("p")
plt.gca().set_xscale("log")
plt.yscale('log', nonposy='clip')
plt.xlabel("r (Gpc)")

# curve fitting
bin_centers = bins[:-1] + 0.5 * (bins[1:] - bins[:-1])
def f(x, a, b):
    return a*x**b
popt, pcov = curve_fit(f, bin_centers, hist, p0 = [50000000,1])
print popt

text = "a*x^b a="+str(popt[0])+",b="+str(popt[1])+" vs. a="+str((4.0)*np.pi*density)+",b=2"

plt.figure(7)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax7.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.figure(8)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax8.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.figure(9)
plt.plot(bin_centers,f(bin_centers, *popt),label="fit function")
ax9.set_title(text,fontsize=9)
plt.legend(loc='upper right')

plt.show()






"""

distance to first BBH
"""

'''

y_range = 20
grad=40
rho=4

plt.figure(4)

class Distance_pdf(rv_continuous):
    def _pdf(self, x, rho):
        return (4*np.pi*rho*x**2)*(np.exp((-4*rho*np.pi*x**3)/3.0))

xbins = np.linspace(0,y_range,grad*(y_range))

#random distribution
distance_pdf= Distance_pdf(name="Distance distribution", a=0)
x=np.array([np.random.rand(num_events)])
y=distance_pdf.ppf(x,rho)[0]

plt.hist(y,normed=True,bins=xbins,label="probability")

#actual distribution
y=(4*np.pi*rho*np.power(xbins,np.full((1,grad*(y_range)),2)))*(np.exp((-4*rho*np.pi*np.power(xbins,np.full((1,grad*(y_range)),3)))/3.0))
y=y[0]
plt.plot(xbins,y,label="expected function")

plt.title("Expected Distance to first BBH")
plt.ylabel("p")
plt.xlabel("r")
plt.legend(loc='upper right')

plt.show()
'''