# polymerCrunch.py Program for performing Fourier analysis of conductive polymer images:

# Paper: "Real space imaging of semi-crystalline order and single strands of polythiophene using ultra-high resolution atomic force microscopy"
#
# Author: Alex Summerfield - University of Nottingham
# Contact: alex.summerfield@nottingham.ac.uk / alexsummerfield@me.com / www.researchgate.net/profile/Alex_Summerfield
# last revision:  06/02/2019
# Notes: Commented version for supplementary information


# Program description:

# Program is run to perform sliding 2DFFT on sub-kernels of an image. Each kernel
# is rotationally interpolated around a torus of allowed spatial frequency
# values (selected by the user by changing appropriate program variables). The
# position of peaks corresponding to strong periodic features are then
# identified by the program and extracted along with the angle and spatial
# freuqncy that they occur at for each kernel position. The data is then
# exported to a pickle object to be re-loaded in the analysis and plotting
# script 'polymerAnalysis.py'. Typically takes a couple of hours to run.

################################################################################

# importing modules

import numpy as np # importing NumPy
import math # math functions
import sys
import cPickle as pickle # cPickle module for export of processed data

from scipy import interpolate # interpolation function from SciPy: https://www.scipy.org/
import detect_peaks # detect_peaks function (Author: Marcus Duarte, Link: https://github.com/demotu/BMC)
import loadBar # function for printing out progress of the script - from https://stackoverflow.com/questions/3173320/text-progress-bar-in-the-console

# FUNCTIONS ####################################################################

# normaliseProfile: Function for normalising a single line profile ###############################

# Input: inputProfile = 1D numpy array of a line profile

# Output: normalisedProfile = inputProfile normalised between 0 and 1.

def normaliseProfile(inputProfile):

    normalisedProfile = (inputProfile - np.min(inputProfile)) / (np.max(inputProfile) - np.min(inputProfile))

    return normalisedProfile

# End of functions ##############################################################


# Main program

# loading in data file

 # Defining file name to load nb: desired image is saved as 2D grid as .txt file using Gwyddion
fname = 'polymerPhaseImage.txt'

rawData = np.flipud(np.genfromtxt(fname)) # loading and flipping the image due to Gwyddion export

imSize = 100.0e-9 # Image size in nm
imShape = np.shape(rawData) # shape of input array
imPix = imSize/imShape[0] # spatial size of single pixel (metres)

deltaTheta = 1.0 # angle step size for interpolation of 2DFFT
sizeRange = np.array([4.0e-10, 1.5e-9]) # spatial size of the periodic features to filter
freqRange = 1./sizeRange # converting to spatial frequency
thetaVals = np.arange(0.0,360,deltaTheta) # angle values


# Defining kernel to perform FFT on ############################################

kNDesired = 8e-9 # desired kernel size (desired in m) - determined manually by trying different values.
kN = int(np.ceil(kNDesired / imPix)) # rounded size of kernel in pix

if kN % 2 != 0: # if kN is not divisible by two then add one to make 2DFFT calculation quicker
    kN = kN + 1

kSize = kN * imPix # actual kernel size (m)

slideVal = 5 # number of pixels the kernel is moved between calculating each 2DFFT

# Defining kernel sizes and pre-alocating memory to speed up script.

kernel = np.empty((kN,kN),dtype=np.float) # raw image Data kernel
kernelFFT = np.empty((kN,kN),dtype=np.float) # 2DFFT of kernel

imNyq = np.shape(kernel)[0]/(2*kSize) # spatial Nyquist frequency
freqVals = np.linspace(-imNyq,imNyq,np.shape(kernel)[0],dtype=np.float) # generating the spatial frequency value array

xGrid,yGrid = np.meshgrid(freqVals,freqVals) # Creating grid of spatial frequency values
freqGrid = np.sqrt(xGrid**2 + yGrid**2) # funding the 'distance' to centre of freqency grid
wavelengthGrid = 1./freqGrid # converting above to spatial wavelength
angleGrid = np.arctan2(yGrid,xGrid) + math.pi # calculating the angle at each positon

angleGrid1D = angleGrid.flatten() # flattened angleGrid for masking opereations later
freqGrid1D = freqGrid.flatten() # as above for spatial fequency

conFilter = np.logical_and(freqGrid <= max(abs(freqRange)),freqGrid >= min(abs(freqRange))) # creating filter
conFilter1D = conFilter.flatten() # flattening logical convolution filter

# creating angles of each positon

nX = imShape[0] - kernel.shape[0] # maximum x grid value of the iteration over the grid
nY = imShape[1] - kernel.shape[1] # maximum y grid value of the iteration over the grid

totalImages = len(np.arange(0,nY,slideVal))*len(np.arange(0,nX,slideVal)) # total number of images to process

imagePosition = [] # empty container
imageValue = [] # image value at given position for sanity check

# Main calculation loop ########################################################

massiveDump = [] # creatively named container for the final data to be exported as pickle object

imageNo = 0 # variable for images currently processed

loadBar.printProgress(imageNo,totalImages,suffix= str(imageNo) + ' of ' + str(totalImages)) # print percentage progress of the calculation

