import sys
import importlib.metadata
def test_package(package_name):
    """Test if package exists and returns version or -1"""
    try:
        version = importlib.metadata.version(package_name)
    except importlib.metadata.PackageNotFoundError:
        version = '-1'
    return version

if test_package('pyTEMlib') < '0.2024.1.0':
    print('installing pyTEMlib')
    !{sys.executable} -m pip install git+https://github.com/pycroscopy/pyTEMlib.git@main -q --upgrade

if 'google.colab' in sys.modules:
    !{sys.executable} -m pip install numpy==1.24.4
print('done')

%matplotlib  widget
import matplotlib.pyplot as plt
import numpy as np
import sys
if 'google.colab' in sys.modules:
    from google.colab import output
    output.enable_custom_widget_manager()
    
# additional package 
import  itertools 
import scipy.constants as const

import ipywidgets as ipyw

# Import libraries from pyTEMlib
import pyTEMlib
import pyTEMlib.kinematic_scattering as ks         # Kinematic scattering Library
                             # Atomic form factors from Kirklands book

### And we use the image tool library of Quantifit
import pyTEMlib.file_tools as ft
import pyTEMlib

__notebook_version__ = '2024.02.21'

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

#####################
# Plot ZOLZ SAED Pattern #
#####################

## Get information as dictionary
tagsD = atoms.info['diffraction']

#We plot only the allowed diffraction spots
points = tagsD['allowed']['g']
# we sort them by order of Laue zone
ZOLZ = tagsD['allowed']['ZOLZ']
HOLZ = tagsD['allowed']['HOLZ']

# Plot
fig = plt.figure()
ax = fig.add_subplot(111)
# We plot the x,y axis only; the z -direction is set to zero - this is our projection
ax.scatter(points[ZOLZ,0], points[ZOLZ,1], c='red', s=20)
ax.scatter(points[HOLZ,0], points[HOLZ,1], c='green', s=20)

# zero spot plotting
ax.scatter(0,0, c='red', s=100)
ax.scatter(0,0, c='white', s=40)

ax.axis('equal')
FOV = 18
plt.ylim(-FOV,FOV); plt.xlim(-FOV,FOV); plt.show()

import pyTEMlib.animation  as animate

plt.figure()
animate.deficient_holz_line(exact_bragg=False, laue_zone=1)
#animate.deficient_holz_line(exact_bragg='True', laue_zone=1, color='blue')
animate.deficient_holz_line(exact_bragg='True', laue_zone=1, color='red', shift=True)

atoms.info['experimental']['nearest_zone_axes']
nearest = atoms.info['experimental']['nearest_zone_axes']
print('Next nearest zone axes are:')
for i in range(1, nearest['amount']):
    print(f"  {nearest[str(i)]['hkl']}: mistilt: {np.rad2deg(nearest[str(i)]['mistilt_alpha']):6.2f}, "
          f"{np.rad2deg(nearest[str(i)]['mistilt_beta']):6.2f}")

#Calculate angle between K0 and deficient cone vector
#For dynamic calculations K0 is replaced by Kg
K0 = np.linalg.norm(atoms.info['experimental']['incident_wave_vector'])
g_allowed = atoms.info['diffraction']['allowed']['g']
g_norm_allowed = np.linalg.norm(g_allowed, axis = 1)
dtheta = np.arcsin(g_norm_allowed/K0/2.)

np.arcsin(np.abs(g_allowed[:,2])/g_norm_allowed)

#calculate length of distance of deficient cone to K0 in ZOLZ plane

gd_length =2*np.sin((dtheta)/2 )*K0

#Calculate nearest point of HOLZ and Kikuchi lines

gd = g_allowed.copy()

gd[:,0] = -gd[:,0]*gd_length/g_norm_allowed

gd[:,1] = -gd[:,1]*gd_length/g_norm_allowed

gd[:,2] = 0.

###calculate and save line in Hough space coordinates (distance and theta)

slope = gd[:,0]/(gd[:,1]+1e-20)

distance = gd_length

theta = np.arctan(slope)