; Copyright (c) 1998-2017 A.P. Hitchcock  All rights reserved
;+
;NAME:
;		AX_FRC.PRO
;
;LAST CHANGED: ----------------------------------- 	01 Apr 2017
;
; PURPOSE:
;	This procedure computes the Fourier Ring Correlation of an image
; as a means to estimate spatial resolution
; translated in Feb-17 from FRC.py written by Xiaohui Zhu (McMaster) Jan 2017
;
; CATEGORY:
;	Image processing. runs stand alone or called from aXis2000.
;
; CALLING SEQUENCE:
;	Result = AX_FRC(image=image, file=file, group=groupID)
;
; INPUTS: no positional inputs
;
; KEYWORDS:
;	IMAGE:  image from inside axis2000 if start from a buffer with an image
;	FILE: 	filename (if not provided user is prompted
;
; OUTPUTS:
;	overplot of the FRC output, a smoothed version and the 1/2-bit and 0.5 lines
;
; COMMON BLOCKS:
;	AXIS_COM	standard set of common blocks
;
; SIDE EFFECTS:
;	A window is created to diplay results if run outside of aXis2000
;
; MODIFICATION HISTORY:
; (01 Apr-17 aph) first version, adapted from XHZ python version
;       - NB  frc.py python  code is at end;  does not agree with ZHZ result
;-

 function ax_find_ring_area, map_dist, r, width_ring
 ; from XHZ, Jan2017 python
  on_error,2
  @axis_com

  check_r = where(map_dist GE r AND map_dist LE r+width_ring, countr)
  ring_area = reverse(array_indices(map_dist, check_r))
;  print, 'ax_find_ring_area ', size(ring_area,/dimensions)
  return, ring_area
end

;------------------------------------------------------------------------
;
;------------------------------------------------------------------------
;
function ax_frc,image=image, file=file, check=check, group=group
on_error,2
@axis_com

s = ''
;goto, test			; BY-PASS to ensure test on same jpg file python uses
if keyword_set(image) then begin
; ---------- image supplied from axis2000 buffer
	if image.t EQ '2d' then begin
		tmp = image
		im = image.d
		fileshort = strmid(image.dl,0,strpos(image.dl,' '))
		if n_elements(lPath) GT 0 then path = lPath
		if n_elements(WritePath) GT 0 then path = WritePath
		if n_elements(DefPath) GT 0 then path = DefPath
		if n_elements(path) EQ 0 then path = CodePath
;		print, 'Path is ', path
	endif else begin
		axis_log, 'FRC: only works with images'
		return, s
	endelse
endif else begin
; ---------- get image from file
	IF NOT keyword_set(file) then begin
		fltr='*.axb'
		if keyword_set(nof) then fltr='*'
		file = pickfile2(/READ, title=title, FILTER=fltr, /LPATH, DEFPATH=defpath)
	ENDIF

	if file EQ '' THEN RETURN, s  ; bail-out if no filename
	test = findfile(file)		 ; bail-out if non-existent file
	if strlen(test(0)) EQ 0 then begin
		axis_log, 'Cannot find ' + file
		return, s
	endif

	tmp = axb_load(file=file)
;	print, 'Read axb file ', file
	t = ax_name(file)
	fileshort = t(1)
	path = t(0)
	im = tmp.d
endelse

; ---- test equivlence to python
test:
;path = 'Y:\aXis-Pass-Test-data\04-IMAGES\13a-FRC\'
;path = 'E:\aXis-Pass-Test-data\04-IMAGES\13a-FRC\'
;name = 'C160916086' & pixnm = 250.
;name = 'raw709m' & pixnm = 8.3
;name = '11_140828024-mag-sq' & pixnm = 6.6
;name = 'tmp'  &  pixnm= 6.71352
;read_jpeg,path+name+'.jpg', im
;fileshort = name
;t = size(im, /dimensions) & nx = t(0) & ny = t(1)
;x = findgen(nx)*pixnm/1000. & y = findgen(ny)*pixnm/1000.
;tmp = {t:2d, x:x, y:y, d:im,  xl:'um', yl:'image', dl:fileshort}

; ------ force size so that 1/2 images are square with an odd number of pixels (then FFT centre is integer)
 nx = n_elements(tmp.x)  & ny = n_elements(tmp.y)
if nx EQ ny then goto, skip_square
if float(nx)/float(ny) GT 0.99 AND float(nx)/float(ny) LT 1.01 then begin
	make_square, tmp
	nx = n_elements(tmp.x)  & ny = n_elements(tmp.y)
endif else begin
	axis_log, 'Image must be square - use Zoom~numerical'
	return, s
endelse
skip_square:

im = tmp.d
axis_log, 'Image size is ' + strtrim(string(nx),2) +' * ' + strtrim(string(ny),2)
pixnm = 1000.*(tmp.x(1) - tmp.x(0))
axis_log, 'pixel size (nm) is ' +  strtrim(string(pixnm),2)

; ------- split the original image into two subsets, even (i+j) and odd(i+j)
nx = n_elements(tmp.x)  & ny = n_elements(tmp.y)
hx = nx/2   & hy = ny/2
imo = fltarr(hx, hy)
ime = fltarr(hx, hy)
for i = 0, nx - 1 do begin
	for j = 0, ny - 1 do begin
		if (i+j) mod 2 eq 1 then begin
			if i/2 LT hx AND j/2 LT hy then imo(i/2,j/2) = im(i,j)
		endif else begin
			if i/2 LT hx AND j/2 LT hy then ime(i/2,j/2) = im(i,j)
		endelse
	endfor
endfor
; help, hx, hy, ime, imo
ime = rotate(ime,7)  & imo = rotate(imo,7)
ime=shift(ime,-1,0)
; write_tiff, path+'IDL-half_even-'+fileshort+'.tif', ime
; write_tiff, path+'IDL-half_odd-'+fileshort+'.tif' , imo

; --------- shift alonng appropriate axis to match (even: shift column by 1 to right)
image1 = shift(imo,0,0)
image2 = shift(ime,-1,0)		; NB in python code, this shift is to axis 1	;

; ----- test ti see if IDENTICAL 1/2 images gives same result
;image2 = read_tiff(path+'half_even-'+fileshort+'.tif')
;image1 = read_tiff(path+'half_odd-'+fileshort+'.tif')

; --------- Get the Nyquist frequency
nx = hx & ny = hy
if nx GT ny then begin
    image1=imo[0:ny-1,0:ny-1]
    image2=ime[0:ny-1,0:ny-1]
    maxSize=ny
endif else begin
    image1=imo[0:nx-1,0:nx-1]
    image2=ime[0:nx-1,0:nx-1]
   maxSize=nx
endelse

freq_nyq = fix(floor(maxSize/2.0))
n_mesh = ceil(maxSize/2.0) + floor(maxSize/2.0)
; ---------  Create Fourier grid
t = -1.*floor(maxSize/2.0) + findgen(n_mesh)
nt = n_elements(t)
x = fltarr(nt,nt)
y = fltarr(nt,nt)
for i = 0, nt-1 do x(*,i) = t
for i = 0, nt-1 do y(i,*) = shift(t,0)
map_dist = fltarr(n_mesh,n_mesh)

for i = 0, n_mesh-1 do begin
	for j = 0, n_mesh-1 do begin
		map_dist(i,j) = sqrt(x(i,j)*x(i,j) + y(i,j)*y(i,j))
	endfor
endfor
;write_tiff, path+'IDL-map_dist-'+fileshort+'.tif', map_dist
;help, map_dist
;print, 'map_dist ', map_dist

; ----- FFT transforms of the  input images
fft_image1=shift(fft(image1),hx/2, hy/2)	; IDL's FFT routine puts zero freq at output array(i,j) with i = j = 0
fft_image2=shift(fft(image2),hx/2, hy/2)

; --- check if the FFT looks the same
if keyword_set(check) then begin
	mag1 = alog10(abs(fft_image1))
	t = size(mag1, /dimensions)
	xm = findgen(t(0))  &  ym = findgen(t(1))
	FFT1={t:'2d', x:xm, y:ym, xl:'inverse nm', yl:'FFT1 - half-odd', d:mag1, dl:fileshort + ' FFT1 half odd '}
	fftfile=path+fileshort+'-FFT1-odd.axb'
;	fftfile = dialog_pickfile(file = fftfile, /write, title='Name for FFT of ODD-pixels of image')
	if fftfile NE '' then tmp = axb_save(FFT1, file=fftfile)

	mag2 = alog10(abs(fft_image2))
	t = size(mag2, /dimensions)
	xm = findgen(t(0))  &  ym = findgen(t(1))
	FFT2= {t:'2d', x:xm, y:ym, xl:'inverse nm', yl:'FFT2 - half-even', d:mag2, dl:fileshort + ' FFT2-even'}
	fftfile=path+fileshort+'-FFT2-even.axb'
;	fftfile = dialog_pickfile(file = fftfile, /write, title='Name for FFT of even-pixels of image')
	if fftfile NE '' then tmp = axb_save(FFT2, file=fftfile)
endif

; ------- Allow user to modify thickness, and size of the rings	 ;what is basis of these values ???

if n_elements(xlock) EQ 0 then width_ring = 0.8 else width_ring = xlock(1)
if width_ring EQ 0 then width_ring = 0.8	; because xlock is set to [0,0]
if keyword_set(group) then begin
	width_ring = get_num(prompt='FRC - ring width (0.2 - 1)', val = width_ring, group=group)
endif else  width_ring = get_num(prompt='FRC - ring width (0.2 - 1)', val = width_ring)
xlock(1) = width_ring
if n_elements(ylock) EQ 0 then radius = 3.0 else radius  = ylock(1)
if radius EQ 0 then radius = 3.0					; because ylock is set to [0,0]
if keyword_set(group) then begin
	radius = get_num(prompt='FRC - ring width (0.2 - 1), def = 0.8', val = radius, group=group)
endif else  radius = get_num(prompt='radius (1 - 5), def = 3.0', val = radius)
ylock(1) = radius
axis_log, 'ring width = ' + string(width_ring, format='(F4.1)')
axis_log, ' radius = ' + string( radius, format='(F4.1)')

; ---------  Calculate FRC
t = 0
; while t LT 2 do begin
while radius + width_ring LT freq_nyq do begin
    ring = ax_find_ring_area(map_dist , radius , width_ring )
;	print, t, '   ring ', size(ring, /dimensions)
    aux1 = reform(fft_image1[ring[0,*],ring[1,*]])
    aux2 = reform(fft_image2[ring[0,*],ring[1,*]])
;    print,t,'  aux1 ', size(aux1, /dimensions)
; --------------------
; ----   #FRC=aux1* conjugate(aux2)/(aux1)**2 * aux2**2 (was commented out in XHZ version)
; --------------------
    if t EQ 0 then C1 = total(aux1 * conj(aux2)) else $
	    C1 = [C1,total(aux1 * conj(aux2))]
;    print, t, 'C1  ', C1
	if t EQ 0 then C2 = total(abs(aux1)*abs(aux1))else $
	    C2 = [C2,total(abs(aux1)*abs(aux1))]
;    print,  'C2  ', C2
	if t EQ 0 then C3 = total(abs(aux2*abs(aux2)))else $
	    C3 = [C3,total(abs(aux2)*abs(aux2))]
;    print,  'C3  ',C3
    if t EQ 0 then n = n_elements(aux1) else n = [n,n_elements(aux1)]
    t = t+1
    radius = radius + width_ring
    FRC = abs(C1)/sqrt(C2*C3)
;    print, t, ' FRC ', size(FRC, /dimensions)
endwhile

; ---------- smooth using Savitsky-Golay routine
savgolFilter = SAVGOL(3, 3, 0, 2)
FRCs = CONVOL(FRC, savgolFilter, /EDGE_TRUNCATE)

; ------------- generate the half-bit line
n = n[1:n_elements(n)-1]
half_bit=fltarr(1)
for i =0, n_elements(n)-1 do begin
    y = (0.2071 + 1.9102/sqrt(n(i))) / (1.2071 + 0.9102/sqrt(n(i)))
 	half_bit =[half_bit,y]
endfor
half_bit = half_bit[1:n_elements(half_bit)-1]

; -------- prepare FRC, and associated plots
nbins=n_elements(FRC)
xmax=1.414/pixnm
xstep=xmax/nbins
x = findgen(nbins)*xstep
half_height = fltarr(nbins) + 0.5
FRCaxb={t:1d, x:x, d:FRC, dn:FRC, xl:'inverse nm', yl:'Fourier Ring Correlation', dl:fileshort + ' FRC'}
FRCaxbs={t:1d, x:x, d:FRCs, dn:FRCs, xl:'inverse nm', yl:'Fourier Ring Correlation smoothed', dl:fileshort + ' FRC-S7'}
halfh= {t:1d, x:x, d:half_height, dn:half_height, xl:'inverse nm', yl:' ', dl:fileshort + ' half-height'}
halfb= {t:1d, x:x, d:half_bit, dn:half_bit, xl:'inverse nm', yl:' ', dl:fileshort + ' half-bit'}

if keyword_set(group) then begin

	HANDLE_VALUE, Data(7), FRCaxb, /SET
	Label(7) = FRCaxb.dl
	Plotbuf, 7
	HANDLE_VALUE, Data(8),halfh, /SET
	Label(8) = halfh.dl
	Plotbuf, 8
	HANDLE_VALUE, Data(9),  halfb, /SET
	Label(9) = halfb.dl
	PlotBuf, 9
	HANDLE_VALUE, Data(0),FRCaxbs, /SET
	CurBuf = 0
	Label(0) = FRCaxbs.dl
	PlotBuf, 0
;	pBufs = [0,7,8,9]
;	Refresh
endif
; plot  outside aXis2000
WINDOW, /free, title='Fourier ring correlation'
splot, FRCaxbs, thick=4, /zero
splot, FRCaxb, psym=1, symsize =2, /o
splot, halfh, /o
splot, halfb, /o, thick = 2, color = 64


return, FRCaxbs
end


;---------------------------- complete python code -------------------------
;# -*- coding: utf-8 -*-
;"""
;Created on Wed Jan 18 13:28:00 2016
;@author: Xiaohui Zhu
;"""
;from numpy import fft
;import numpy as np
;import pylab as plt
;from PIL import Image
;
;# Read the ptychography image
;image_path='c:\\Downloads\\FRC_ALS\\' # Image location
;image_name='fuel-cell-8bit' # Image name
;im=Image.open(image_path+image_name+'.tif')
;print "Size of the image is: "+ str(im.size)
;
;# split the original images into two subsets
;
;# get the even subset
;half_even=Image.new('L',tuple([int(d/2) for d in im.size]))
;try:
;    for i in range(im.size[0]):
;        for j in range(im.size[1]):
;            if i%2==0:
;                half_even.putpixel((int(i/2),int(j/2)),im.getpixel((i,j)))
;    half_even.save(image_path+'half_even-'+image_name+'.tif')
;except IndexError:
;    print 'Note: they are not both even numbers'
;
;# get the odd subset
;try:
;    half_odd=Image.new('L',tuple([int(d/2) for d in im.size]))
;    for i in range(im.size[0]):
;        for j in range(im.size[1]):
;            if i%2!=0:
;                half_odd.putpixel((int(i/2),int(j/2)),im.getpixel((i,j)))
;    half_odd.save(image_path+'half_odd-'+image_name+'.tif')
;except IndexError:
;    print 'Note: they are not both even numbers'
;
;## Read the odd and even subsets
;image1 = np.roll(half_odd,0,axis=0)
;image2 = np.roll(half_even,1,axis=1)
;
;#  Get the Nyquist frequency
;ny,nx=image1.shape
;if nx>ny:
;    image1=image1[0:ny,0:ny]
;    image2=image2[0:ny,0:ny]
;    maxSize=ny
;else:
;    image1=image1[0:nx,0:nx]
;    image2=image2[0:nx,0:nx]
;    maxSize=nx
;freq_nyq=int(np.floor(maxSize/2.0))
;
;##  Create Fourier grid
;x = np.arange( -np.floor( maxSize/2.0 ) , np.ceil( maxSize/2.0 ) )
;y = np.arange( -np.floor( maxSize/2.0 ) , np.ceil( maxSize/2.0 ) )
;x,y = np.meshgrid( x , y )
;map_dist = np.sqrt( x*x + y*y )
;
;##  FFT transforms of the input images
;fft_image1=fft.fftshift(fft.fftn(image1))
;fft_image2=fft.fftshift(fft.fftn(image2))
;
;## Smooth the curve
;def savitzky_golay(y, window_size, order, deriv=0, rate=1):
;    from math import factorial
;    try:
;        window_size = np.abs(np.int(window_size))
;        order = np.abs(np.int(order))
;    except ValueError:
;        raise ValueError("window_size and order have to be of type int")
;    if window_size % 2 != 1 or window_size < 1:
;        raise TypeError("window_size size must be a positive odd number")
;    if window_size < order + 2:
;        raise TypeError("window_size is too small for the polynomials order")
;    order_range = range(order+1)
;    half_window = (window_size -1) // 2
;    # precompute coefficients
;    b = np.mat([[k**i for i in order_range] for k in range(-half_window, half_window+1)])
;    m = np.linalg.pinv(b).A[deriv] * rate**deriv * factorial(deriv)
;    # pad the signal at the extremes with
;    # values taken from the signal itself
;    firstvals = y[0] - np.abs( y[1:half_window+1][::-1] - y[0] )
;    lastvals = y[-1] + np.abs(y[-half_window-1:-1][::-1] - y[-1])
;    y = np.concatenate((firstvals, y, lastvals))
;    return np.convolve( m[::-1], y, mode='valid')
;
;##  Get thickness, and size of the rings
;width_ring = 0.8
;r = 3.0
;
;def find_ring_area( map_dist, r, width_ring):
;    check_r=map_dist>=r
;    check_r_width=map_dist<=r+width_ring
;    ring_area = np.argwhere( check_r * check_r_width)
;    return ring_area
;
;##  Calculate FRC
;C1 = []
;C2 = []
;C3 = []
;n = []
;
;while r + width_ring < freq_nyq :
;    ring = find_ring_area( map_dist , r , width_ring )
;    aux1 = fft_image1[ring[:,0],ring[:,1]]
;    aux2 = fft_image2[ring[:,0],ring[:,1]]
;    #FRC=aux1* conjugate(aux2)/(aux1)**2 * aux2**2
;    C1.append( np.sum( aux1 * np.conjugate(aux2) ) )
;    C2.append( np.sum( np.abs( aux1 )**2 ) )
;    C3.append( np.sum( np.abs( aux2 )**2 ) )
;    n.append(len(aux1))
;    r += width_ring
;    n = np.array( n )
;    n=list(n)
;    FRC = np.abs( np.array( C1 ) )/ np.sqrt( np.array( C2 ) * np.array( C3 ) )
;    FRC = savitzky_golay(FRC, 7, 1)
;half_bit=[]
;for i in n:
;    y= (0.2071 + 1.9102 / np.sqrt( i ) ) / ( 1.2071 + 0.9102 / np.sqrt( i ) )
;    half_bit.append(y)
;half_bit=np.array(half_bit)
;
;# Pixel size read from ptycho image:
;pixnm = np.float(8.3) #Change it for different images
;nbins=len(FRC)
;freqs = np.linspace(0,1.414 / pixnm, nbins)
;half_height = np.ones((nbins)) * 0.5
;
; # plot results
;plt.plot(freqs,FRC,marker='o',mfc='blue',mec='white',lw=1,linestyle='-',color='black')
;plt.plot(freqs,half_height,linestyle = '--', color = 'grey')
;plt.plot(freqs,half_bit,'r-')
;axes = plt.gca()
;axes.set_xlim([0,0.18])
;axes.set_ylim([0,1.05])
;plt.show()
;f = open(image_name+'-frc.txt','w')
;for i in range(nbins - 1):
;    f.write(str(FRC[i]) + '\n')
;f.close()
;f = open(image_name+'-freqs.txt','w')
;for i in range(nbins - 1):
;    f.write(str(freqs[i]) + '\n')
;f.close()
;f = open(image_name+'-half_height.txt','w')
;for i in range(nbins - 1):
;    f.write(str(half_height[i]) + '\n')
;f.close()
;f = open(image_name+'-half_bit.txt','w')
;for i in range(nbins - 1):
;    f.write(str(half_bit[i]) + '\n')
;f.close()
;im.show() # Show which image processed\
