import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from multiprocessing import Process, Pool


# Finite illuminated volume case
def finite_volume(focal_distance,
                  focal_lengths,
                  concentrations,
                  particle_area=np.pi*(.030/2)**2,
                  z_depth=100,
                  blur_factor=3,
                  stereo_baseline=30,
                  sensor_size=(1920, 1080, 0.0045),
                  f_number=2.8,
                  particle_contrast=0.5):

    pixel_size = sensor_size[2]  # mm
    apertures = focal_lengths / f_number  # mm
    blur_angles = np.arctan(apertures / focal_distance / 2)

    # output data
    results = np.zeros((len(blur_angles), len(concentrations)))
    for row, b in enumerate(blur_angles):

        dz = z_depth
        # dx and dy from focal length and sensor size at focal distance
        dx = 2 * (focal_distance + dz / 2) * np.tan(np.arctan(sensor_size[0] * pixel_size / (2 * focal_lengths[row])))
        dy = 2 * (focal_distance + dz / 2) * np.tan(np.arctan(sensor_size[0] * pixel_size / (2 * focal_lengths[row])))


        print('Focal Length: ' + str(focal_lengths[row]) + ', ' + 'Volume: ' + str([dx, dy, dz]))

        for col, con in enumerate(concentrations):
            print('Concentration: ' + str(con))
            # numer of particles from volume and concentration
            n_particles = int(dx * dy * dz * con)
            z_vals = focal_distance - dz / 2 + dz * np.random.random(n_particles)
            # x and y vals after camera projection
            x_rand = -dx / 2 + dx * np.random.random(n_particles)
            y_rand = -dy / 2 + dy * np.random.random(n_particles)
            x_vals = focal_lengths[row] * x_rand / z_vals
            y_vals = focal_lengths[row] * y_rand / z_vals

            x_vals_l = focal_lengths[row] * (x_rand-stereo_baseline) / z_vals
            y_vals_l = focal_lengths[row] * (y_rand) / z_vals

            # contrast estimate
            particle_dist = np.sqrt(x_vals ** 2 + y_vals ** 2)
            blur_rad = np.abs(z_vals - focal_distance) * np.tan(b)
            overlapping_signals = (blur_rad >= particle_dist).astype('uint8')
            blur_area = np.pi * blur_rad ** 2
            blur_signal = particle_contrast * particle_area / (blur_area + particle_area)
            ideal_contrast = particle_contrast / blur_factor
            if particle_contrast > 0:
                results[row, col] = (ideal_contrast - np.sum(blur_signal * overlapping_signals)) / ideal_contrast
            else:
                results[row, col] = 0

            scale = 10

            # stereo 3D position error estimate
            x_pix = np.round(scale*x_vals/pixel_size)
            y_pix = np.round(scale*y_vals/pixel_size)
            x_pix_l = np.round(scale*x_vals_l/pixel_size)
            y_pix_l = np.round(scale*y_vals_l/pixel_size)

            avg_error = 0
            for ind, p in enumerate(x_pix):
                err = np.abs(y_pix[ind] - y_pix_l)
                err_sorted = sorted(err)
                match_ind = 0
                while True:

                    zh = focal_lengths[row] / pixel_size * stereo_baseline * scale / (x_pix[ind] - x_pix_l[match_ind])
                    if zh > 450 and zh < 550:
                        break
                    match_ind += 1
                match_ind = np.argwhere(err == np.min(err))[0][0]
                zh = focal_lengths[row]/pixel_size*stereo_baseline*scale/(x_pix[ind]-x_pix_l[match_ind])
                xh = x_pix[ind] * zh / (focal_lengths[row] / pixel_size)/scale
                yh = y_pix[ind] * zh / (focal_lengths[row] / pixel_size)/scale
                reproj_error = np.sqrt((zh-z_vals[ind])**2 + (xh-x_rand[ind])**2 + (yh-y_rand[ind])**2)
                avg_error += reproj_error

            avg_error /= len(x_pix)
            print(avg_error)



    return results


if __name__ == "__main__":

    # Simulation settings
    concentrations = np.logspace(-5, 0, num=50)  # #/mm^3
    focal_lengths = np.array([25, 50, 100, 200])  # mm
    focal_distance = 500  # mm

    particle_contrast = 0.5

    results = np.zeros((len(focal_lengths), len(concentrations)))
    iterations = 50
    for i in range(0, iterations):
        print('iteration: ' + str(i))
        results += finite_volume(focal_distance, focal_lengths, concentrations)
    results /= iterations

    fig, ax = plt.subplots()
    ax.semilogx(concentrations * 1000, np.transpose(results))
    plt.xlabel('Particle Concentration (#/ml)')
    plt.ylabel('Relative Contrast')
    plt.legend([str(x) + ' mm' for x in focal_lengths])
    plt.ylim([.1, 1])
    plt.grid(True)
    plt.show()