#### Principal retardance and 3D orientation in color

orientation_3D_image = np.transpose(np.array([azimuth%(2*np.pi)/2/np.pi, inclination,  np.clip((ne_map-no_map)*psz*2*np.pi/lambda_illu/np.pi/2*lambda_illu*1e3,0, 1.5)/1.5]),(3,1,2,0))
orientation_3D_image_RGB = wo.orientation_3D_to_rgb(orientation_3D_image, interp_belt = 20/180*np.pi, sat_factor = 1)

plt.figure(figsize=(10,10))
plt.imshow(orientation_3D_image_RGB[z_layer], origin='lower')
plt.figure(figsize=(10,10))
plt.imshow(orientation_3D_image_RGB[:,y_layer], origin='lower')
plt.figure(figsize=(3,3))
wo.orientation_3D_colorwheel(wheelsize=128, circ_size=50, interp_belt=20/180*np.pi, sat_factor=1,  discretize=True)
plt.show()

#### Angular histogram of 3D orientation
wo.orientation_3D_hist(azimuth.flatten(), \
                       inclination.flatten(), \
                       np.abs(target).flatten(),\
                       bins=36, num_col=1, size=10, contour_level = 100, hist_cmap='gnuplot2', top_hemi=True)
plt.show()


## Compute permittivity and scattering potential tensors
### assign the permittivity tensor according to the sample properties

epsilon_mean = (ne_map**2 + no_map**2)/2
epsilon_del = (ne_map**2 - no_map**2)/2

epsilon_tensor = np.zeros((3,3,N,M,L))
epsilon_tensor[0,0] = epsilon_mean - epsilon_del*(np.cos(inclination)**2 - np.sin(inclination)**2 * np.cos(2*azimuth))
epsilon_tensor[0,1] = epsilon_del * np.sin(inclination)**2 * np.sin(2*azimuth)
epsilon_tensor[0,2] = epsilon_del * np.sin(2*inclination) * np.cos(azimuth)
                     linelength=spacing * 1.8,
                     linewidth=1.3,
                     cmapAzimuth='hsv',
                     alpha=0.4)
plt.show()

ret_mask = np.abs(retardance_pr_PT[1]).copy()
ret_mask[ret_mask < 0.0075] = 0
ret_mask[ret_mask > 0.0075] = 1

plt.figure(figsize=(10, 10))
plt.imshow(ret_mask[:, :, z_layer], cmap='gray', origin='lower')
plt.figure(figsize=(10, 10))
plt.imshow(np.transpose(ret_mask[y_layer, :, :]),
           cmap='gray',
           origin='lower',
           aspect=z_step / ps)
plt.show()

# Angular histogram of 3D orientation
wo.orientation_3D_hist(azimuth[1].flatten(),
                       theta[1].flatten(),
                       ret_mask.flatten(),
                       bins=72,
                       num_col=1,
                       size=10,
                       contour_level=100,
                       hist_cmap='gray',
                       top_hemi=True)
plt.show()
                      aspect=z_step / ps)
plt.show()

# 3D orientation histograms for the structure and permittivity tensor
mask = np.abs(ret_crop[0].copy())
mask[mask < 0.0001] = 0
mask[mask >= 0.0001] = 1

az = azimuth_crop.flatten()
th = theta_crop.flatten()
val = np.abs(ret_crop * mask).flatten()
wo.orientation_3D_hist(az,
                       th,
                       val,
                       bins=36,
                       num_col=1,
                       size=5,
                       contour_level=100,
                       hist_cmap='gnuplot2',
                       top_hemi=True,
                       colorbar=False)

# ## Load fluorescence data
# Load fluorescence data
fluor_file_name = 'Cardiomyocytes_infected_2_fluor_decon.zarr'
reader_fluor = WaveorderReader(fluor_file_name, 'zarr')
I_fluor = np.squeeze(np.transpose(reader_fluor.get_array(0), (0, 1, 3, 4, 2)))

I_rgb = np.zeros((N_full, M_full, N_defocus, 3))
I_rgb[:, :, :, 1] = np.clip((I_fluor[1])**0.8 / np.max(
    (I_fluor[1])**0.8) * 255 * 1.5, 0, 255)
I_rgb[:, :, :, 2] = np.clip(I_fluor[0] / np.max(I_fluor[0]) * 255 * 2, 0, 255)
                     spacing=spacing,
                     window=spacing,
                     linelength=spacing * linelength_scale,
                     linewidth=1,
                     cmapAzimuth='hsv',
                     alpha=0.4,
                     subplot_ax=ax[1, 0])
ax[1, 0].set_title(f'XZ section (y = {y_layer})')

ax[1, 1].remove()
plt.show()

### Angular histogram of computed 3D orientation

ret_mask = np.abs(retardance_pr_PT[0]).copy()

## threshold parameters for 3D dataset ##
# ret_mask[ret_mask<0.15]=0
######################################

## threshold parameters for 2D dataset ##
ret_mask[ret_mask < 0.00125] = 0
######################################

plt.figure(figsize=(10, 10))
plt.imshow(ret_mask[:, :, z_layer], cmap='gray', origin='lower')
wo.orientation_3D_hist(azimuth[0].flatten(), \
                       theta[0].flatten(), \
                       ret_mask.flatten(),\
                       bins=36, num_col=1, size=10, contour_level = 100, hist_cmap='gnuplot2', top_hemi=True)
plt.show()