# -*- coding: utf-8 -*-
"""
Created on Mon Mar 23 16:04:32 2020

This script takes an HDF5 dataset (structured in a py4DSTEM compatible format) and extracts the data.
The "Center of Mass" (COM) of each diffraction pattern in the dataset is calculated
The COM is vizualized by overlaying a Quiver Plot on a COM magnitude image, BF image, or another image of the user's choice. 
This displays both the COM direction and magnitude on one Figure.  
Alternatively, COM may be displayed as a color plot, with brightness representing magnitude and hue representing direction.


@author: Barnaby Levin
"""

# -*- coding: utf-8 -*-
"""
Created on Tue Mar 17 11:35:44 2020

@author: Barnaby Levin
"""

import numpy as np
from numpy import pi
import h5py as h5
import matplotlib.pyplot as plt
import matplotlib as mpl
import pims as ps
from tifffile import imsave
from colorsys import hls_to_rgb

## First Read a py4DSTEM and LiberTEM compatible DE5 (HDF5) file into Python

Filename=r'C:\Users\Barnaby Levin\Documents\Microscope Data\4DSTEM\Oak Ridge Nion Data\Lietal-NPJCompMat-2019.h5'
#Filename=r'C:\Users\Barnaby Levin\Documents\Microscope Data\4DSTEM\Harvard Test Data\DE16_9.h5'


## Next, extract the 4D data from the file
Dataset=h5.File(Filename,'r')
base_items=list(Dataset.items())
print('Items in Base Directory: ', base_items )
G1=Dataset.get('4DSTEM_experiment')
G1_items=list(G1.items())
print('Items in 4DSTEM_experiment: ', G1_items )
G2=G1.get('data')
G2_items=list(G2.items())
print('Items in data: ', G2_items )
G3=G2.get('datacubes')
G3_items=list(G3.items())
print('Items in datacubes: ', G3_items )
G4=G3.get('datacube_0')
G4_items=list(G4.items())
print('Items in datacube_0: ', G4_items )
G5=G4.get('data')

Data=np.array(G5)
Datasize=np.shape(Data)

## Now we want to sum all of the patterns to make a PACBED Pattern.
## We will use the centre of mass of the PACBED pattern to define the origin 
## of the coordinates for the rest of our calculations. 
## We can also make a quick BF Image by summing all of the pixels in each diffraction pattern

SumIm=np.sum(Data,axis=(2,3)) # This is the image generated by summing all of the pixels in each pattern
PACBED=np.sum(Data,axis=(0,1)) # This is the diffraction pattern obtained by summing the patterns from every pixel of the STEM scan
# fig = plt.imshow(PACBED)
# plt.set_cmap('viridis')
# plt.show() 


## Now calculate the center of mass of the PACBED pattern. This will act as the Origin for Center of Mass Calculations in individual patterns
KX = np.arange(0,Datasize[3])
KY= np.arange(0,Datasize[2])
PANorm=np.sum(PACBED)

PACOMX=np.sum(np.multiply(KX,np.sum(PACBED,axis=1)))/PANorm 
PACOMY=np.sum(np.multiply(KY,np.sum(PACBED,axis=0)))/PANorm 

## These are a couple of plots just to check where the Center of Mass of the PACBED pattern is
#plt.scatter(PACOMX,PACOMY, color='w')
# ax=plt.quiver((Datasize[3]+1)/2,(Datasize[2]+1)/2,PACOMX-(1+Datasize[3])/2,-(PACOMY-(1+Datasize[2])/2), linewidth=5, scale=1, scale_units='x', width=0.02, headwidth=5, color='w')
# plt.show()

## Preallocate memory for upcoming calculations
COMData=np.zeros((Datasize[0]*Datasize[1],6))
COMXImage=np.zeros((Datasize[0],Datasize[1]))
COMYImage=np.zeros((Datasize[0],Datasize[1]))
COMMagImage=np.zeros((Datasize[0],Datasize[1]))
COMAngImage=np.zeros((Datasize[0],Datasize[1]))


