import sys
from pkg_resources import get_distribution, DistributionNotFound

def test_package(package_name):
    """Test if package exists and returns version or -1"""
    try:
        version = get_distribution(package_name).version
    except (DistributionNotFound, ImportError) as err:
        version = '-1'
    return version

if test_package('pyTEMlib') < '0.2025.1.0':
    print('installing pyTEMlib')
    !{sys.executable} -m pip install  --upgrade pyTEMlib
print('done')

%matplotlib  widget
import matplotlib.pyplot as plt
import matplotlib
import numpy as np

# additional package 
import itertools 
import scipy.constants as const
import os
import sys
import skimage # for polar coordinates of diffraction pattern
import scipy


if 'google.colab' in sys.modules:
    from google.colab import output
    output.enable_custom_widget_manager()
    from google.colab import drive

# Import libraries from the book
import pyTEMlib
# it is a good idea to show the version numbers at this point for archiving reasons.
__notebook_version__ = '2026.02.12'

print('pyTEM version: ', pyTEMlib.__version__)
print('notebook version: ', __notebook_version__)

if 'google.colab' in sys.modules:
    drive.mount("/content/drive")
fileWidget = pyTEMlib.file_tools.FileWidget(sum_frames=True)

diff_pattern = fileWidget.selected_dataset
print(f"alpha tilt {np.degrees(diff_pattern.metadata['experiment']['stage']['tilt']['alpha']):.2f}°")  
print(f"beta tilt {np.degrees(diff_pattern.metadata['experiment']['stage']['tilt']['beta']):.2f}°")  
view = np.log2(1+np.abs(diff_pattern)).plot()

diff_pattern

diff_pattern.view_metadata()
print( f"length of incident wavevector {1/pyTEMlib.utilities.get_wavelength(200000, unit='A'):.1f} 1/A")

print(f'Camera length: {float(diff_pattern.original_metadata['Optics']['CameraLength'])*1000:.2f} mm')
print(f'Pixel size: {diff_pattern.u.slope:.4f} 1/nm')

## Access the data of the loaded image

#diff_pattern = np.array(main_dataset.sum(axis=0))
diff_pattern = diff_pattern-diff_pattern.min()
radius = int( diff_pattern.shape[1]/6)
center = np.array([diff_pattern.shape[0]/2, diff_pattern.shape[1]/2])

center= np.unravel_index(np.argmax(np.array(diff_pattern), axis=None), diff_pattern.shape)


plt.figure(figsize=(8, 6))
plt.imshow(np.log(3.+diff_pattern).T, origin = 'upper')
current_axis = plt.gca()
selector = matplotlib.widgets.EllipseSelector(current_axis, 
                           None,
                           interactive=True , 
                           minspanx=5, minspany=5,
                           spancoords='pixels')  # gca get current axis (plot)

center = np.array(center)

selector.extents = (center[0]-radius-3,center[0]+radius-3,center[1]-radius-3,center[1]+radius-3)

plt.show()

diff_pattern.shape[0]/2, center, np.array(selector.extents)+radius

xmin, xmax, ymin, ymax = selector.extents
x_center, y_center = selector.center
x_shift = x_center - diff_pattern.shape[0]/2
y_shift = y_center - diff_pattern.shape[1]/2
print(f'radius = {(xmax-xmin)/2:.0f} pixels')

center = (x_center, y_center )
print(f'new center = {center} [pixels]')

out_tags ={}
out_tags['center'] = center

center = (x_center, y_center)
scale = diff_pattern.u.slope
polar_projection = skimage.transform.warp_polar(diff_pattern, center=center).T
polar_projection[polar_projection<0.] =0.
polar_projection += 1e-12
log_polar =np.log2(polar_projection)
log_polar -= log_polar.min()
profile = polar_projection.sum(axis=1)*100


plt.figure()
im = plt.imshow(np.log(1+polar_projection),extent=(0,360,polar_projection.shape[0]*scale,scale),cmap="gray")
plt.colorbar(im)
ax = plt.gca()
ax.set_aspect("auto");
plt.xlabel('angle [degree]');
plt.ylabel('distance [1/nm]')

plt.plot(profile/profile.max()*8000,np.linspace(1,len(profile),len(profile))*scale,c='r');
plt.xlim(0,360)

