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.2025.2.0':
    print('installing pyTEMlib')
    !{sys.executable} -m pip install  --upgrade pyTEMlib -q
print('done')

%matplotlib widget
import numpy as np
import matplotlib.pylab as plt 
import sys

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

import sys
sys.path.insert(0, '../../')
import sidpy

%load_ext autoreload
%autoreload 2
if 'google.colab' in sys.modules:
    drive.mount("/content/drive")

import pyTEMlib.file_tools      # File input/ output library
import pyTEMlib.image_tools 
import pyTEMlib.probe_tools
import pyTEMlib.atom_tools

    
# For archiving reasons it is a good idea to print the version numbers out at this point
print('pyTEM version: ', pyTEMlib.__version__)

## Do all of registration

notebook_tags= {'notebook': 'Image_Registration',
                'notebook_version': '2025_03_13',
                'pyTEM version': pyTEMlib.__version__}

fileWidget = pyTEMlib.file_tools.FileWidget()

dataset = fileWidget.selected_dataset
v = dataset.plot()

rigid_registered = pyTEMlib.image_tools.rigid_registration(dataset, normalization='phase')

v = rigid_registered.plot()

non_rigid= pyTEMlib.image_tools.demon_registration(rigid_registered)

view2 = non_rigid.plot()

demon_registered, rigid_registered = pyTEMlib.image_tools.complete_registration(dataset)
demon_registered.data_type = 'image_stack'
view2 = demon_registered.plot()

view3 = rigid_registered.plot()

rigid_registered.metadata['analysis']['rigid_registration']['drift']

drift = rigid_registered.metadata['analysis']['rigid_registration']['drift']
polynom_degree = 2 # 1 is linear fit, 2 is parabolic fit, ...

x = np.linspace(0,drift.shape[0]-1,drift.shape[0])

line_fit_x = np.polyfit(x, drift[:,0], polynom_degree)
poly_x = np.poly1d(line_fit_x)
line_fit_y = np.polyfit(x, drift[:,1], polynom_degree)
poly_y = np.poly1d(line_fit_y)

plt.figure()
plt.axhline(color = 'gray')
plt.plot(x, drift[:,0], label = 'drift x')
plt.plot(x, drift[:,1], label = 'drift y')
plt.plot(x, poly_x(x),  label = 'fit_drift_x')
plt.plot(x, poly_y(x),  label = 'fit_drift_y')

plt.legend();
ax_pixels = plt.gca()
ax_pixels.step(1, 1)

scaleX = (rigid_registered.x[1]-rigid_registered.x[0])*1000.  #in pm

ax_pm = ax_pixels.twinx()
x_1, x_2 = ax_pixels.get_ylim()

ax_pm.set_ylim(x_1*scaleX, x_2*scaleX)

ax_pixels.set_ylabel('drift [pixels]')
ax_pm.set_ylabel('drift [pm]')
ax_pixels.set_xlabel('image number');
plt.tight_layout()

# ------- Input ------
atoms_size = 0.08 # in nm
# --------------------

image = non_rigid.sum(axis=0)
image.data  = skimage.filters.difference_of_gaussians(image, 5, 20)
out_tags = {}
image2.metadata['experiment']= {'convergence_angle': 30, 'acceleration_voltage': 200000.}

scale_x = image.x.slope
gauss_diameter = atoms_size/scale_x
print(gauss_diameter)
if gauss_diameter < 1:
    print('smal')
    gauss_diameter = 1
print(gauss_diameter, gauss_diameter*scale_x)
gauss_probe = pyTEMlib.probe_tools.make_gauss(image.shape[0], image.shape[1], gauss_diameter)

print('Deconvolution of ', dataset.title)
LR_dataset = pyTEMlib.image_tools.decon_lr(image, gauss_probe, verbose=False)

extent = LR_dataset.get_extent([0,1])
fig, ax = plt.subplots(1, 2, figsize=(8, 4), sharex=True, sharey=True)
ax[0].imshow(image.T, extent = extent,vmax=np.median(np.array(image))+3*np.std(np.array(image)))
ax[1].imshow(LR_dataset.T, extent = extent, vmax=np.median(np.array(LR_dataset))+3*np.std(np.array(LR_dataset)));

