# polymerAnalysis.py Program for data analysis and plotting 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 used to plot data previously processsed using 'polymerCrunch.py'

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

# importing libraries

import numpy as np
import math

import sys
import os
import glob
import cPickle as pickle
import itertools
import random


from matplotlib import pyplot as plt # matplitlib for plotting: www.matplotlib.org
from matplotlib import cm # colourmap import
from scipy import interpolate # interpolation function from SciPy: www.scipy.org
from scipy.ndimage.filters import gaussian_filter # as above
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

# loading data files

fname = 'polymerPhaseImage.txt' # file name of the raw data used for analysis
rawImage = np.genfromtxt(fname) # loading in raw data (used for overlay later)
positions = np.genfromtxt('Positions.csv',delimiter=',') # loading the centroid locations of every FFT
imageValues = np.genfromtxt('Values.csv',delimiter=',') # loading the values at each position

# Defining constants and variables for later use ###############################

imSize = 100.0e-9 # size of the image (m)
imShape = np.shape(rawImage) # shape of input arrays
imPix = imSize/imShape[0] # spatial size of a single pixel (m)

# generating meshgrids
imGridx,imGridy = np.meshgrid(np.linspace(-imSize/2,imSize/2,np.shape(rawImage)[1]),np.linspace(-imSize/2,imSize/2,np.shape(rawImage)[1]))

slideVal = 5 # number of pixels betwwen each step during sliding
kNDesired = 8e-9 # defining kernel size

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 FFT calculation quicker
    kN = kN + 1
kSize = kN * imPix # kernel size (m)

# extracting upper and lower bounds of data

minXPos = np.min([position[0] for position in positions])
minYPos = np.min([position[1] for position in positions])

maxXPos = np.max([position[0] for position in positions])
maxYPos = np.max([position[1] for position in positions])

Xim = np.linspace(-imSize/2,imSize/2,np.shape(rawImage)[1])
Yim = np.linspace(-imSize/2,imSize/2,np.shape(rawImage)[0])

noPoints = 300 # size for image interpolation

x_grid,y_grid = np.meshgrid(np.linspace(minXPos,maxXPos,noPoints),np.linspace(minYPos,maxYPos,noPoints))

imSize = 100.0e-9 # dividing the number of pixels by the splitting of the main image

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 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
freqVals = np.linspace(freqRange[1],freqRange[0],kN)

processedData = []

print "\nLoading in raw data from exported files:\n"

processedData = pickle.load(open('outputData.pickle.','rb'))
interpFreqVals = np.linspace(np.min(freqRange),np.max(freqRange),kN/2.0)

print "\nLoading complete!\n"

# main processing loop #########################################################

print "\nProcessing Data:\n"

meanFreqHistData = list(itertools.chain.from_iterable([data['meanPeakFreqs']for data in processedData]))
meanAngleHistData = list(itertools.chain.from_iterable([value for value in [data['meanPeakAngles']for data in processedData]]))

maxFreqHistData = list(itertools.chain.from_iterable([value for value in [data['maxPeakFreqs']for data in processedData]]))
maxAngleHistData = list(itertools.chain.from_iterable([value for value in [data['maxPeakAngles']for data in processedData]]))

meanPeriodHistData = np.divide(1,meanFreqHistData)
maxPeriodHistData = np.divide(1,maxFreqHistData)

# isolating max strength fourier components in 0-180 degree interval sue to symmetry

maskedMaxAngles = [np.array(data['maxPeakAngles']) <= 180 for data in processedData]

# selecting out the relevant peaks for analysis

strongestPeakAngle =[]
strongestPeakFrequency = []
angleMask2 =[]
strongestPeakPeriodicity = []

secondPeakAngle = []
secondPeakFrequency = []
secondPeakPeriodicity = []
filteredPositions = []

i = 0
for data in processedData:

    intensitiesArgs = np.argsort(np.argsort(data['maxPeakFreqVals'])) # finding the order of the intensities i.e. list the rankings of wach element
    anglesMask = np.array(data['maxPeakAngles']) <= 180 # finding logic mask array for the relevant elements in range 0-180 degrees due to symmetry
    maskedIndexes = np.arange(0,len(intensitiesArgs))[anglesMask] # masked array of indexes

    strongestPeakArg = maskedIndexes[np.argsort(intensitiesArgs[anglesMask])[-1]] # get the last element of the masked array, i.e. the maximum intensity of the mnaximum allowed angle.

    # determining if peak is a root 2 distance at 45 degrees.

    if (1/data['maxPeakFreqs'][strongestPeakArg]) < 0.6e-9:
        strongestPeakAngle.append((data['maxPeakAngles'][strongestPeakArg])%90)
    else:
        strongestPeakAngle.append((data['maxPeakAngles'][strongestPeakArg] - 45)%90)


    strongestPeakFrequency.append(data['maxPeakFreqs'][strongestPeakArg])
    strongestPeakPeriodicity.append(1/data['maxPeakFreqs'][strongestPeakArg])

    try:
        secondPeakArg = maskedIndexes[np.argsort(intensitiesArgs[anglesMask])[-2]] # get the second most intense peak

        secondPeakAngleValue = data['maxPeakAngles'][secondPeakArg]
        secondPeakFrequencyValue = data['maxPeakFreqs'][secondPeakArg]
        secondPeakPeriodicityValue = 1/data['maxPeakFreqs'][secondPeakArg]

        secondPeakAngle.append(secondPeakAngleValue)
        secondPeakFrequency.append(secondPeakFrequencyValue)
        secondPeakPeriodicity.append(secondPeakPeriodicityValue)

        filteredPositions.append(positions[i]) # adding the position

    except: # if there is a problem with doing this default to 0 for all values

        secondPeakAngleValue = 0
        secondPeakFrequencyValue = 0
        secondPeakPeriodicityValue = 0.0

    #secondPeakSearchAngle = (data['maxPeakAngles'][strongestPeakArg] + 90) % 180 # generating secondary search angle for other peak within 0-180 deg

    i = i+1 # incrementing positionsn

