# Starting up a distributed Cluster locally 
# You don't have to do this but it helps to visualize what is happening
from dask.distributed import Client
client = Client()  # set up local cluster on your laptop
client
# Ignore INFO below just giving informataion about the scheduler set up

# setting the plotting backend
%matplotlib widget
dp = hs.load("data/mgo_nanoparticles.zspy", lazy=True) # Use Zarr backend for speed and load data lazily 
dp.plot() # plot the signal DP

mean_cbed = dp.mean() # get the mean diffraction pattern
mean_cbed.plot(vmax=500) # plot the mean diffraction pattern
roi = hs.roi.CircleROI(cx=0.3246, cy=-1.49316, r= 0.2) # create a ROI to be used as a virtual appature
roi.add_widget(mean_cbed) # add roi to plot

virtual = dp.get_integrated_intensity(roi) # integrate the ROI
virtual.plot()

mean_cbed = dp.mean() # get the mean diffraction pattern
mean_cbed.plot(vmax=1000) # plot the mean diffraction pattern
df_roi = hs.roi.CircleROI(cx=0, cy=0, r= 2.0, r_inner=0.5) # create a ROI to be used as a virtual appature
df_roi.add_widget(mean_cbed) # add roi to plot

df =  dp.get_integrated_intensity(df_roi) # integrate the DF ROI 
hs.plot.plot_images([df, virtual], label=["Virtual Dark Field", "Virtual Aperture"], tight_layout=True, colorbar=None) # Plot both ROI's using hyperspy

import numpy as np
dp_sliced = dp_sub.inav[20:30,20:30] # slice that data to make the interactive processing work better
slic_max = np.max(dp_sliced.data) # find max of the data for scaling
dp_sliced = dp_sliced/slic_max * 20 # scale the data for interactive find_peaks
pks = dp_sliced.find_peaks(method="local_max",
                           threshold_abs=10, 
                           min_distance=3,vmax=10) # find peaks interactively 

dp

# Convert to diffraction vectors (This workflow will be simplified once pyxem 1.0.0 is released!)
from pyxem.signals import DiffractionVectors
import numpy as np
pks = dp_sub.find_peaks(interactive=False, 
                    method="local_max",
                    threshold_abs=2.2*slic_max/20,
                    min_distance=3) # Finding peaks on the entire dataset
dv = DiffractionVectors.from_peaks(peaks=pks,
                                   center=np.array(dp.axes_manager.signal_shape)/2-.5,
                                   calibration=dp.axes_manager.signal_axes[0].scale) # convert to Diffraction Vectors Signal

%matplotlib widget
from pyxem.signals.diffraction_vectors import generate_marker_inputs_from_peaks
xx,yy = generate_marker_inputs_from_peaks(dv) # generating hyperspy markers from peaks
dp_sub.plot(vmax=500)
for x,y in zip(xx,yy):
    m = hs.plot.markers.point(y,x,
                              color="blue",
                              size=50,)
    dp_sub.add_marker(m, plot_marker=True, permanent=False,)

# plot the map (number found) of diffraction vectors 
dv.get_diffracting_pixels_map().plot()

distance_threshold = 0.1
min_samples = 7
# cluster the vectors to get only the unique vectors
unique_peaks = dv.get_unique_vectors(method='DBSCAN',
                                     distance_threshold=distance_threshold,
                                     min_samples=min_samples)
print(np.shape(unique_peaks.data)[0], ' unique vectors were found.')

#remove the zero beam
unique_peaks = unique_peaks.filter_magnitude(min_magnitude=.4,
                                   max_magnitude=np.inf) 
print(np.shape(unique_peaks)[0], ' unique vectors.')

%matplotlib widget
dpt = dp.T
dpt.plot()
for x,y in zip(unique_peaks.data[:,1],unique_peaks.data[:,0]):
    markers = hs.plot.markers.point(x= x,y=y, color="r")
    dpt.add_marker(markers, plot_on_signal=False)

