;-----------
; kngao_rtc.pro
;  Simulation of the FPGA implementation of multi-guidestar tomography
;  kngao = Keck Next Generation Adaptive Optics
;  rtc = Real-time control

; one voxel unit contains:
;   the Fourier ccmponents of delta-index
;   the Fourier factors needed to do the transform
;   Cn2 for this layer
;   shift factors for each guidestar
;   the cone beam scale factor for this altitude (assuming 90km sodium guidestars)
;   spatial frequency address (x and y)
;   factor for post-conditioning (k^(11/3))

; one wfs pixel unit contains:
;   the Fourier components of delta-phase
;   the Fourier factors needed to do the transform
;   spatial frequency address (x and y)
;   complex matrix elements for pre-conditioning
;   spatial domain mask for pupil (0 or 1)

; one DM pixel unit contains:
;   the Fourier components of the actuator commands
;   plate-equation Fourier factors
;   actuator 2-entry look-up table
;   alternative to plate equation: Fourier factors for inverting the DM influence function

pro mpp_transform,voxel_array
  s = size(voxel_array)
  d = s[0]
  n = s[1]
  if (d eq 2) then begin
    sx = [-1,0]
    sy = [0,-1]
  endif else begin 
    sx = [-1,0,0]
    sy = [0,-1,0]
  endelse
  ; ====== The mpp Fourier Transform is 4 lines of code ========
  ; ====== done in 2n steps ====================================
  print,'<mpp_transform> doing x transforms' & wait,.01
  for ix = 0,n-1 do begin
    voxel_array.accum1 += voxel_array.phi*voxel_array.c[ix,0]
    voxel_array.phi = shift(voxel_array.phi,sx)
  endfor
  print,'<mpp_transform> doing y transforms' & wait,.01
  for iy = 0,n-1 do begin
    voxel_array.fphi += voxel_array.accum1*voxel_array.c[iy,1]
    voxel_array.accum1 = shift(voxel_array.accum1,sy)
  endfor
  ; ============================================================
  print,'<mpp_transform> done' & wait,.01
end

pro mpp_backpropagate,voxel_array
  s = size(voxel_array)
  dimen = s[0]
  if (dimen eq 2) then n_layers = 1 else n_layers = s[3]
  ; ===== the mpp back propagation is 3 lines of code =====
  ; ===== done in n_layers steps ==========================
  for layer = 0,n_layers-1 do voxel_array[*,*,layer].s = shift(voxel_array[*,*,layer].s,[-layer,0,0])
  voxel_array.accum1 = voxel_array.fwfs
  for layer = 0,n_layers-1 do begin
    voxel_array.fphi += voxel_array.accum1*voxel_array.s[0]
    voxel_array.accum1 = shift(voxel_array.accum1,[0,0,1])
    voxel_array.s = shift(voxel_array.s,[1,0,0,0])
  endfor
  for layer = 0,n_layers-1 do voxel_array[*,*,layer].s = shift(voxel_array[*,*,layer].s,[layer,0,0])
  ; ==========================================================
end

pro mpp_forwardpropagate,voxel_array
  s = size(voxel_array)
  dimen = s[0]
  if (dimen eq 2) then n_layers = 1 else n_layers = s[3]
  n_guidestars = (size(voxel_array[0,0,0].s))[1]
  voxel_array.fwfs *= 0
  ; ===== the mpp forward propagation is 3 lines of code =====
  ; ===== done in n_guidestars * n_layers steps ==============
  for guidestar = 0,n_guidestars-1 do begin
    voxel_array.accum1 = voxel_array.fphi*voxel_array.s[guidestar]
    for layer = 0,n_layers-1 do begin
      voxel_array[*,*,guidestar].fwfs += voxel_array[*,*,guidestar].accum1
      voxel_array.accum1 = shift(voxel_array.accum1,[0,0,-1])
    endfor
  endfor
  ; ==========================================================
end

function ftcube,a,inverse=inverse
  if (n_elements(inverse) eq 0) then inverse = 0
  m = (size(a))[3]
  r = complex(a)
  for k = 0,m-1 do r[*,*,k] = ft(a[*,*,k],inverse = inverse)
  return,r