# final plotting of the data


# grid of the colors of the angle at each position
zGrid = interpolate.griddata(positions,strongestPeakAngle,(x_grid,y_grid))
periodGrid = interpolate.griddata(positions,strongestPeakPeriodicity,(x_grid,y_grid))


# Main figure (Figure 4c) ######################################################

plt.figure(10,frameon=False)

ampFig = plt.imshow(rawImage,extent=(-imSize/2,imSize/2,-imSize/2,imSize/2),cmap='gray',interpolation='none')
orientMap = plt.imshow(np.flipud(zGrid),extent=(minXPos,maxXPos,minYPos,maxYPos),alpha = 0.5,cmap=plt.cm.get_cmap('hsv'))

interpPos =[[-3.923e-8,-1.833e-8,9.08e-9,1.744e-8],[-3.5e-8,4.0e-8,-1.9e-8,-1.9e-8]] # coordinates for line profiles - OLD
#interpPos =[[-3.923e-8,-1.696e-8,9.08e-9,1.605e-8],[-3.5e-8,4.0e-8,-1.9e-8,-1.9e-8]] # positions for line profiles - new for phase image

boxPos = [[-4.22e-8,-1.31e-8,6.39e-9,2.15e-8],[-3.28e-8,-1.00e-8,-2.48e-8,-1.42e-8],[-3.57e-8,-7.29e-9,2.58e-8,3.79e-8]] # coordinates of zoom boxes [xMin,xMax,yMin,yMax]

boxDists = [[boxPos[0][1] - boxPos[0][0]],[boxPos[1][1] - boxPos[1][0]],[boxPos[2][1] - boxPos[2][0]]]


for pos in interpPos: # for each interpolation position

    plt.figure(10) # switch to main figure

    interpPositions = [np.linspace(pos[0],pos[1],100),np.linspace(pos[2],pos[3],100)] # generate coordinates for plotting the line
    orientFunction = interpolate.RectBivariateSpline(np.linspace(minXPos,maxXPos,noPoints),np.linspace(minYPos,maxYPos,noPoints),np.transpose(zGrid)) # create function for interpolation
    interpVals = np.diag(orientFunction(interpPositions[0],interpPositions[1])) # interpolate at the given positions

    profileDistance = np.linspace(0,np.linalg.norm([pos[1]-pos[0],pos[3]-pos[2]]),100)

    # Plotting the position of the line profile on Fig. 4c #####################

    plt.plot(pos[0:2],pos[2:4],color='white',linewidth=3,linestyle='--')

    # Plotting the interpolated angle along the line profile (Figs. 4g,h) ######

    plt.figure(figsize=(2.0,2.0))
    plt.plot(profileDistance*1e9,interpVals)

    plt.xlabel('Distance (nm)')
    plt.ylabel(r'Angle ($\theta^{\circ}$)')

    plt.xticks(np.linspace(0,np.ceil(np.max(profileDistance*1e9)/10)*10,3))

    plt.ylim([0,90])
    plt.tight_layout()

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

plt.figure(10)

plt.xlim([minXPos,maxXPos])
plt.ylim([minYPos,maxYPos])

plt.axis('off')

orientCbar = plt.colorbar(orientMap,use_gridspec=True)
orientCbar.set_alpha(1)
orientCbar.set_label(label=r'Angle ($\theta^\circ$)',fontsize=12)
orientCbar.draw_all()

plt.tight_layout()
fig = plt.gcf()
fig.set_size_inches(5.5,5.5)


# Plot crops of Fig 4c, Figs 4d,e,f ############################################

# plot zoomed images for sum figures

i = 0
for box in boxPos: # for each position

    plt.figure(figsize=(5.5,3.7))

    plt.imshow(rawImage,extent=(-imSize/2,imSize/2,-imSize/2,imSize/2),cmap='gray',interpolation='none')
    plt.imshow(np.flipud(zGrid),extent=(minXPos,maxXPos,minYPos,maxYPos),alpha = 0.5,cmap=plt.cm.get_cmap('hsv'))

    #plt.plot(interpX,interpY,color='white',linewidth=3)

    plt.axis('off')

    plt.tight_layout()

    plt.xlim([box[0],box[1]])
    plt.ylim([box[2],box[3]])

    i += 1 # increment number


plt.show() # showing all figures

# End of script ################################################################