from pyxem.generators import VirtualDarkFieldGenerator
radius=0.1
vdfgen = VirtualDarkFieldGenerator(dp, unique_peaks) # generate the virtual darkfield images
VDFs = vdfgen.get_virtual_dark_field_images(radius=radius) # create Dark Field images using the virtual appatures
VDFs.compute() # compute it and load into RAM
VDFs.plot() # plot the VDFS

# test for the right watershed parameters
from pyxem.utils.segment_utils import separate_watershed
i = 50
sep_i = separate_watershed(
    VDFs.inav[i].data, min_distance=min_distance, min_size=min_size,
    max_size=max_size, max_number_of_grains=max_number_of_grains,
    exclude_border=exclude_border, marker_radius=marker_radius,
    threshold=threshold, plot_on=True)

segs = VDFs.get_vdf_segments(min_distance=min_distance,
                             min_size=min_size,
                             max_size = max_size,
                             max_number_of_grains = max_number_of_grains,
                             exclude_border=exclude_border,
                             marker_radius=marker_radius,
                             threshold=threshold)

corrsegs = segs.correlate_vdf_segments(
    corr_threshold=corr_threshold, vector_threshold=vector_threshold,
    segment_threshold=segment_threshold)
print(np.shape(corrsegs.segments)[0],' correlated segments were found.')
num_segments = np.shape(corrsegs.segments)[0]

import matplotlib.pyplot as plt
f = plt.figure(figsize=(15,6))
virt_diff = corrsegs.get_virtual_electron_diffraction(calibration=dp.axes_manager.signal_axes[0].scale, shape=(144,144), sigma=.05)
virt_diff.set_diffraction_calibration(dp.axes_manager.signal_axes[0].scale)
segments = corrsegs.segments
labels = ["Segment: " + str(i) for i in range(num_segments)] + ["Segment: " + str(i) +"DP" for i in range(num_segments)] # For adding labels
hs.plot.plot_images(segments.split()+virt_diff.split(),
                    per_row=num_segments,
                    fig=f,tight_layout=True,
                    colorbar=None,
                    scalebar="all",
                    axes_decor=None,
                    label =labels) # adjust this function to get nice quality plotting in hyperspy
plt.show()

# Create a grid of roatations
from diffsims.generators.rotation_list_generators import get_beam_directions_grid
resolution = 1.5 # maximum angle in degrees between nearest templates. Pretty rough grid for speed.
grid_cub = get_beam_directions_grid("hexagonal", resolution, mesh="spherified_cube_edge")
print("Number of patterns: ", grid_cub.shape[0])