end

;========== supervisory computer =============
; ------- define physical constants ------
i = complex(0.,1.)
arcsec = !pi/(180.*3600.)
km = 1.e3

; ------- define the problem ---------
n_layers = 5 ; number of atmospheric layers
h = [0,1.,2.,4.,10.]*km
h = [0,4,8,12,16]*km
percent_atmos = [0.517+0.119,0.063,0.061,0.105,0.081+0.054] ; Mauna Kea Ridge Model, Dekany KAON #503 (collapsed to 5 layers -DG)
n_guidestars = 4 ; number of LGS guidestars
gs_angle = [[1.,1.],[1.,-1],[-1.,-1.],[-1.,1.]]*15.*arcsec/sqrt(2.)
na_across = 64 ; number of actuators across DM
ns_across = 48 ; number of subapertures across WFS
d_tele = 10. ; diameter of telescope, in meters
d_subap = d_tele/ns_across ; diameter of a WFS subaperature, in meters
dx = d_subap
n = nextpow2(2*ns_across) ; number of voxel processors on a layer
; load inital physical parameters, x,y and kx,ky axes
dk = 1./(dx*n)
kx = (findgen(n) - n/2)*dk
ky = kx
x = (findgen(n) - n/2)*dx
y = x
; -------------------------------------------------------------------------

; ---- define the aperture ----
;   the aperture will contain an indication of subap % illumination by pupil
n_apsample = 1024
ap = keckap(n_apsample,du=d_tele/(n_apsample/2))
ap = smooth(float(ap),n_apsample/n)
ap = rebin(ap,n,n)
; -----------------------------

; ------ define a voxel structure ----------
; one voxel unit contains:
;   phi: the delta-index
;   fphi: the Fourier components of delta-index
;   wfs: the phase at the wavefront sensor  (one wfs assigned per layer, assuming n_layers > n_guidestars)
;   fwfs: the Fourier components of phase at the wavefront sensor
;   accum1: a temporary accumulator
;   c: the Fourier factors needed to do the transform (x and y)
;   Cn2: Cn2 for this layer
;   s: complex shift factors for each guidestar: exp(i*2*pi*(sx*kx+sy*ky))
;   sf: the cone beam scale factor for this altitude (assuming 90km sodium guidestars)
;   k: spatial frequency address (x and y)
;   p: factor for post-conditioning (k^(11/3))
voxel_definition = {voxel, phi:real(0.), accum1:complex(0.), fphi:complex(0.), wfs:float(0.), fwfs:complex(0.), c: complexarr(n,2), Cn2:float(1.), $
            s: complexarr(n_layers), sf: float(1.), k: fltarr(2), p: float(0.)}
voxel_array = replicate({voxel},n,n,n_layers)
; ------------------------------------------

; ------- load the voxel array -----------
; load up the Fourier factors - note: later we will re-shuffle storage
; to interleave + and - spatial frequencies, but for now, we'll do the usual 4-quadrant organization
print,'loading Fourier factors' & wait,.1
t = systime(1)
for layer = 0, n_layers-1 do begin ; load up with...
  for ix = 0,n-1 do voxel_array[ix,0,layer].k[0] = kx[ix] ; spatial frequency addresses
  for iy = 0,n-1 do voxel_array[0,iy,layer].k[1] = ky[iy]
  for ix = 0,n-1 do begin
    voxel_array[ix,0,layer].c[*,0] = exp(i*2*!pi*x*kx[ix]) ; Fourier factors
    x = shift(x,-1)
  endfor
  for iy = 0,n-1 do begin
    voxel_array[0,iy,layer].c[*,1] = exp(i*2*!pi*y*ky[iy])
    y = shift(y,-1)
  endfor
