import numpy as np
import matplotlib.pyplot as plt

z0 = 0.45
zf = 0.55
zmid = np.mean([z0, zf])
hist_bins = range(20, 2000, 5)
scale = 10
pixel_size = 0.001
focal_length = 21/pixel_size
samples = 10000

# 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()

size_classes = [50, 100, 200, 400, 800, 1600, 3200]

distros = []

sizes_true = []
sizes_est = []
sizes_theta = []
sizes_theta_est = []
for i in range(samples):
    ind = int(len(size_classes)*np.random.random())
    size_true = np.random.normal(size_classes[ind], scale)

    z_val = z[i]

    proj_size = focal_length*size_true/(z_val*1000000)
    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)))
    sizes_est.append(proj_size*zmid*1000000/focal_length)
    sizes_theta_est.append(proj_size_theta*zmid*1000000/focal_length)

plt.hist(sizes_true, bins=hist_bins)
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(['True size','2D estimate'])
#plt.legend(['True size with unknown orientation and 0.5 aspect ratio','2D estimate'])  
plt.show()