import numpy as np
import powerlaw
import matplotlib.pyplot as plt

z0 = 0.45
zf = 0.55
zmid = np.mean([z0, zf])
hist_bins = np.linspace(.000001, .01, 200)
scale = 10
pixel_size = 0.001
focal_length = 21/pixel_size
samples = 1000000

# random samples from quadratic pdf of z
r = np.random.random(samples)
z = (r*(zf**3 - z0**3) + z0**3)**(1/3)
# pdf(z) = 3/(zf^3 - z0^3)*z^2
# CDF(z) = (z^3 - z0^3) / (zf^3 - z0^3)
# CDF^-1(r) = (r*(zf^3 - z0^3) + z0^3) ^(1/3)
#plt.hist(z, bins=100)
#plt.show()

slopes = [.7, .8, .9, .99]

distros = []

sizes_true = []
sizes_est = []
sizes_theta = []
sizes_theta_est = []
plot_colors = 'brgk'
xmin = 0.00001
for i in range(len(slopes)):

    #size_true = np.random.power(slopes[i], samples)
    #y = np.random.random(samples)
    #size_true = xmin * (1-r) ** (-1/(slopes[i]-1))
    x1 = 10000
    x0 = 20
    n = slopes[i]
    #size_true = ((x1 ** (n + 1) - x0 ** (n + 1)) * y + x0 ** (n + 1)) ** (1 / (n + 1))
    theoretical_distribution = powerlaw.Power_Law(xmin=.01, parameters=[slopes[i]])
    size_true = theoretical_distribution.generate_random(samples)


    r = np.random.random(samples)
    z = (r * (zf ** 3 - z0 ** 3) + z0 ** 3) ** (1 / 3)
    #z_val = z[i]

    proj_size = focal_length*size_true/(z*1000000)
    print(proj_size.shape)
    #proj_size_theta = proj_size/2*(1 + np.cos(np.random.random()*np.pi/2))

    #sizes_true.append(size_true)
    #sizes_theta.append(size_true/2*(1 + np.cos(np.random.random()*np.pi/2)))
    size_est = proj_size*zmid*1000000/focal_length
    #sizes_theta_est.append(proj_size_theta*zmid*1000000/focal_length)

    n, binc = np.histogram(size_true, bins=hist_bins, density=True)
    plt.loglog(1000000*binc[1:], n[0:], plot_colors[i] + '-')
    n, binc = np.histogram(size_est, bins=hist_bins, density=True)
    plt.loglog(1000000*binc[1:], n[0:], plot_colors[i] + '.')
    #plt.plot(binc)

    #plt.semilogy(size_true)

    #results = powerlaw.Fit(size_true)
    #print(results.power_law.alpha)
    #print(results.power_law.xmin)

    #plt.hist(sizes_est, bins=hist_bins)

#plt.hist(sizes_theta, bins=hist_bins)
#plt.hist(sizes_theta_est, bins=hist_bins)
plt.grid(True)
plt.xlabel('Particle Size (um)')
plt.ylabel('Counts')
plt.legend(['slope=-.7', '2D Estimate', 'slope=--.8', '2D Estimate', 'slope=-.9', '2D Estimate', 'slope=-.99', '2D Estimate'])
#plt.legend(['True size with unknown orientation and 0.5 aspect ratio','2D estimate'])  
plt.xlim([100, 9000])
plt.show()