import sys
import importlib.metadata
def test_package(package_name):
    """Test whether 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.2025.12.0':
    print('installing pyTEMlib')
    !{sys.executable} -m pip install git+https://github.com/pycroscopy/pyTEMlib.git@main -q --upgrade

print('done')

%matplotlib  widget
import os
import sys
single

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

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

# Import microscopt library
sys.path.insert(0, "../../")  # point to the pyTEMlib location
%load_ext autoreload
%autoreload 2
import pyTEMlib
    
# it is a good idea to show the version numbers at this point for archiving reasons.
__notebook_version__ = '2025.12.06'
print('pyTEM version: ', pyTEMlib.__version__)
print('notebook version: ', __notebook_version__)

# ------Input -------------
load_your_own_data = False
# -------------------------
if 'google.colab' in sys.modules:
        drive.mount("/content/drive")

if load_your_own_data:
    fileWidget = pyTEMlib.file_tools.FileWidget()

if load_your_own_data:
    datasets = fileWidget.datasets
    main_dataset = fileWidget.selected_dataset
else:  # load example
    datasets = pyTEMlib.file_tools.open_file(os.path.join("../../example_data", "GOLD-NP-DIFF.dm3"))
    main_dataset =  datasets['Channel_000']

view = main_dataset.plot(vmax=20000)

center = np.array(main_dataset.shape)/2

plt.figure(figsize=(8, 6))
plt.imshow(np.log(3.+main_dataset).T, origin = 'upper')
# gca get current axis (plot)
selector = matplotlib.widgets.EllipseSelector(plt.gca(), None,interactive=True )
# selector.to_draw.set_visible(True)
radius = 559 
center = np.array(center)

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

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

center = (x_center, y_center)
print(f'center = {center[0]:.1f}, {center[1]:.1f} [pixels]')

tags = pyTEMlib.diffraction_tools.scattering_profiles(main_dataset, center)

# ###################
# Make some profiles
# ###################
profile_0 = tags['polar_projection'][:,0:10].sum(axis=1)
profile_0 += tags['polar_projection'][:,350:360].sum(axis=1)
profile_180 = tags['polar_projection'][:,190:210].sum(axis=1)
profile_270 = tags['polar_projection'][:,260:280].sum(axis=1)
profile_90 = tags['polar_projection'][:,80:100].sum(axis=1)

# ############### 
# Plot

scale = main_dataset.get_image_dims(return_axis=True)[0].slope
extent = (0, 360, tags['polar_projection'].shape[0]*scale, scale)
log_polar = np.log2(tags['polar_projection']+1e-12)
log_polar -= log_polar.min()
plt.figure()
plt.imshow(log_polar, extent=extent, cmap="gray")
ax = plt.gca()
ax.set_aspect("auto");
plt.xlabel('angle [degree]');
plt.ylabel('distance [1/nm]')

plt.plot(tags['radial_average']/tags['radial_average'].max()*200,
         np.linspace(1,len(tags['radial_average']),len(tags['radial_average']))*scale,
         c='r')
plt.plot(profile_0/profile_0.max()*200,
         np.linspace(1,len(profile_0),len(profile_0))*scale,
         c='orange');
plt.plot(profile_90/profile_90.max()*200,
         np.linspace(1,len(profile_90),len(profile_90))*scale,c='purple');
plt.plot(profile_270/profile_270.max()*200,
         np.linspace(1,len(profile_270),len(profile_270))*scale,c='b');
plt.plot([0,360],[3.8,3.8])
plt.plot([0,360],[6.3,6.3])
plt.colorbar();

import scipy as sp
import scipy.signal as signal

scale = main_dataset.u.slope*4.28/3.75901247*1.005



# find_Bragg peaks in profile
peaks, g= signal.find_peaks(tags['radial_average'],rel_height =0.7, width=7)  # np.std(second_deriv)*9)

print(peaks*scale)

tags['ring_radii_px'] = peaks


plt.figure()

plt.imshow(log_polar,extent=(0,360,tags['polar_projection'].shape[0]*scale,scale),cmap='gray')

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

plt.plot(tags['radial_average']/tags['radial_average'].max()*200,
         np.linspace(1,len(tags['radial_average']),len(tags['radial_average']))*scale,
         c='r')
for i in peaks:
    if i*scale > 3.5:
        plt.plot((0,360),(i*scale,i*scale), linestyle='--', c = 'steelblue')

# -----Input for ring pattern calculation ----
structure = 'gold'
hkl_max = 7
verbose = True
# --------------------------------------------

atoms = pyTEMlib.crystal_tools.structure_by_name(structure)
main_dataset.structures['Structure_000'] = atoms
main_dataset.metadata['experiment']['hkl_max']  = hkl_max

pyTEMlib.diffraction_tools.ring_pattern_calculation(main_dataset, verbose=verbose)

# -------Input of grain size ----
resolution  = .0001 # 1/nm
thickness = 100 # Ang
elements = ['Au', 'C']
# -------------------------------
import scipy

def plot_diffraction_profile(dataset, resolution, thickness, elements=[]):
    tags = dataset.metadata.get('diffraction', {})
    if not tags:
        print("No diffraction metadata found.")
        return  
    unique = tags['Ring_Pattern']['allowed']['g norm']
    family = tags['Ring_Pattern']['allowed']['hkl']
    intensity = np.array(tags['Ring_Pattern']['allowed']['structure_factor'])

    scale = main_dataset.get_image_dims(return_axis=True)[0].slope*1.085*1.0/10*1.01

    width = (1/thickness + resolution) / scale
    
    intensity2 = intensity/intensity.max()*10

    profile = tags['radial_average']
    gauss = scipy.signal.windows.gaussian(len(profile), std=width)
    simulated_profile = np.zeros(len(profile))
    reciprocal_dist = np.linspace(1,len(profile),len(profile))*scale

    x  =[]
    form_factors = np.zeros((len(elements), len(reciprocal_dist)))
    for j, element in enumerate(elements):
        for i, r_dist in enumerate(reciprocal_dist):
            form_factors[j,i] = pyTEMlib.diffraction_tools.form_factor(element, r_dist)
    yAu = form_factors[elements.index('Au')]
    plt.figure()
    plt.plot(reciprocal_dist,profile/profile.max()*150, color='blue', label='experiment');
    for j in range(len(unique)-1):
        if unique[j] < len(profile)*scale:
            # 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]*np.array(yAu)*1.3#rec_dist**2*10
            simulated_profile = simulated_profile + g
    plt.plot(np.linspace(1,len(profile),len(profile))*scale,simulated_profile, label='simulated');
    plt.plot(reciprocal_dist,np.array(yAu)**2, label='form_factor')
    plt.xlabel(r'angle (1/$\AA$)')
    plt.legend()
#plt.ylim(-35,210);

# pyTEMlib.diffraction_tools.
plot_diffraction_profile(main_dataset, resolution, thickness, elements)

# pyTEMlib.diffraction_tools.
import sidpy
def set_center(main_dataset, center, scale=scale):
    if isinstance(scale, float):
        scale = (scale, scale)
    x_axis = np.linspace(0, main_dataset.shape[0]-1, main_dataset.shape[0])-center[1]
    x_axis *= scale[0]
    y_axis = np.linspace(0, main_dataset.shape[1]-1, main_dataset.shape[1])-center[0]
    y_axis *= -scale[1]
    x = sidpy.Dimension(name='u', values=x_axis)
    
    main_dataset.set_dimension(0, sidpy.Dimension(name='u', values=x_axis, units='1/nm',
                                                  dimension_type='spatial', quantity='reciprocal distance'))
    main_dataset.set_dimension(1, sidpy.Dimension(name='v', values=y_axis, units='1/nm',
                                                  dimension_type='spatial', quantity='reciprocal distance'))

v =main_dataset.plot()

scale  = main_dataset.get_image_dims(return_axis=True)[0].slope

pyTEMlib.diffraction_tools.set_center(main_dataset, center, scale=scale*1.07)
v = pyTEMlib.diffraction_tools.plot_ring_pattern(main_dataset.metadata['diffraction'], main_dataset)

plt.close('all')

main_dataset.metadata['diffraction']['analysis'] = 'Indexing_Diffraction_Rings'
main_dataset.metadata['diffraction']['scale'] = scale
main_dataset.metadata['diffraction'].keys()

h5_group = pyTEMlib.file_tools.save_dataset(datasets, filename='Gold_SAED.hf5')
h5_group.file.close()