# Parameters necessary for simulating a template library
# half size of the images
diffraction_calibration=dp.axes_manager.signal_axes[0].scale
half_shape = (dp.data.shape[-2]//2, dp.data.shape[-1]//2)
# maximum radius in reciprocal space to calculate spot intensities for
reciprocal_radius = np.sqrt(half_shape[0]**2 + half_shape[1]**2)*diffraction_calibration
print("Max Recip. Radius: ", reciprocal_radius, "Inverse Angstroms")


# simulate the diffraction patterns for MgO
structure_matrix = diffpy.structure.loadStructure("data/MgO.cif")

# "The microscope = the parameters that determine how the templates are calculated"
diff_gen = DiffractionGenerator(accelerating_voltage=300,
                                precession_angle=0,
                                scattering_params=None,
                                shape_factor_model="linear",
                                minimum_intensity=0.1,
                                )

lib_gen = DiffractionLibraryGenerator(diff_gen)

# Generating a library
# "Library of structures and orientations"
library_phases_mgo = StructureLibrary(["MgO"], [structure_matrix], [grid_cub])
# Calculate the actual library

diff_lib_mgo = lib_gen.get_diffraction_library(library_phases_mgo,
                                              calibration=diffraction_calibration,
                                              reciprocal_radius=2,
                                              half_shape=half_shape,
                                              with_direct_beam=False,
                                              max_excitation_error=0.07)

frac_keep = 0.8 
n_keep = None

# if frac_keep < 1 or 1 < n_keep < number of templates then indexation
# templates in "indexes" that have the highest "fast" correlation
n_best = 2 # keep the 5 best rotations
result, phasedict = iutls.index_dataset_with_template_rotation(dp_sub,
                                                    diff_lib_mgo,
                                                    phases = ["MgO"],  # if we have multiple phases we can also specify which ones we want to consider. If it's not specified, all phases are used.
                                                    n_best = n_best,
                                                    frac_keep = frac_keep,
                                                    n_keep = n_keep,
                                                    delta_r = delta_r,
                                                    delta_theta = delta_theta,
                                                    max_r = 50,
                                                    intensity_transform_function=intensity_transform_function,
                                                    normalize_images = normalize_image,
                                                    normalize_templates=normalize_templates,
                                                    )

# plotting the orientation maps and the Inverse Pole Figures
solution = result["orientation"]

import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
from orix.projections import StereographicProjection

# map a vector onto the fundamental zone of the hexagon
def to_fundamental(data_sol):
    data_sol = np.abs(data_sol)
    data_sol = np.sort(data_sol, axis=-1)
    column = data_sol[...,0].copy()
    data_sol[..., 0] = data_sol[...,1]
    data_sol[..., 1] = column
    return data_sol


def get_ipf_color(vectors):
    # the following column vectors should map onto R [100], G [010], B[001], i.e. the identity. So the inverse of 
    # this matrix maps the beam directions onto the right color vector
    color_corners = np.array([[0, 1, 1],
                              [0, 0, 1],
                              [1, 1, 1]])
    color_mapper = np.linalg.inv(color_corners)
    # a bit of wrangling
    data_sol = to_fundamental(vectors.data)
    flattened = data_sol.reshape(np.product(data_sol.shape[:-1]), 3).T
    rgb_mapped = np.dot(color_mapper, flattened)
    rgb_mapped = np.abs(rgb_mapped / rgb_mapped.max(axis=0)).T
    rgb_mapped = rgb_mapped.reshape(data_sol.shape)
    return rgb_mapped
    
    
from orix.quaternion.rotation import Rotation
from orix.vector.vector3d import Vector3d

# draw IPF - Z (row 1), IPF - Y (row 2), IPF - Z (row 3)
fig, ax = plt.subplots(ncols = solution.shape[2], nrows = 3, figsize = (10, 6))

for i in range(solution.shape[2]):
    solution_vectors_z = Rotation.from_euler(np.deg2rad(solution[:,:,i,:]))*Vector3d.zvector()  
    solution_vectors_y = Rotation.from_euler(np.deg2rad(solution[:,:,i,:]))*Vector3d.yvector()   
    solution_vectors_x = Rotation.from_euler(np.deg2rad(solution[:,:,i,:]))*Vector3d.xvector()   
    ax[0, i].set_title(f"Solution {i}")
    ax[0, i].imshow(get_ipf_color(solution_vectors_z)*mask.data[:,:, np.newaxis])
    ax[1, i].imshow(get_ipf_color(solution_vectors_y)*mask.data[:,:, np.newaxis])
    ax[2, i].imshow(get_ipf_color(solution_vectors_x)*mask.data[:,:, np.newaxis])

ax[0,0].set_ylabel("IPF-Z")
ax[1,0].set_ylabel("IPF-Y")
ax[2,0].set_ylabel("IPF-X")

fig.tight_layout()

from pyxem.signals.diffraction_vectors import generate_marker_inputs_from_peaks
xx,yy = generate_marker_inputs_from_peaks(v) # generating hyperspy markers from peaks
dp_sub.plot(vmax=500)
for x,y in zip(xx,yy):
    m = hs.plot.markers.point(y,x,
                              color="blue",
                              size=50,)
    dp_sub.add_marker(m, plot_marker=True, permanent=False,)