for j in np.arange(0,nY,slideVal): # for each y position
    for i in np.arange(0,nX,slideVal): # for each x position

        imagePosition.append([-imSize/2.0 + imPix*(i + kN/2),-imSize/2.0 + imPix*(j+ kN/2)]) # appending the centre position of each image
        imageValue.append(rawData[j+kN/2,i+kN/2]) # appending the image value at the centre of the kernel (useful for checking coordinate system)

        # Fourier processing of data

        kernel[:,:] = np.transpose(rawData[j:j+kN,i:i+kN]) # extract image from the main image into kernel
        kernelFFT[:,:] =  np.abs((np.fft.fft2(rawData[j:j+kN,i:i+kN]))) # perform 2D FFT on data and extract absolute values
        kernelFFT[:,:] = np.fft.fftshift(kernelFFT) #  shift zero frequency to the centre of the spectrum

        # generate 2D interpolation function for interpolation around the sub-image
        interpfunc = interpolate.interp2d(freqVals,freqVals,kernelFFT)

        FFTImage = [] # empty list for each line of the FFT image

        # interpolating around torous of 2DFFT image
        for theta in thetaVals: # for each angle around the torus - theta

            X = np.linspace(0.0,freqVals[-1]*math.cos(math.radians(theta)),np.shape(kernelFFT)[0]/2.0)# generate X vals for interpolation
            Y = np.linspace(0.0,freqVals[-1]*math.sin(math.radians(theta)),np.shape(kernelFFT)[1]/2.0)# generate Y vals for interpolation

            line = [] # empty list for each line
            masekdFreqVals =[] #

            for k in range(len(X)): # iterate through each frequency value - k

                freqInterpVal = np.sqrt(X[k]**2 + Y[k]**2) # the spatial frequency value at that position.

                # create each frequency element of the line

                if(np.logical_and(freqInterpVal >= min(freqRange),freqInterpVal <= max(freqRange))): # if frequency is within allowed range
                    line.append(interpfunc(X[k],Y[k])[0]) # interpolating 2DFFT at each frequency value
                    masekdFreqVals.append(freqInterpVal) # add


            FFTImage.append(line) # Append each line to image

        processedData = dict() # creating dict of all the processed data for each kernel

        processedData['maxProfile'] = normaliseProfile(np.max(FFTImage,1)) # averaging along the frequency direction
        processedData['meanProfile'] = normaliseProfile(np.mean(FFTImage,1)) # finding the maximum vals along the frequency direction

        # detecting and extracting values of peaks for later analysis

        processedData['maxPeakIndices'] = detect_peaks.detect_peaks(processedData['maxProfile'],mph=0.3,mpd=10) # find indices of peaks along averaged profiles
        processedData['meanPeakIndices'] = detect_peaks.detect_peaks(processedData['meanProfile'],mph=0.3,mpd=10) # find indices from maximum profile

        processedData['meanPeakAngles'] = [thetaVals[index] for index in processedData['meanPeakIndices']] #  extract the angle at which the peaks occur
        processedData['maxPeakAngles'] = [thetaVals[index] for index in processedData['maxPeakIndices']] #  extract the angle at which the peaks occur

        processedData['maxPeakFreqVals'] = [np.max(FFTImage[index]) for index in processedData['maxPeakIndices']] # find the actual value at the position at which the peaks occur
        processedData['maxPeakFreqIndex'] = [np.argmax(FFTImage[index]) for index in processedData['maxPeakIndices']] # find the index at which the peak frequency occurs
        processedData['maxPeakFreqs'] = [masekdFreqVals[index] for index in processedData['maxPeakFreqIndex']] # find the frequency at the position of the most intense peaks

        processedData['meanPeakFreqVals'] = [np.max(FFTImage[index]) for index in processedData['meanPeakIndices']] # find the actual value at the position at which the peaks occur
        processedData['meanPeakFreqIndex'] = [np.argmax(FFTImage[index]) for index in processedData['meanPeakIndices']] # find the index at which the peak frequency occurs
        processedData['meanPeakFreqs'] = [masekdFreqVals[index] for index in processedData['meanPeakFreqIndex']] # find the frequency at the position of the most intense peaks


        processedData['maxPeakFreq'] = processedData['maxPeakFreqs'][np.argmax(processedData['maxPeakFreqVals'])]# the spatial frequency of the peak with the maximum intensity
        processedData['maxPeakAngle'] = processedData['maxPeakAngles'][np.argmax(processedData['maxPeakFreqVals'])]# angle at which this peak occurs

        processedData['meanPeakFreq'] = processedData['meanPeakFreqs'][np.argmax(processedData['meanPeakFreqVals'])]# the spatial frequency of the peak with the maximum intensity
        processedData['meanPeakAngle'] = processedData['meanPeakAngles'][np.argmax(processedData['meanPeakFreqVals'])]# angle at which this peak occurs

        processedData['Number'] = imageNo

        # exporting the processed data #########################################

        # reducing disk space by not saving all the profiles i.e. making the lists empty

        processedData['maxProfile'] = []
        processedData['meanProfile'] = []

        massiveDump.append(processedData) # append the processed data dict object to the final object to be exported.

        # updating the progress of the main loop

        imageNo = imageNo + 1 # incrementing image number for progress update

        # print progress to terminal
        loadBar.printProgress(imageNo,totalImages,suffix= str(imageNo) + ' of ' + str(totalImages))

# End of main program loop #####################################################

# Informing the user that the program is about to export the data
print 'Taking Massive Dump!'

# writing data to Python pickle object for loading using analysis program

# NB: output file will be ~200 mb for these parameters and the image used in this work.

pickle.dump(massiveDump,open('outputData.pickle','wb')) # change file name in analysis program

print 'Data dump successfull!!!''

print 'Exporting position and height values'
np.savetxt('Positions.csv',imagePosition,delimiter=',') # saving the x,y positions of all the images
np.savetxt('Values.csv',imageValue,delimiter=',') # saving the positions of all the images
print 'Position and height values saved!'

print 'End of Program!'

# End of Program ###############################################################