;+
; NAME:
;  vt3d_thermwave
; PURPOSE: (one line)
;  Return thermal wave
; DESCRIPTION:
;  Return thermal wave
; CATEGORY:
;  Volatile Transport
; CALLING SEQUENCE:
;  temp = vt3d_thermwave(temp_terms, phase, dtemp_dzeta, z_skin, zeta)
; INPUTS:
;  temp_terms = [temp_0, temp_1, ...], K : complex[n_term+1] or complex[n_loc,n_term+1]
;   s.t., at the surface, temp(phase) = temp_0 + temp_1 exp(i*phase) + temp_2 exp(2*i*phase) + ...
;  phase: float or float[n_t]. freq * time. Radian
;  dtemp_dzeta: float, float[n_loc]
;   - internal heat flux / (sqrt(freq) * therminertia)
;  z_skin: float or float[n_z] or float[n_loc,n_z]. 
;   If z_skin is an array or matrix, then the skin depth varies with depth, and
;   this attempts to modify the wave equation with the WKB approximation.
;   This has no effect if z_skin is a scalar (constant with depth), so
;   it is common to simply set z_skin = 1.
;  zeta: scaled depth, float, float[n_z], float[n_loc, n_z]
; OUTPUTS:
;  temp: float, float[n_t], float[n_z], or float[n_t,n_z]
;        float[n_t,n_loc], 
;
;   temp_terms      phase  substrate  temp  
;   n_term+1        scalar scalar     scaler
;   n_loc,n_term+1  scalar scalar     n_loc
;   n_term+1        n_t    scalar     n_t
;   n_loc,n_term+1  n_t    scalar     n_t,n_loc
;   n_term+1        scalar n_z        n_z
;   n_loc,n_term+1  scalar n_z        n_loc,n_z
;   n_term+1        n_t    n_z        n_t,n_z
;   n_loc,n_term+1  n_t    n_z        n_t,n_loc,n_z
;   n_term+1        scalar n_loc,n_z  n_loc,n_z
;   n_loc,n_term+1  scalar n_loc,n_z  n_loc,n_z
;   n_term+1        n_t    n_loc,n_z  n_t,n_loc,n_z
;   n_loc,n_term+1  n_t    n_loc,n_z  n_t,n_loc,n_z
;   
; RESTRICTIONS:
; PROCEDURE:
;  Young, L. A. 2016, Volatile transport on inhomogeneous surfaces: II. Numerical calculations (VT3D)
;   Resubmitted to Icarus.
;   Eq 3.2-7
;  If skin depth varies with z, then the wave is calculated with a wkb approximation.
; MODIFICATION HISTORY:
;  Written 2011 Dec 30, by Leslie Young, SwRI
;  Modified 2012 Sep 30, LAY.  Added m inside sqrt(_i) term.
;  2016 Mar 22 LAY. Modified for inclusion in vt3d library.
;-
function vt3d_thermwave, temp_terms, phase, dtemp_dzeta, z_skin, zeta

    _i = complex(0,1)

    ; get the dimensions of temp_t, phase, substrate
    temp_0 = extract_last(temp_terms,0,t_dim,n_term) ; scalar or n_loc
    n_term -= 1

    p_dim = size(phase,/n_dim)  ; 0 or 1
    n_t = n_elements(phase) 

    ; each of the extract's are scalar or n_loc
    foo = (temp_0 * $
           extract_last(dtemp_dzeta,0,d_dim) * $
           extract_last(z_skin,0,s_dim,s_n_z) * $
           extract_last(zeta,0,z_dim,z_n_z) )
    n_loc = n_elements(foo)
    n_z = max([s_n_z, z_n_z])

    ;increase dimension of each into local copy
    if t_dim lt 2 then temp_t = replicate(1.,n_loc)##temp_terms else temp_t=temp_terms
    if p_dim lt 1 then p = [phase] else p=phase
    case d_dim of
       0: dt = replicate(dtemp_dzeta, n_loc,n_z)
       1: dt = dtemp_dzeta ## replicate(1., n_loc)
       2: dt = dtemp_dzeta
    endcase
    case s_dim of
       0: s = replicate(z_skin, n_loc,n_z)
       1: s = z_skin ## replicate(1., n_loc) 
       2: s = z_skin
    endcase
    case z_dim of
       0: z = replicate(zeta,n_loc,n_z)
       1: z = zeta ## replicate(1., n_loc)
       2: z = zeta
    endcase

    ; calculate temp at phase = 0
    temp_0 = replicate(1.,n_z) ## extract_last(temp_t,0); n_loc, n_z

    if n_elements(dtemp_dzeta) eq 1 then begin
       temp_f = dt * z
    endif else begin
       temp_f = fltarr(n_loc,n_z)
       for i = 1, n_z-1 do begin
          temp_f[*,i] = temp_f[*,i-1] + (z[*,i]-z[*,i-i]) * mean(dt[*,i-1:i])
       endfor
    endelse

    if n_elements(z_skin) eq 1 then begin
       wkb = replicate(1,n_loc,n_z)
    endif else begin
       s_0 = fltarr(n_loc)
       for i=0,n_loc-1 do s_0[i] = interpol(s[i,*], z[i,*], 0.)
       wkb = sqrt(s/( replicate(1., n_z) ## s_0) )
    endelse
    
    temp = replicate(1,n_t)#complex(reform(temp_0 + temp_f,n_loc*n_z),0)
; help, n_t, n_loc, n_z, temp   
    for m = 1,n_term do begin
       temp_m = extract_last(temp_t,m) # replicate(1.,n_z); n_loc, n_z
       temp += exp(_i * m * p) # reform(temp_m*wkb*exp(sqrt(_i*m)*z), n_loc*n_z)
    endfor

    case 1 of 

       n_t eq 1 and n_loc eq 1 and n_z eq 1: return, re(temp[0])
       n_t eq 1 and n_loc gt 1 and n_z eq 1: return, reform(re(temp),n_loc)
       n_t eq 1 and n_loc eq 1 and n_z gt 1: return, reform(re(temp),n_z)
       n_t eq 1 and n_loc gt 1 and n_z gt 1: return, reform(re(temp),n_loc,n_z)

       n_t gt 1 and n_loc eq 1 and n_z eq 1: return, reform(re(temp),n_t)
       n_t gt 1 and n_loc gt 1 and n_z eq 1: return, reform(re(temp),n_t,n_loc)
       n_t gt 1 and n_loc eq 1 and n_z gt 1: return, reform(re(temp),n_t,n_z)
       n_t gt 1 and n_loc gt 1 and n_z gt 1: return, reform(re(temp),n_t,n_loc,n_z)

    endcase

end

pro _test_vt3d_thermwave

    print, 'test 1 - wave at phase 0'
    zeta = -findgen(40)/4.
    temp_terms = [200., complex(50,0)]
    dtemp_dzeta = 0.
    z_skin = -1.
    phase = 0.
    temp = vt3d_thermwave(temp_terms, phase, dtemp_dzeta, z_skin, zeta)
    plot, temp, zeta
    stop

    print, 'test 2 - wave at phase 0 and !pi'
    zeta = -findgen(40)/4.
    temp_terms = [200., complex(50,0)]
    dtemp_dzeta = 0.
    z_skin = -1.
    phase = 0.
    temp = vt3d_thermwave(temp_terms, phase, dtemp_dzeta, z_skin, zeta)
    plot, temp, zeta, xr=[150,250]
    phase = !pi
    temp = vt3d_thermwave(temp_terms, phase, dtemp_dzeta, z_skin, zeta)
    oplot, temp, zeta, li=2
    stop

    print, 'test 3 - wave at phase 0 and !pi, one call'
    zeta = -findgen(40)/4.
    temp_terms = [200., complex(50,0)]
    dtemp_dzeta = 0.
    z_skin = -1.
    phase = [0.,!pi]
    temp = vt3d_thermwave(temp_terms, phase, dtemp_dzeta, z_skin, zeta)
    plot, temp[0,*], zeta, xr=[150,250]
    oplot, temp[1,*], zeta, li=2
    stop

    print, 'test 4 - surface temp at many phases'
    zeta = 0.
    temp_terms = [200., complex(50,0)]
    dtemp_dzeta = 0.
    z_skin = -1.
    phase = findgen(360)*!pi/180
    temp = vt3d_thermwave(temp_terms, phase, dtemp_dzeta, z_skin, zeta)
    plot, phase, temp

    print, 'test 4 - surface temp at many phases, multiple terms'
    zeta = 0.
    temp_terms = [200., complex(50,0), complex(25,0)]
    dtemp_dzeta = 0.
    z_skin = -1.
    phase = findgen(360)*!pi/180
    temp = vt3d_thermwave(temp_terms, phase, dtemp_dzeta, z_skin, zeta)
    plot, phase, temp

end