# --- Input ------
sensitivity = 1.0
# ----------------

# find_Bragg peaks in profile
peaks, g= scipy.signal.find_peaks(np.log(1+profile),rel_height=sensitivity, width=7)  # np.std(second_deriv)*9)

print('Peaks are at pixels:')
print(peaks)

out_tags['ring_radii_px'] = peaks


plt.figure()

plt.imshow(np.log2(1.+polar_projection),extent=(0,360,polar_projection.shape[0]*scale,scale),cmap='gray', vmin=np.max(np.log2(1+diff_pattern))*0.5)

ax = plt.gca()
ax.set_aspect("auto");
plt.xlabel('angle [degree]');
plt.ylabel('distance [1/nm]')

plt.plot(profile/profile.max()*5000,np.linspace(1,len(profile),len(profile))*scale,c='r');
plt.xlim(0,360)
for i in peaks:
    if i*scale > 3.5:
        plt.plot((0,360),(i*scale,i*scale), linestyle='--', c = 'steelblue')

# -------Input  -----
material = 'aluminium'
# -------------------

# Initialize the dictionary with all the input# Initialize the dictionary with all the input
atoms = pyTEMlib.crystal_tools.structure_by_name(material)

diff_pattern.structures['Structure_000'] = atoms


#Reciprocal Lattice 
# We use the linear algebra package of numpy to invert the unit_cell \"matrix\"
reciprocal_unit_cell = atoms.cell.reciprocal() # transposed of inverted unit_cell

#INPUT
hkl_max = 7#  maximum allowed Miller index

acceleration_voltage = 200.0 *1000.0 #V
wave_length  = pyTEMlib.utilities.get_wavelength(acceleration_voltage, unit='A')


h  = np.linspace(-hkl_max,hkl_max,2*hkl_max+1)   # all to be evaluated single Miller Index
hkl  = np.array(list(itertools.product(h,h,h) )) # all to be evaluated Miller indices
g_hkl = np.dot(hkl,reciprocal_unit_cell)  

# Calculate Structure Factors

structure_factors = []

base = atoms.positions # in Carthesian coordinates
for j  in range(len(g_hkl)):
    F = 0
    for b in range(len(base)):
        # Atomic form factor for element and momentum change (g vector)
        f = pyTEMlib.diffraction_tools.get_form_factor(atoms[b].symbol,np.linalg.norm(g_hkl[j]))  
        F += f * np.exp(-2*np.pi*1j*(g_hkl[j]*base[b]).sum())        
    structure_factors.append(F)
F = structure_factors = np.squeeze(np.array(structure_factors))

# Allowed reflections have a non zero structure factor F (with a  bit of numerical error)
allowed = np.absolute(structure_factors) > 0.001

distances = np.linalg.norm(g_hkl, axis = 1)

print(f' Of the evaluated {hkl.shape[0]} Miller indices {allowed.sum()} are allowed. ')
# We select now all the 
zero = distances == 0.
allowed = np.logical_and(allowed,np.logical_not(zero))

F = F[allowed]
g_hkl = g_hkl[allowed]
hkl = hkl[allowed]
distances = distances[allowed]

sorted_allowed = np.argsort(distances)

distances = distances[sorted_allowed]
hkl = hkl[sorted_allowed]
F = F[sorted_allowed]

# How many have unique distances and what is their muliplicity

unique, indices  = np.unique(distances, return_index=True)

print(f' Of the {allowed.sum()} allowed Bragg reflections there are {len(unique)} families of reflections.')

intensity = np.absolute(F[indices]**2*(np.roll(indices,-1)-indices))
print('\n index \t  hkl \t      1/d [1/Ang]       d [pm]     F      multip.  intensity' )
family = []