## For loop to calculate COM for each diffraction pattern in the dataset
for nx in range(0,Datasize[1]):
    for ny in range(0,Datasize[0]):
        CBED=Data[ny,nx,:,:]
        Norm=np.sum(CBED)
       
        ## Center of mass of each CBED pattern
        RawCOMX=np.sum(np.multiply(KX,np.sum(CBED,axis=1)))/Norm 
        RawCOMY=np.sum(np.multiply(KY,np.sum(CBED,axis=0)))/Norm 
        
        ## Center of mass of each CBED pattern corrected so origin is center of mass of PACBED pattern
        COMX=RawCOMX-PACOMX
        COMY=RawCOMY-PACOMY
        
        ## Magnitude and direction of COM shift
        COMMag=np.sqrt(COMX*COMX+COMY*COMY)
        COMAng=np.arctan2(-COMY,COMX)
        
        ## Copy the data into a 2D array for use in the quiver arrow overlay plot
        COMData[(nx*Datasize[0])+ny,:]=[nx,ny,COMX,COMY,COMMag,COMAng]
        
        ## Copy the data into 2D arrays to make different types of COM images
        COMXImage[ny,nx]=COMX
        COMYImage[ny,nx]=COMY
        COMMagImage[ny,nx]=COMMag
        COMAngImage[ny,nx]=COMAng

## Comment and Uncomment the plotting code below to make different types of overlay plots. 

## Plot COM Magnitude image with perceptually uniform Viridis colormap
fig = plt.imshow(COMMagImage, cmap='viridis')
plt.show()

## Plot COM Angle (also called direction or phase) image with cyclic hsv colormap
# fig = plt.imshow(COMAngImage)
# plt.set_cmap('hsv')
# plt.show() 

## Plot COM-X image with divergent seismic colormap
# fig = plt.imshow(COMXImage)
# plt.set_cmap('seismic')
# plt.show() 

## Plot COM-Y image with divergent seismic colormap
# fig = plt.imshow(COMYImage)
# plt.set_cmap('seismic')
# plt.show() 

## Plot summed Bright Field Image with perceptually uniform Viridis colormap
# fig = plt.imshow(SumIm, cmap='viridis')
# plt.show()

## Plot ADF Image associated with the data (uploaded separately from the 4D Dataset) 
# fig = plt.imshow(expr_haadf, cmap='viridis')
# plt.show()

## Plot COM as Vector Field represented by arrows (quiver plot)
## There are a number of different parameters that control the style of the plot. 
## These are documented here: https://matplotlib.org/api/_as_gen/matplotlib.pyplot.quiver.html  
ax=plt.quiver(COMData[:,0],COMData[:,1],COMData[:,2],COMData[:,3], color='w', headlength=3,headwidth=2,minlength=3, scale=Datasize[3]/40, scale_units='x')#, scale=Datasize[3]/400, scale_units='x', width=0.004, headwidth=5, color='w')
plt.show()


## This visualization plots the COM vector field in one colormap, where brightness 
## is proportional to magnitude and hue is proportional to angle (also called direction or phase)
# def colorize(z):
#     r = np.abs(z)
#     arg = np.angle(z) 

#     h = (arg + pi)  / (2 * pi) + 0.5
#     l = 1.0 - 1.0/(1.0 + r**0.3)
#     s = 0.8

#     c = np.vectorize(hls_to_rgb) (h,l,s) # --> tuple
#     c = np.array(c)  # -->  array of (3,n,m) shape, but need (n,m,3)
#     c = c.swapaxes(0,2) 
#     return c

# z=COMXImage+1j*COMYImage
# ## w = 1/(z+1j)**2 + 1/(z-2)**2
# img = colorize(z)
# plt.imshow(img)
# plt.show()