v = LR_dataset.plot(cmap ='gray', vmax=500000)

new_image = np.array(image)
new_image = new_image - new_image.min()

fft_transform = (np.fft.fftshift(np.fft.fft2(np.array(new_image))))
dset  = image

image_dims = dset.get_image_dims(return_axis=True)

units_x = '1/' + image_dims[0].units
units_y = '1/' + image_dims[1].units

fft_dset = sidpy.Dataset.from_array(fft_transform)
fft_dset.quantity = dset.quantity
fft_dset.units = 'a.u.'
fft_dset.data_type = 'IMAGE'
fft_dset.source = dset.title
fft_dset.modality = 'fft'

fft_dset.set_dimension(0, sidpy.Dimension(np.fft.fftshift(np.fft.fftfreq(new_image.shape[0],
                                                                         d=dset.x[1]-dset.x[0])),
                                          name='u', units=units_x, dimension_type='RECIPROCAL',
                                          quantity='reciprocal_length'))
fft_dset.set_dimension(1, sidpy.Dimension(np.fft.fftshift(np.fft.fftfreq(new_image.shape[1],
                                                                         d=dset.y[1]- dset.y[0])),
                                          name='v', units=units_y, dimension_type='RECIPROCAL',
                                          quantity='reciprocal_length'))

filtered_power_spectrum = pyTEMlib.image_tools.power_spectrum(image)

filtered_power_spectrum.view_metadata()
print('source: ', filtered_power_spectrum.source)
view = filtered_power_spectrum.plot()

### Log Deconvolution

tags = {'analysis': {'name': 'Lucy_Richardson',
                     'input': dataset.title,
                     'probe_diameter': gauss_diameter,
                     'kind_of_probe': 'Gauss',
                     'probe_width': atoms_size
                     }}
LR_dataset.metadata['analysis'].update(notebook_tags)


datasets = fileWidget.datasets
datasets['LR_decon'] = LR_dataset
datasets['Sum_non_ridgid'] = image

import skimage
# ------- Input ------
threshold = 20.9 #usally between 0.01 and 0.9  the smaller the more atoms
atom_size = .1 #in nm  
# ----------------------


image = LR_dataset
#image = image_choice.dataset
# scale_x = pyTEMlib.file_tools.get_slope(image.dim_1)
blobs =  skimage.feature.blob_log(image, max_sigma=atom_size/scale_x, threshold=threshold)

fig1, ax = plt.subplots(1, 1,figsize=(8,7), sharex=True, sharey=True)
plt.title("blob detection ")

plt.imshow(image.T, interpolation='nearest',cmap='gray', vmax=np.median(np.array(image))+3*np.std(np.array(image)))
plt.scatter(blobs[:, 0], blobs[:, 1], c='r', s=20, alpha = .5);

out_tags  =  {}
out_tags['analysis']= 'Atom Positions'
out_tags['notebook']= __notebook__ 
out_tags['notebook_version']= __notebook_version__

out_tags['atoms'] = blobs
out_tags['atom_size'] = atom_size #in nm gives the size of the atoms or resolution
out_tags['threshold'] =  threshold  #between 0.01 and 0.1 
out_tags['pixel_size'] = scale_x


out_tags['name'] = 'Atom finding'
out_tags['title'] = out_tags['name']

atom_group = ft.log_results(current_channel, attributes=out_tags)

for key in current_channel:
    if 'Log' in key:
        if 'analysis' in current_channel[key]:
            print(f"{key} includes analysis: {current_channel[key]['analysis'][()]}")

import joblib 

def process_data(data):
    # Simulate a time-consuming data processing step
    import time
    time.sleep(2)
    return [data ** 2]*15

data = [1, 2, 3, 4, 5, 6]

# Parallel(n_jobs=workjob_num)(delayed(func_be_applied)(aug) for elem in array
results = joblib.Parallel(n_jobs=6)(joblib.delayed(process_data)(data[d]) for d in range(6))
print(results)

from joblib import Parallel, delayed 

class A(object):
    def __init__(self, x):
        self.x = x
    def p(self):
        self.y = self.x**2
        return self.y