reflection = 0
out_tags['reflections'] = {}
multiplicitity = (np.roll(indices,-1)-indices)
intensity = np.absolute(F[indices]**2*multiplicitity)
print(f"\n index \t     {'hkl'.ljust(12)}\t 1/d [1/Ang]   d [pm] \t  F \t multip. intensity")
family = []
index = 0
for j in range(0, len(unique)-2):
    i = indices[j]    
    i2 = indices[j+1]   
    family.append(hkl[i+np.argmax(hkl[i:i2].sum(axis=1))])
    print(f'{i:3g}\t  {str(family[j]).ljust(13)} \t  {distances[i]:.4f}  \t {1/distances[i]*100:.0f} \t {np.absolute(F[i]):.2f}',
          f'\t  {indices[j+1]-indices[j]:3g} \t {intensity[j]:.2f}') 
    
    out_tags['reflections'][str(reflection)] = {'index': index,
                                                'recip_distances': distances[i],
                                                'structure_factor': np.absolute(F[i]),
                                                'multiplicity': indices[j+1]-indices[j],
                                                'intensity': intensity[j]}
    reflection +=1
    index+=1
diff_pattern.metadata['SAED'] = out_tags

# -------Input of grain size ----
first_peak_pixel = 100
first_peak_reciprocal_distance = 0.5
pixel_size = first_peak_reciprocal_distance/first_peak_pixel
resolution  = 0 # 1/nm
thickness = 100 # Ang
# -------------------------------

print(f'Pixel size is {pixel_size:.5f} 1/Ang')
print(f'Indicated pixel size is {diff_pattern.u.slope/10:.5f} 1/Ang')

width = (1/thickness + resolution) / scale
# scale = ft.get_slope(main_dataset.dim_0.values)  *1.085*1.0/10
scale = pixel_size
intensity2 = intensity/intensity.max()*10

gauss = scipy.signal.windows.gaussian(len(profile), std=width)
simulated_profile = np.zeros(len(profile))
rec_dist = np.linspace(1,len(profile),len(profile))*pixel_size


plt.figure()
plt.plot(rec_dist,profile/profile.max()*150, color='blue', label='experiment');
for j in range(len(unique)-1):
    if unique[j] < len(profile)*scale and j <20:
        # plot lines
        plt.plot([unique[j],unique[j]], [0, intensity2[j]],c='r')
        # plot indices
        index = '{'+f'{family[j][0]:.0f} {family[j][1]:.0f} {family[j][2]:.0f}'+'}' # pretty index string
        plt.text(unique[j],-3, index, horizontalalignment='center',
              verticalalignment='top', rotation = 'vertical', fontsize=8, color = 'red')
        
        # place Gaussian with appropriate width in profile
        g = np.roll(gauss,int(-len(profile)/2+unique[j]/scale))* intensity2[j]*10#rec_dist**2*10
        simulated_profile = simulated_profile + g
plt.plot(np.linspace(1,len(profile),len(profile))*scale,simulated_profile/50, label='simulated');
plt.xlabel('angle (1/$\AA$)')
plt.legend()
plt.ylim(-.5,10)

plot_profile = profile.copy()
plot_profile[:first_peak_pixel-20] = 0
fig = plt.figure(figsize=(9, 6)) 

extent= np.array([-center[0], diff_pattern.shape[0]-center[0],-diff_pattern.shape[1]+center[1], center[1]])*scale

plt.imshow(np.log(3.+diff_pattern).T,cmap='gray', extent=(extent*1.0)) #, vmin=np.max(np.log2(1+diff_pattern))*0.5)
plt.xlabel(r'reciprocal distance [nm$^{-1}$]')
ax = fig.gca()
#ax.add_artist(circle1);
plt.plot(np.linspace(1,len(profile),len(profile))*scale,plot_profile/plot_profile.max(), color='y');
plt.plot((0,len(profile)*scale),(0,0),c='r')

for j in range(len(unique)-1):
    i = indices[j]   
    if distances[i] < len(profile)*scale:
        plt.plot([distances[i],distances[i]], [0, intensity2[j]/20],c='r')
        arc = matplotlib.patches.Arc((0,0), distances[i]*2, distances[i]*2, angle=90.0, theta1=0.0, theta2=270.0, color='r', fill= False, alpha = 0.5)#, **kwargs)
        ax.add_artist(arc);
plt.scatter(0,0);

for i in range(6):
    index = '{'+f'{family[i][0]:.0f} {family[i][1]:.0f} {family[i][2]:.0f}'+'}' # pretty index string
    plt.text(unique[i],-0.05, index, horizontalalignment='center',
             verticalalignment='top', rotation = 'vertical', fontsize=8, color = 'white')

plt.xlim(diff_pattern.u[0]/10, diff_pattern.u[-1]/10)