#!/usr/bin/env python
# coding: utf-8

# This notebook is part of the `kikuchipy` documentation https://kikuchipy.org.
# Links to the documentation won't work from the notebook.

# # Geometrical EBSD simulations
# 
# In this tutorial, we will inspect and visualize the results from EBSD indexing by plotting Kikuchi lines and zone axes onto an EBSD signal.
# We consider this a *geometrical* EBSD simulation, since it is only positions of Kikuchi lines and zone axes that are computed.
# These simulations are based on the work by Aimo Winkelmann in the supplementary material to <cite data-cite="britton2016tutorial">Britton et al. (2016)</cite>.
# 
# These simulations can be helpful when checking whether indexing results are correct and for interpreting them. 
# 
# Let's import the necessary libraries

# In[ ]:


# Exchange inline for notebook or qt5 (from pyqt) for interactive plotting
get_ipython().run_line_magic('matplotlib', 'inline')

import matplotlib.pyplot as plt
import numpy as np

from diffpy.structure import Atom, Lattice, Structure
from diffsims.crystallography import ReciprocalLatticeVector
import hyperspy.api as hs
import kikuchipy as kp
from orix.crystal_map import Phase
from orix.quaternion import Rotation


# Plotting parameters
plt.rcParams.update(
    {"figure.figsize": (10, 10), "font.size": 20, "lines.markersize": 10}
)


# We'll inspect the indexing results of a small nickel EBSD dataset of (3, 3) patterns

# In[ ]:


# Use kp.load("data.h5") to load your own data
s = kp.data.nickel_ebsd_small()
s


# Let's enhance the Kikuchi bands by removing the static and dynamic backgrounds

# In[ ]:


s.remove_static_background()
s.remove_dynamic_background()


# In[ ]:


_ = hs.plot.plot_images(
    s, axes_decor=None, label=None, colorbar=False, tight_layout=True
)


# To project Kikuchi lines and zone axes onto our detector, we need
# 
# 1. a description of the crystal phase
# 
# 2. the set of Kikuchi bands to consider, e.g. the sets of planes {111}, {200}, {220}, and {311}
# 
# 3. the crystal orientations with respect to the reference frame
# 
# 4. the position of the detector with respect to the sample, in the form of a sample-detector model which includes the sample and detector tilt and the projection center (shortes distance from the source point on the sample to the detector), given here as (PC$_x$, PC$_y$, PC$_z$)

# ## Prepare phase and reflector list
# 
# We'll store the crystal phase information in an [orix.crystal_map.Phase](https://orix.readthedocs.io/en/stable/reference/generated/orix.crystal_map.Phase.html) instance

# In[ ]:


phase = Phase(
    space_group=225,
    structure=Structure(
        atoms=[Atom("Ni", [0, 0, 0])],
        lattice=Lattice(3.52, 3.52, 3.52, 90, 90, 90),
    ),
)

print(phase)
print(phase.structure)


# We'll build up the reflector list using [diffsims.crystallography.ReciprocalLatticeVector](https://diffsims.readthedocs.io/en/latest/reference/generated/diffsims.crystallography.ReciprocalLatticeVector.html)

# In[ ]:


rlv = ReciprocalLatticeVector(
    phase=phase, hkl=[[1, 1, 1], [2, 0, 0], [2, 2, 0], [3, 1, 1]]
)
rlv


# We'll obtain the symmetrically equivalent vectors and plot each family of vectors in a distinct colour in the stereographic projection

# In[ ]:


rlv = rlv.symmetrise().unique()
rlv.size


# In[ ]:


rlv.print_table()


# In[ ]:


# Dictionary with {hkl} as key and indices into ReciprocalLatticeVector as values
hkl_sets = rlv.get_hkl_sets()
hkl_sets


# In[ ]:


hkl_colors = np.zeros((rlv.size, 3))
for idx, color in zip(
    hkl_sets.values(),
    [
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, 1],
        [0.75, 0, 0.75],
    ],  # Red, green, blue, magenta
):
    hkl_colors[idx] = color


# In[ ]:


hkl_labels = []
for hkl in rlv.hkl.round(0).astype(int):
    hkl_labels.append(str(hkl).replace("[", "(").replace("]", ")"))


# In[ ]:


rlv.scatter(c=hkl_colors, grid=True, ec="k", vector_labels=hkl_labels)


# We can also plot the plane traces, i.e. the Kikuchi lines, in both hemispheres (they are identical for Ni)

# In[ ]:


rlv.draw_circle(
    color=hkl_colors, hemisphere="both", figure_kwargs=dict(figsize=(15, 10))
)


# ## Specify rotations and detector-sample geometry
# 
# Rotations and the detector-sample geometry for the nine nickel EBSD patterns are stored in the kikuchipy h5ebsd file.
# These were found by Hough indexing using *PyEBSDIndex* followed by orientation and PC refinement using a nickel EBSD master pattern simulated with *EMsoft*.
# See the [Hough indexing](hough_indexing.ipynb) and [Pattern matching](pattern_matching.ipynb) tutorials for details on indexing and the [reference frame](reference_frames.ipynb) tutorial for details on the definition of the detector-sample geometry.

# In[ ]:


rot = s.xmap.rotations
rot = rot.reshape(*s.xmap.shape)
rot  # Quaternions


# We describe the sample-detector model in an [kikuchipy.detectors.EBSDDetector](../reference/generated/kikuchipy.detectors.EBSDDetector.rst) instance.
# The sample was tilted $70^{\circ}$ about the microscope X direction towards the detector, and the detector normal was orthogonal to the optical axis (beam direction).
# Using this information, the projection center (PC) was found using *PyEBSDIndex* as described in the Hough indexing tutorial.

# In[ ]:


s.detector


# In[ ]:


s.detector.pc


# ## Create simulations on detector
# 
# Now we're ready to create geometrical simulations.
# We create simulations using the [kikuchipy.simulations.KikuchiPatternSimulator](../reference/generated/kikuchipy.simulations.KikuchiPatternSimulator.rst), which takes the reflectors as input

# In[ ]:


simulator = kp.simulations.KikuchiPatternSimulator(rlv)


# In[ ]:


sim = simulator.on_detector(s.detector, rot)


# By passing the detector and crystal orientations to
# [KikuchiPatternSimulator.on_detector()](../reference/generated/kikuchipy.simulations.KikuchiPatternSimulator.on_detector.rst),
# we've obtained a
# [kikuchipy.simulations.GeometricalKikuchiPatternSimulation](../reference/generated/kikuchipy.simulations.GeometricalKikuchiPatternSimulation.rst),
# which stores the detector and gnomonic coordinates of the Kikuchi lines and
# zone axes for each crystal orientation

# In[ ]:


sim


# We see that not all 50 of the reflectors in the reflector list are present in some pattern.

# ## Plot simulations
# 
# These geometrical simulations can be plotted one-by-one by themselves

# In[ ]:


sim.plot()


# Or, they can be plotted on top of patterns in three ways: passing a pattern to [GeometricalKikuchiPatternSimulation.plot()](../reference/generated/kikuchipy.simulations.GeometricalKikuchiPatternSimulation.plot.rst)

# In[ ]:


sim.plot(index=(1, 2), pattern=s.inav[2, 1].data)


# Or, we can obtain collections of lines, zone axes and zone axes labels as *Matplotlib* objects via [GeometricalKikuchiPatternSimulation.as_collections()](../reference/generated/kikuchipy.simulations.GeometricalKikuchiPatternSimulation.as_collections.rst) and add them to an existing Matplotlib axis

# In[ ]:


fig, ax = plt.subplots(ncols=3, nrows=3, figsize=(15, 15))

for idx in np.ndindex(s.axes_manager.navigation_shape[::-1]):
    ax[idx].imshow(s.data[idx], cmap="gray")
    ax[idx].axis("off")

    lines, zone_axes, zone_axes_labels = sim.as_collections(
        idx,
        zone_axes=True,
        zone_axes_labels=True,
        zone_axes_labels_kwargs=dict(fontsize=12),
    )
    ax[idx].add_collection(lines)
    ax[idx].add_collection(zone_axes)
    for label in zone_axes_labels:
        ax[idx].add_artist(label)

fig.tight_layout()


# Or, we can obtain the lines, zone axes, zone axes labels and PCs as *HyperSpy* markers via [GeometricalKikuchiPatternSimulation.as_markers()](../reference/generated/kikuchipy.simulations.GeometricalKikuchiPatternSimulation.as_markers.rst) and add them to a signal of the same navigation shape as the simulation instance.
# This enables navigating the patterns *with* the geometrical simulations

# In[ ]:


markers = sim.as_markers()

# To delete previously added permanent markers, do
# del s.metadata.Markers

s.add_marker(markers, plot_marker=False, permanent=True)


# In[ ]:


s.plot()