if  __name__ == '__main__':
    runs = [A(x) for x in range(20)]
    with Parallel(n_jobs=2, verbose=5) as parallel:
        delayed_funcs = [delayed(lambda x:x.p())(run) for run in runs]
        run_A = parallel(delayed_funcs)

a = np.arange(12).reshape(2,2,3)

it = np.nditer(a, flags=['multi_index'])
it.remove_axis(2)
for x in it:
    print(x, a[it.multi_index], it.multi_index)

a[1,0]

a[np.indices([2,2])]

blobs

import ase
structure = ase.Atoms('Cu'*len(blobs), positions=blobs*[1,1,0], cell=[image.x[-1],image.y[-1],1], pbc=True)

st

image_choice = ft.ChooseDataset(current_channel)  

# ------- Input ------
atom_radius = 3  # in pixel
# ----------------------

if image_choice.dataset.data_type.name == 'IMAGE_STACK':
    image = image_choice.dataset.sum(axis=0)
else:
    image = image_choice.dataset

#atoms = atom_group['atoms'][()]
atoms = blobs
image = image-image.min()
print(image)

#atom_radius = 2
MaxInt = 0
MinInt = 0 
maxDist = 2
sym = pyTEMlib.atom_tools.atom_refine(np.array(image), atoms, atom_radius, max_int = 0, min_int = 0, max_dist = 2)
refined_atoms = np.array(sym['atoms'])

fig, ax = plt.subplots(1, 2, figsize=(8, 4), sharex=True, sharey=True)
ax[0].imshow(image.T)
ax[0].scatter(refined_atoms[:,0],refined_atoms[:,1],  s=10, alpha = 0.3, color = 'red')
ax[0].set_title('refined atom postion')
ax[1].imshow(image.T)
ax[1].scatter(atoms[:, 0], atoms[:, 1], c='r', s=10, alpha = .3);
ax[1].set_title('blobs on image');

fileWidget.datasets

h5_file = dataset.h5_dataset.file

print(h5_file.filename)
h5_file.close()

import simpleITK as sitk

def DemonReg(cube, verbose = False):
    """
    Diffeomorphic Demon Non-Rigid Registration 
    Usage:
    ------
    DemReg = DemonReg(cube, verbose = False)

    Input:
        cube: stack of image after rigid registration and cropping
    Output:
        DemReg: stack of images with non-rigid registration

    Dempends on:
        simpleITK and numpy
    
    Please Cite: http://www.simpleitk.org/SimpleITK/project/parti.html
    and T. Vercauteren, X. Pennec, A. Perchant and N. Ayache
    Diffeomorphic Demons Using ITK\'s Finite Difference Solver Hierarchy
    The Insight Journal, http://hdl.handle.net/1926/510 2007
    """
    
    DemReg =  np.zeros_like(cube)
    nimages = cube.shape[0]
    print(nimages)
    # create fixed image by summing over rigid registration

    fixed_np = np.average(current_dataset, axis=0)

    fixed = sitk.GetImageFromArray(fixed_np)
    fixed = sitk.DiscreteGaussian(fixed, 2.0)

    #demons = sitk.SymmetricForcesDemonsRegistrationFilter()
    demons = sitk.DiffeomorphicDemonsRegistrationFilter()

    demons.SetNumberOfIterations(200)
    demons.SetStandardDeviations(1.0)

    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(fixed);
    resampler.SetInterpolator(sitk.sitkBspline)
    resampler.SetDefaultPixelValue(0)

    done = 0
        
    for i in range(nimages):
        if done < int((i+1)/nimages*50):
            done = int((i+1)/nimages*50)
            sys.stdout.write('\r')
            # progress output :
            sys.stdout.write("[%-50s] %d%%" % ('*'*done, 2*done))
            sys.stdout.flush()
        
        moving = sitk.GetImageFromArray(cube[i])
        movingf = sitk.DiscreteGaussian(moving, 2.0)
        displacementField = demons.Execute(fixed,movingf)
        outTx = sitk.DisplacementFieldTransform( displacementField )
        resampler.SetTransform(outTx)
        out = resampler.Execute(moving)
        DemReg[i,:,:] = sitk.GetArrayFromImage(out)
        #print('image ', i)
        
    
    print(':-)')
    print('You have succesfully completed Diffeomorphic Demons Registration')
    
    return DemReg