;  now replicate them across the array
  for ix = 1,n-1 do voxel_array[ix,*,layer].k[1] = voxel_array[ix-1,*,layer].k[1]
  for iy = 1,n-1 do voxel_array[*,iy,layer].k[0] = voxel_array[*,iy-1,layer].k[0]
  for ix = 1,n-1 do voxel_array[ix,*,layer].c[*,1] = voxel_array[ix-1,*,layer].c[*,1]
  for iy = 1,n-1 do voxel_array[*,iy,layer].c[*,0] = voxel_array[*,iy-1,layer].c[*,0]
endfor
print,'time to load Fourier factors: ',systime(1)-t,' seconds' & wait,.1
print,'loading guidestar shift factors' & wait,.1
t = systime(1)
; load the guidestar shift factors
for layer = 0,n_layers-1 do begin
  for guidestar = 0,n_guidestars-1 do begin
    voxel_array[*,*,layer].s[guidestar] = exp(i*2*!pi*(voxel_array[*,*,layer].k[0]*gs_angle[0,guidestar]*h[layer] + voxel_array[*,*,layer].k[1]*gs_angle[1,guidestar]*h[layer]))
  endfor
endfor
print,'time to load shift factors: ',systime(1)-t,' seconds' & wait,.1
; ------------------------------

;=======================================================================================================
;  test the massively parallel processor (mpp) Fourier transform algorithm by transforming the aperture
;-------------------------------------------------------------------------------------------------------
test_fourier = 1 & fourier_method = 'fft'
test_forwardProp = 1
test_backProp = 1

if (test_fourier) then begin
; for layer = 0,n_layers-1 do voxel_array[*,*,layer].phi = ap ; load aperture into phase
  voxel_array[*,*,n_layers-1].phi = ap
  print,'starting Fourier test' & wait,.01
  
  if ((fourier_method eq 'mpp') or (fourier_method eq 'both')) then begin
  ; forward transform using mpp algorithm
    t = systime(1)
    voxel_array.accum1 = 0
    voxel_array.fphi = 0
    mpp_transform,voxel_array
    voxel_array.fphi /= float(n)^2 ; scaling for the forward transform
    ; for the inverse transform, we don't scale, and the Fourier coefficient must be pre-conjugated
    print,'mpp ft time elapsed: ',systime(1)-t,' seconds'
    ; disp,abs(voxel_array.fphi),'mpp algorithm'
  endif
  
  if ((fourier_method eq 'fft') or (fourier_method eq 'both')) then begin
  ; forward transform using fft algorithm
    t = systime(1)
    for layer = 0,n_layers-1 do voxel_array[*,*,layer].fphi = ft(voxel_array[*,*,layer].phi)
    print,'fft time elapsed: ',systime(1)-t,' seconds'
    ; disp,abs(voxel_array.fphi),'fft algorithm'
  endif
  
endif  ; if (test_fourier)

if (test_forwardProp) then begin
  ; assumes voxels already have Fourier transform coefficients
  print,'starting forward propagation test' & wait,.01 & t = systime(1)
  mpp_forwardPropagate,voxel_array
  print,'mpp propagation time elapsed: ',systime(1)-t,' seconds' & wait,.01
  ; check the answer by inverse transforming the wfs
  for guidestar = 0,n_guidestars-1 do voxel_array[*,*,guidestar].wfs = ft(voxel_array[*,*,guidestar].fwfs,/inverse)
  disp,voxel_array[*,*,0:n_guidestars-1].wfs,'wfs'
endif  ; if (test_forwardProp)

if (test_backProp) then begin
  voxel_array.fphi *= 0
  voxel_array.wfs *= 0
  voxel_array.fwfs *= 0
  for gs = 0,n_guidestars-1 do voxel_array[*,*,gs].wfs = (gs+1)*ap ; put dummy data in the wfs for testing up-propagation
  voxel_array.fwfs = ftcube(voxel_array.wfs)
  print,'starting back propagation test' & wait,.01 & t = systime(1)
  mpp_backPropagate,voxel_array
  print,'mpp propagation time elapsed: ',systime(1)-t,' seconds' & wait,.01
  ; check the answer by inverse transforming the wfs
  voxel_array.phi = ftcube(voxel_array.fphi,/inverse)
  disp,voxel_array.phi,'phi'  
endif

end