module gpt_interface_mod

use bmad_struct
use em_field_mod

implicit none

type gpt_lattice_param_struct
  integer :: fieldmap_dimension = 1    ! Dimensions for field map. 1 or 3
  logical :: only_write_autophase_parameters = .false. ! Option to only write phasing info
end type

contains

!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+ 
! Subroutine write_gpt_lattice_file (gpt_file_unit, lat, gpt_lattice_param, err)
!
! Subroutine to write an gpt lattice file using the information in a lat_struct.
!
! Input:
!   gpt_file_unit -- Integer: unit number to write to
!   lat             -- lat_struct: Holds the lattice information.
!
! Output:
!   err    -- Logical, optional: Set True if, say a file could not be opened.
!-

subroutine write_gpt_lattice_file (gpt_file_unit, lat, gpt_lattice_param,  err)
type (lat_struct), target :: lat
type (ele_struct), pointer :: ele
type (str_index_struct) :: fieldgrid_names
type (gpt_lattice_param_struct) :: gpt_lattice_param

integer :: n, ie, ix_start, ix_end, iu
integer :: gpt_file_unit

character(4000) :: line
character(40) :: name
character(40) :: r_name = 'write_gpt_lattice_file'

logical :: err

!

print *, 'This is '//r_name



ix_start = 1
ix_end = lat%n_ele_max ! Include lords


! Initialize fieldgrid filename list
n = ix_end - ix_start + 1
allocate ( fieldgrid_names%name(n) ) 
allocate ( fieldgrid_names%index(n) )
fieldgrid_names%n_max = 0


! Loop over cavities
write(gpt_file_unit, '(a)') '# GPT lattice'
write(gpt_file_unit, '(2a)') '# ', trim(lat%lattice)

do ie = ix_start, ix_end
  ele => lat%ele(ie)
  if (skip_ele()) cycle
  name = ele%name
  if (ele%lord_status == multipass_lord$) ele => pointer_to_slave (ele, 1)

  ! Clean up symbols in element name
  call str_substitute (name, "#", "_part_")
  call str_substitute (name, "\", "_and_")  
  call str_substitute (name, ".", "_")  
  
  if (gpt_lattice_param%only_write_autophase_parameters) then
    if (ele%key == lcavity$ .or. ele%key == rfcavity$ .or. ele%key == e_gun$)  then
      call write_gpt_ele(gpt_file_unit, ele, name, fieldgrid_names, dimensions = gpt_lattice_param%fieldmap_dimension, only_phasing = .true.)
    endif  
  else
    call write_gpt_ele(gpt_file_unit, ele, name, fieldgrid_names, dimensions = gpt_lattice_param%fieldmap_dimension)
  endif
enddo
write(gpt_file_unit, '(a)') '# END GPT lattice'

contains

!--- Function to skip non-physical elements
function skip_ele() result (skip)
implicit none
logical skip
  ! Skip these elements:
if (ele%slave_status == super_slave$ .or. &
      ele%slave_status == multipass_slave$ .or. &
      ele%key == drift$ .or. &
      ele%key == girder$ .or. &
      ele%key == overlay$ .or. &
      ele%key == group$) then
  skip = .true.
else
  skip = .false.
endif
end function

end subroutine write_gpt_lattice_file

!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+ 
subroutine write_gpt_ele(iu, ele, name, fieldgrid_names, dimensions, only_phasing)
use bmad
type (ele_struct) :: ele
type (floor_position_struct) :: floor1, floor2
type (str_index_struct), optional :: fieldgrid_names
real(rp) :: w_mat(3,3)
real(rp) :: s, x(3), dx(3), d(3), d1(2), d2(2), d3(2), d4(2), w1, e_angle, ds_slice
real(rp) :: gap, strength, x0, y0, z0, x_vec(3), y_vec(3), x_center, y_center, z_center, theta_center
real(rp) :: max_field, freq, phase_lag, ref_time
integer :: i, iu, n_slice, i_slice, q_sign, i_dim
integer, optional :: dimensions
character(40)   :: fieldgrid_output_name, name
character(1) :: component
character(8), parameter :: rfmt = '(es15.7)'
character(40)  :: r_name = 'write_gpt_ele'
logical, optional :: only_phasing

!

i_dim = integer_option(1, dimensions)

q_sign = sign(1,  charge_of(ele%branch%param%particle) ) 

! Get global position and rotation of origin of the field map, at the center of the element
floor1%r = [0.0_rp, 0.0_rp, ele%value(L$)/2]
floor2 = coords_local_curvilinear_to_floor (floor1, ele, in_body_frame = .true., w_mat = w_mat)
x0 = floor2%r(1)
y0 = floor2%r(2)
z0 = floor2%r(3)
! Get x, y unit vectors from w_mat
x_vec = w_mat(:,1)
y_vec = w_mat(:,2)



write (iu, '(2a)')       '#  Element: ', trim(name)
write (iu, '(2a)')       '#  Bmad name: ', trim(ele%name)
write (iu, '(2a)')       '#  Bmad key: ', key_name(ele%key)
write (iu, '(a, i0)')    '#  ix_ele: ', ele%ix_ele
write (iu, '(a, f10.5)')    '#  theta: ', floor2%theta


select case (ele%key)


case (lcavity$, rfcavity$, e_gun$) 
if (.not. present(fieldgrid_names)) call out_io (s_error$, r_name, 'fieldgrid_names required')
!---------------------------------------------------------------------
!Example: 1D cavity
!---------------------------------------------------------------------
!XCTB01 = 0.00;
!YCTB01 = 0.00;
!ZCTB01 = 467456456456;
!Map1D_TM("wcs", XCTB01,YCTB01,ZCTB01,  1,0,0, 0,1,0, "/home/colwyn/ued/fields/eindhoven_1D.gdf", "Z", "Ez", ECTB01, phiCTB01, 2*pi*Master_RF);
!
!arguments:
! (1     )  coordinate system - specify which coordinate system you are going to use to place the element in. If you don't want to make user defined systems, then use "wcs" for world coordinate system [x,y,z]
! (2 -  4)  offset vector - the 3D offset of the origin of the field map [ (0,0,0) in the field map coordinate system ]
! (5 -  7)  x-unit vector in the wcs basis for the element, in above example, there is no rotation so this is (1,0,0).  This doesn't require the vector to be normalized.
! (8 - 10)  y-unit vector in the wcs basis for the element, in above example, there is no rotation so this is (0,1,0).  This doesn't require the vector to be normalized. 
!
!From these two unit vectors, GPT computes the z-unit vector = (x-unit) X (y-unit), which forms the rotation matrix internally used to rotate the fields in the map into the wcs.

!(11) "/home/colwyn/ued/fields/eindhoven_1D.gdf" - field map file 
!(12-13) "Z", "Ez" are the names GPT will look for in the GDF file for the columns specifying the z, Ez data.  You can call your columns whatever you want, but the order you see in the function call is important: it will assign any column named "Z" to the z data, which is the 12th arguement.
!(14) scale factor - ECTB01 is the number that GPT will scale the field by.
!(15) Phase in radians to apply to the map = phi_concrests + phi_relative_offset.
!(16) Angular frequency of the cavity = 2*pi*f.  (f=1.3 GHz)
  
  ! Frequency
  if ( associated(ele%grid_field)) then
    freq = ele%value(rf_frequency$) * ele%grid_field(1)%harmonic
  else
    ! No grid present
    freq = ele%value(rf_frequency$) 
  endif
  
  ! Phase 
  !phase_lag = twopi*(ele%value(phi0$) +  ele%value(phi0_err$))
  !! adjust the lag for 'zero-crossing' 
  !if (ele%key == rfcavity$) phase_lag = twopi*( ele%grid_field%mode(1)%phi0_ref - ele%value(phi0_max$) ) - phase_lag
  !ref_time
  

  if (logic_option(.false., only_phasing)) then
    call gpt_field_grid_scaling(ele, i_dim, max_field, ref_time)
    call write_property('field_scale', max_field, 'Maximum on-axis Ez in V/m')
    call write_property('phase', ref_time*twopi*freq)
    return
  endif


  call get_gpt_fieldgrid_name_and_scaling( ele, fieldgrid_names, fieldgrid_output_name, &
                               max_field, ref_time, dimensions = i_dim)
  call write_property('frequency', freq)
  call write_property('field_scale', max_field, 'Maximum on-axis Ez in V/m')
  call write_property('phase', ref_time*twopi*freq)
  
  select case (dimensions)
  case(1)
    write (iu, '(a)', advance='NO'), 'Map1D_TM('
    call write_geometry()
    write (iu, '(3a)', advance='NO')  '"', trim(fieldgrid_output_name)//'.gdf', '", '
    write (iu, '(1a)', advance='NO') ' "z", "Ez", '
    call write_property('field_scale')
    call write_comma()
    call write_property('phase')
    call write_comma()
    call write_property('frequency*2*pi')
    write (iu, '(a)') ');'  
    
  case(2)
    write (iu, '(a)', advance='NO'), 'Map25D_TM('
    call write_geometry()
    write (iu, '(3a)', advance='NO')  '"', trim(fieldgrid_output_name)//'.gdf', '", '
    write (iu, '(1a)', advance='NO') ' "r", "z", "Er", "Ez", "Bphi", '
    call write_property('field_scale')
    call write_comma()
    write (iu, '(i0)', advance='NO') 0 ! Longitudinal wavenumber k in [m-1]. Must be zero for the simulation of a standing wave cavity.
    call write_comma()
    call write_property('phase')
    call write_comma()
    call write_property('frequency*2*pi')
    write (iu, '(a)') ');'      
  
  case(3)
    ! Note: GPT fields oscillate as exp(+i wt)
    do i=1,2
    if (i==1) then
      component = 'E'
    else
      if (ele%key==e_gun$ .and. freq==0) exit ! No H field for dc guns!
      component = 'H'
    endif
    
    write (iu, '(a)', advance='NO'), 'Map3D_'//component//'complex('
    call write_geometry()
    write (iu, '(3a)', advance='NO')  '"', trim(fieldgrid_output_name)//'_'//component//'.gdf', '", '
    if (i==1) then
      write (iu, '(1a)', advance='NO')  '"x", "y", "z", "ExRe", "EyRe", "EzRe", "ExIm", "EyIm", "EzIm", '
    else 
      write (iu, '(1a)', advance='NO')  '"x", "y", "z", "HxRe", "HyRe", "HzRe", "HxIm", "HyIm", "HzIm", '
    endif
    
    call write_property('field_scale')
    call write_comma()
    call write_property('phase')
    call write_comma()
    call write_property('frequency*2*pi')
    write (iu, '(a)') ');'    
    enddo
    
  end select

case (solenoid$) 
if (.not. present(fieldgrid_names)) call out_io (s_error$, r_name, 'fieldgrid_names required')
call get_gpt_fieldgrid_name_and_scaling( ele, fieldgrid_names, fieldgrid_output_name, &
                               max_field, ref_time, dimensions = i_dim)
  
  call write_property('field_scale', max_field, 'signed abs max on-axis Bz in T')
  select case (dimensions)
  case(1)
    write (iu, '(a)', advance='NO'), 'Map1D_B('
    call write_geometry()
    write (iu, '(3a)', advance='NO')  '"', trim(fieldgrid_output_name)//'.gdf', '", '
    write (iu, '(1a)', advance='NO') '"z", "Bz", '
    call write_property('field_scale')
    write (iu, '(a)') ');'    

  case(2)
    write (iu, '(a)', advance='NO'), 'Map2D_B('
    call write_geometry()
    write (iu, '(3a)', advance='NO')  '"', trim(fieldgrid_output_name)//'.gdf', '", '
    write (iu, '(1a)', advance='NO') ' "r", "z", "Br", "Bz", '
    call write_property('field_scale')
    write (iu, '(a)') ');'    

  case(3)
    write (iu, '(a)', advance='NO'), 'Map3D_B('
    call write_geometry()
    write (iu, '(3a)', advance='NO')  '"', trim(fieldgrid_output_name)//'_B.gdf', '", '
    write (iu, '(1a)', advance='NO') ' "x", "y", "z", "Bx", "By", "Bz", '
    call write_property('field_scale')
    write (iu, '(a)') ');'    
  
  end select

case (sbend$) 
!!! TODO
if (.not. present(fieldgrid_names)) call out_io (s_error$, r_name, 'fieldgrid_names required')
   !!!call get_gpt_fieldgrid_name_and_scaling( ele, fieldgrid_names, fieldgrid_output_name, &
   !!!                            max_field, ref_time, dimensions = i_dim)
  write(*, *) '#TEMP: NO FIELD MAP, GEOMETRY: '
  write(iu, '(a)', advance = 'NO') '#GEOMETRY: '
  call write_geometry()
  !call write_property('field_scale', max_field, 'signed abs max on-axis By in T')
  !select case (dimensions)
  !case(3)
  !  write (iu, '(a)', advance='NO'), 'Map3D_B('
  !  call write_geometry()
  !  write (iu, '(3a)', advance='NO')  '"', trim(fieldgrid_output_name)//'_B.gdf', '", '
  !  write (iu, '(1a)', advance='NO') ' "x", "y", "z", "Bx", "By", "Bz", '
  !  call write_property('field_scale')
  !  write (iu, '(a)') ');'    
  

case (quadrupole$)   
   call write_property('gradient', ele%value(b1_Gradient$), 'T/m')
   call write_property('L', ele%value(L$), 'm') 
   write (iu, '(a)', advance='NO'), 'quadrupole('
   call write_geometry()
   call write_property('L')
   call write_comma()
   call write_property('gradient')
   write (iu, '(a)') ' );'  

end select

  
!case default
!  print *, 'ele not coded: ', ele%name
!end select

! Footer
write (iu, '(a)') '' 

contains

subroutine write_geometry()
write (iu, '(a)', advance = 'NO') '"wcs", '
call write_real(x0); call write_comma()
call write_real(y0); call write_comma()
call write_real(z0); call write_comma()
call write_real(w_mat(1,1)); call write_comma()
call write_real(w_mat(2,1)); call write_comma()
call write_real(w_mat(3,1)); call write_comma()
call write_real(w_mat(1,2)); call write_comma()
call write_real(w_mat(2,2)); call write_comma()
call write_real(w_mat(3,2)); call write_comma()
end subroutine


subroutine write_property(property, value, comment)
character(*) :: property
real(rp), optional :: value
character(*), optional :: comment
if (present(value)) then
  if (value == nint(value)) then
    write (iu, '(a, i0, a)', advance = 'NO') trim(name)//'_'//trim(property)// ' = ', nint(value), ';'   
  else
    write (iu, '(a, '//rfmt//', a)', advance = 'NO') trim(name)//'_'//trim(property)// ' = ', value, ';'
  endif
  
  if (present(comment)) then
    write (iu, '(a)') ' # '//trim(comment)
  else
    write (iu, '(a)') ''
  endif
else
  write (iu, '(a)', advance = 'NO') trim(name)//'_'//trim(property)
endif
end subroutine


subroutine write_comma()
write (iu, '(a)', advance='NO') ', '
end subroutine

subroutine write_real(x)
real(rp) :: x
if (x == nint(x)) then
  write (iu, '(i0 )', advance = 'NO') nint(x)
else
  write (iu, rfmt, advance = 'NO') x
endif
end subroutine

end subroutine


!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+ 
! Subroutine  get_gpt_fieldgrid_name_and_scaling(ele, name_indexx, output_name, 
!                                    field_scale, ref_time,  dimensions)
!
! Subroutine to get a field grid filename and its scaling. Calls write_gpt_field_grid_file.
!   If the field grid file does not exist, it is written.
!
!   Note: This is very similar to get_opal_fieldgrid_name_and_scaling
!
!
! Input:
!   ele              -- ele_struct: element to make map
!   name_indexx      -- str_index_struct: contains field grid filenames
!   dimensions       -- integer, optional: 1, 2, or 3 dimensions. Default: 1
!
! Output:   
!   name_indexx      -- str_index_struct: updated if new name is added
!   output_name      -- real(rp): output filename. 
!   field_scale      -- real(rp): the scaling of the field grid
!   ref_time         -- real(rp): time that the field was evaluated at
!-

subroutine get_gpt_fieldgrid_name_and_scaling(&
             ele, name_indexx, output_name, field_scale, ref_time, dimensions)
type (ele_struct) :: ele
type (str_index_struct) :: name_indexx
character(*)  :: output_name
real(rp)      :: field_scale
real(rp)      :: ref_time
character(200) :: unique_grid_file

integer :: ix_match, iu_fieldgrid, ios, i_dim
integer, optional :: dimensions
character(40), parameter:: r_name = 'get_gpt_fieldgrid_name_and_scaling'
character(10) :: suffix = '_ASCII.gpt'

!

i_dim = integer_option(1, dimensions)



output_name = ''

! Check for field grid. Force sbends to have unique fieldmaps
! because their geometry can change how the field is output.
if (associated (ele%grid_field) .and. ele%key /= sbend$) then
  unique_grid_file = ele%grid_field(1)%ptr%file
else
  ! Just us the element's index as an identifier
  write(unique_grid_file, '(a,i0)') 'ele_', ele%ix_ele
endif

! Check field map file. If file has not been written, create a new file. 
call find_index (unique_grid_file, name_indexx, ix_match)
! Check for match with existing grid
if (ix_match > 0) then
  ! File should exist  
  call get_output_name()
  call gpt_field_grid_scaling(ele, i_dim, field_scale, ref_time)
else
  ! File does not exist.
  ! Add name to list  
  call find_index (unique_grid_file, name_indexx, ix_match, add_to_list = .true.)
  ix_match = name_indexx%n_max
  call get_output_name()

  
  select case(i_dim)
  case(1)
    iu_fieldgrid = lunget()
    open(iu_fieldgrid, file = trim(output_name)//trim(suffix), iostat = ios)
    call write_gpt_field_grid_file_1D (iu_fieldgrid , ele, field_scale, ref_time)
    close(iu_fieldgrid)
  
  case(2)
    iu_fieldgrid = lunget()
    open(iu_fieldgrid, file = trim(output_name)//trim(suffix), iostat = ios)
    call write_gpt_field_grid_file_2D (iu_fieldgrid, ele, field_scale, ref_time)
    close(iu_fieldgrid)
  
  case(3)
    ! This will write several files
    call write_gpt_field_grid_file_3D (output_name, ele, field_scale, ref_time)
  
  case default
    call out_io (s_error$, r_name, 'Bad dimension request for field map ')
    if (global_com%exit_on_error) call err_exit
  end select  
  
end if

contains 





subroutine get_output_name()
    ! Names will be suffixed elsewhere
    write(output_name, '(i1, a, i0)') i_dim, 'D_fieldmap_', ix_match
end subroutine

end subroutine get_gpt_fieldgrid_name_and_scaling


!---
! Get field_scale and ref_time depending on the type of grid dimension used
!---
subroutine gpt_field_grid_scaling(ele, i_dim, field_scale, ref_time)
type(ele_struct) :: ele
integer :: i_dim
real(rp) :: field_scale, ref_time
character(40) :: r_name = 'gpt_field_grid_scaling'
!
! Call with iu=0 or '' to get field_scale and ref_time without writing a file
select case(i_dim)
  case(1)
    call write_gpt_field_grid_file_1D (0, ele, field_scale, ref_time)
  case(2)
    call write_gpt_field_grid_file_2D (0, ele, field_scale, ref_time)
  case(3)
    call write_gpt_field_grid_file_3D ('', ele, field_scale, ref_time)
  case default
    call out_io (s_error$, r_name, 'Bad dimension request for field map ')
    if (global_com%exit_on_error) call err_exit
end select
end subroutine



!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+ 
! Subroutine write_gpt_field_grid_file_1D (gpt_file_unit, ele, maxfield, ref_time, dz, err)
!
!   Write 1-D field map files for gpt. The format is:
!   z field
!   ...
!  
!   Note: Simplified from write_opal_field_grid_file
!
! Input:
!   gpt_file_unit -- Integer: unit number to write to, if > 0
!                        if < 0, nothing is written, and only maxfield is returned
!   ele            -- ele_struct: element to make map
!   dz             -- real(rp), optional: z step size in m. Default: 0.001 m
!
!
! Output:          
!   maxfield       -- Real(rp): absolute maximum found for element field scaling
!   ref_time       -- real(rp): time that the field was evaluated at
!   err            -- Logical, optional: Set True if, say a file could not be opened.
!-



subroutine write_gpt_field_grid_file_1D (gpt_file_unit, ele, maxfield, ref_time, dz, err)
integer      :: gpt_file_unit
type (ele_struct) :: ele
type (lat_param_struct) :: param
real(rp)        :: maxfield
logical, optional :: err

character(40)  :: r_name = 'write_gpt_field_grid_file'
character(10)   ::  rfmt 

type (coord_struct) :: orb
type(em_field_struct) :: field_re, field_im
type (grid_field_pt1_struct), allocatable :: pt(:)
type (grid_field_pt1_struct) :: ref_field
real(rp) :: z_step, z_min, z_max, z0
real(rp) :: freq,  z, phase_ref
real(rp) :: gap, edge_range
real(rp) :: ref_time
real(rp), optional :: dz
complex ::  phasor_rotation

integer :: nx, nz, iz, ix

real(rp) :: Ex_factor, Ez_factor, Bx_factor, By_factor, Bz_factor

logical loc_ref_frame

!
if (present(err)) err = .true.


loc_ref_frame = .true. 

param = ele%branch%param

! Format for numbers
rfmt = 'es13.5'

! Output will be relative to z0
z0 = ele%value(L$)/2

! Step size
z_step = real_option(0.001_rp, dz)

z_min = 0.0_rp
z_max = ele%value(L$)
 
nz = ceiling(z_max/z_step)

! These are 1D fields, so the field will be probed on-axis
orb%vec(1) = 0
orb%vec(3) = 0

Ez_factor = 1
Bz_factor = 1

! Reference time, will be adjusted by oscillating elements below  
ref_time = 0  
  
select case (ele%key)

!-----------
! LCavity, RFCavity, E_GUN
!-----------
case (lcavity$, rfcavity$, e_gun$) 
  if ( associated(ele%grid_field)) then
    freq = ele%value(rf_frequency$) * ele%grid_field(1)%harmonic
  else
    ! No grid present
    freq = ele%value(rf_frequency$) 
  endif
  ! if (freq .eq. 0) freq = 1e-30_rp ! To prevent divide by zero

  ! Allocate temporary pt array
  allocate (pt(0:nz))
  ! Write data points
  
  ! initialize maximum found field
  maxfield = 0
  do iz = 0, nz
    z = z_step * iz 
    ! Calculate field at \omegat*t=0 and \omega*t = \pi/2 to get real and imaginary parts
    call em_field_calc (ele, param, z, orb, loc_ref_frame, field_re, rf_time = 0.0_rp)
    ! if frequency is zero, zero out field_im
    if(freq == 0) then
      field_im%E=0
      field_im%B=0
    else 
      call em_field_calc (ele, param, z, orb, loc_ref_frame, field_im, rf_time = 0.25/freq)
    endif
    pt(iz)%E(:) = cmplx(field_re%E(:), field_im%E(:))
    pt(iz)%B(:) = cmplx(field_re%B(:), field_im%B(:))
    ! Update ref_field if larger Ez is found
    if(abs(pt(iz)%E(3)) > maxfield) then
      ref_field = pt(iz)
      maxfield = abs(ref_field%E(3))
    end if 
  end do
  
  ! Reference time
  phase_ref = atan2( aimag(ref_field%E(3) ), real(ref_field%E(3) ) )
  if (freq /= 0 ) ref_time = -phase_ref /(twopi*freq)
  
  ! Write to file
  if (gpt_file_unit > 0 )  then
    ! Normalize
    if (maxfield > 0) Ez_factor = 1/maxfield

    ! Calculate complex rotation number to rotate Ez onto the real axis    
    phasor_rotation = cmplx(cos(phase_ref), -sin(phase_ref) )
    
    write (gpt_file_unit, '(2a13)') 'z', 'Ez'
    do iz = 0, nz   
      write (gpt_file_unit, '(2'//rfmt//')') z_step * iz - z0, Ez_factor * real ( pt(iz)%E(3) * phasor_rotation )
    enddo 
  end if
  deallocate(pt)

  !-----------
  ! Solenoid
  !-----------
  ! Note: This is similar to code for lcavity/rfcavity
  case (solenoid$) 
                                         
  ! Allocate temporary pt array
  allocate (pt(0:nz))
  ! Write data points
  
  ! initialize maximum found field
  maxfield = 0
  
  do iz = 0, nz
    z = z_step * iz 
    call em_field_calc (ele, param, z, orb, loc_ref_frame, field_re, rf_time = 0.0_rp)
    field_im%E = 0
    field_im%B = 0
    pt(iz)%E(:) = cmplx(field_re%E(:), field_im%E(:))
    pt(iz)%B(:) = cmplx(field_re%B(:), field_im%B(:))

    ! Update ref_field if larger Bz is found
    if (abs(pt(iz)%B(3)) > maxfield) then
         ref_field = pt(iz)
         maxfield = abs(ref_field%B(3))
    end if 
  end do
  
  ! Restore the sign
  maxfield = ref_field%B(3)
  
  ! Write to file
  if (gpt_file_unit > 0 )  then
    write (gpt_file_unit, '(2a13)') 'z', 'Bz'

    ! Scaling
    if (maxfield /= 0) Bz_factor = 1/maxfield
    
    do iz = 0, nz
      write (gpt_file_unit, '(2'//rfmt//')') z_step * iz - z0, Bz_factor * real ( pt(iz)%B(3))          
    enddo
  end if

  ! cleanup 
  deallocate(pt)


  !-----------
  ! Default (gives an error)
  !-----------
  case default
  call out_io (s_error$, r_name, 'MISSING gpt FIELD GRID CODE FOR ELEMENT TYPE: ' // key_name(ele%key), &
             '----')
  if (global_com%exit_on_error) call err_exit
  
end select 


if (maxfield == 0) then
  call out_io (s_error$, r_name, 'ZERO MAXIMUM FIELD IN ELEMENT: ' // key_name(ele%key), &
             '----')
  if (global_com%exit_on_error) call err_exit
end if



end subroutine write_gpt_field_grid_file_1D



!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+ 
! Subroutine write_gpt_field_grid_file_2D (gpt_file_unit, ele, maxfield, ref_time, 
!                                                              dr, dz,  err)
!
! Subroutine to write an GPT lattice file using the information in
! a lat_struct. Optionally only part of the lattice can be generated.
!
!
! Input:
!   gpt_file_unit -- Integer: unit number to write to, if > 0
!                        if < 0, nothing is written, and only maxfield is returned
!   ele            -- ele_struct: element to make map
!   dr             -- real(rp), optional: r step size in m. Default: 0.001 m
!   dz             -- real(rp), optional: z step size in m. Default: 0.001 m
!   r_max          -- real(rp), optional: maximum radius in m. Default: 0.02 m
!
! Output:          
!   maxfield       -- Real(rp): absolute maximum found for element field scaling
!   ref_time       -- real(rp): time that the field was evaluated at
!   err            -- Logical, optional: Set True if, say a file could not be opened.
!-

subroutine write_gpt_field_grid_file_2D (gpt_file_unit, ele, maxfield, ref_time, &
                                         dr, dz, r_max,  err)

integer      :: gpt_file_unit
integer         :: dimensions
type (ele_struct) :: ele
type (lat_param_struct) :: param
real(rp)      :: maxfield, ref_time
logical, optional :: err

character(40)  :: r_name = 'write_gpt_field_grid_file_2D '
character(10), parameter ::  rfmt  =  'es13.5'


type (coord_struct) :: orb
type(em_field_struct) :: field_re, field_im
type (grid_field_pt1_struct), allocatable :: pt(:,:,:)
type (grid_field_pt1_struct) :: ref_field
real(rp), optional :: dr, dz, r_max
real(rp) :: x_step, z_step, x_min, x_max, z_min, z_max, z0
real(rp) :: freq, x, z, phase_ref
real(rp) :: gap, edge_range
complex ::  phasor_rotation

integer :: nx, nz, iz, ix

real(rp) :: Ex_factor, Ez_factor, Bx_factor, By_factor, Bz_factor

logical loc_ref_frame

!
if (present(err)) err = .true.

loc_ref_frame = .true. 

param = ele%branch%param

! Output will be relative to z0
z0 = ele%value(L$)/2
  
! Step size
x_step = real_option(0.001_rp, dr)
z_step = real_option(0.001_rp, dz)

z_min = 0.0_rp
z_max = ele%value(L$)
 
nz = ceiling(z_max/z_step)

x_min = 0.0_rp
x_max = real_option(0.02_rp, r_max)

z_min = 0.0_rp
z_max = ele%value(L$)

nx = ceiling(x_max/x_step)  
nz = ceiling(z_max/z_step)


! Reference time, will be adjusted by oscillating elements below  
ref_time = 0  
  
select case (ele%key)

!-----------
! LCavity, RFCavity, E_GUN
!-----------
case (lcavity$, rfcavity$, e_gun$) 

  freq = ele%value(rf_frequency$) * ele%grid_field(1)%harmonic
  ! if (freq .eq. 0) freq = 1e-30_rp ! To prevent divide by zero

  ! Example:
  ! r  z  Er  Ez  Bphi
  ! ...

  ! Allocate temporary pt array
  allocate ( pt(0:nx, 0:nz, 1:1) )
  ! Write data points
  
  ! initialize maximum found field
  maxfield = 0
  
  do ix = 0, nx
    do iz = 0, nz
      x = x_step * ix
      z = z_step * iz 
      orb%vec(1) = x
      orb%vec(3) = 0.0_rp
      
      ! Calculate field at \omegat*t=0 and \omega*t = \pi/2 to get real and imaginary parts
      call em_field_calc (ele, param, z, orb, loc_ref_frame, field_re, rf_time = 0.0_rp)
      ! if frequency is zero, zero out field_im
      if(freq == 0) then
        field_im%E=0
        field_im%B=0
      else 
        call em_field_calc (ele, param, z, orb, loc_ref_frame, field_im, rf_time = 0.25/freq)
      endif

      pt(ix, iz, 1)%E(:) = cmplx(field_re%E(:), field_im%E(:))
      pt(ix, iz, 1)%B(:) = cmplx(field_re%B(:), field_im%B(:))
      
      ! Update ref_field if larger Ez is found
      if(ix == 0 .and. abs(pt(ix, iz, 1)%E(3)) > maxfield) then
         ref_field = pt(ix, iz, 1)
         maxfield = abs(ref_field%E(3))
      end if 
    end do
  end do
  
  ! Reference time
  phase_ref = atan2( aimag(ref_field%E(3) ), real(ref_field%E(3) ) )
  if (freq /= 0 ) ref_time = -phase_ref /(twopi*freq)
  
  ! Write to file
  if (gpt_file_unit > 0 )  then
    write (gpt_file_unit, '(5a13)') 'r', 'z', 'Er', 'Ez', 'Bphi'
    ! Scaling
    Ex_factor =  (1/maxfield)
    Ez_factor =  (1/maxfield)
    By_factor = -(1/maxfield)
  
    ! Calculate complex rotation number to rotate Ez onto the real axis
    phasor_rotation = cmplx(cos(phase_ref), -sin(phase_ref))
  
    do ix = 0, nx
      do iz = 0, nz
        write (gpt_file_unit, '(5'//rfmt//')')  x_step * ix, z_step * iz - z0 , &
        										Ex_factor * real ( pt(ix, iz, 1)%E(1) * phasor_rotation ), &
                                                Ez_factor * real ( pt(ix, iz, 1)%E(3) * phasor_rotation ), &
                                                By_factor * aimag (  pt(ix, iz, 1)%B(2)*phasor_rotation )
      enddo
    enddo
  
  end if
   
   
   
   
   deallocate(pt)

  !-----------
  ! Solenoid
  !-----------
  ! Note: This is similar to code for lcavity/rfcavity
  case (solenoid$) 
                                         
  ! Example:
  ! r z Br Bz
  
  
  ! Allocate temporary pt array
  allocate ( pt(0:nx, 0:nz, 1:1) )
  ! Write data points
  
  ! initialize maximum found field
  maxfield = 0
  
  do ix = 0, nx
    do iz = 0, nz
      x = x_step * ix
      z = z_step * iz 
      orb%vec(1) = x
      orb%vec(3) = 0.0_rp

      call em_field_calc (ele, param, z, orb, loc_ref_frame, field_re, rf_time = 0.0_rp)
      field_im%E = 0
      field_im%B = 0

      pt(ix, iz, 1)%E(:) = cmplx(field_re%E(:), field_im%E(:))
      pt(ix, iz, 1)%B(:) = cmplx(field_re%B(:), field_im%B(:))
      
      ! Update ref_field if larger Bz is found
      if(ix==0 .and. abs(pt(ix, iz, 1)%B(3)) > maxfield) then
         ref_field = pt(ix, iz, 1)
         maxfield = abs(ref_field%B(3))
      end if 


    end do
  end do
  
  ! Restore the sign
  maxfield = ref_field%B(3)
  
  ! Write to file
  if (gpt_file_unit > 0 )  then

    ! Write header
    write (gpt_file_unit, '(4a13)') 'r', 'z', 'Br', 'Bz'

   ! Scaling
   Bx_factor = 1/maxfield
   Bz_factor = 1/maxfield
    
    ! XZ ordering: ix changes fastest (inner loop)
    do ix = 0, nx
      do iz = 0, nz
        write (gpt_file_unit, '(4'//rfmt//')') ,  x_step * ix, z_step * iz - z0, &
          			 Bx_factor * real (  pt(ix, iz, 1)%B(1) ), &
         			 Bz_factor * real (  pt(ix, iz, 1)%B(3) )            
      enddo
    enddo
  
  end if

  ! cleanup 
   deallocate(pt)

  !-----------
  ! Default (gives an error)
  !-----------
  case default
  call out_io (s_error$, r_name, 'MISSING GPT FIELD GRID CODE FOR: ' // key_name(ele%key), &
             '----')
  if (global_com%exit_on_error) call err_exit
  
  
  
end select 

if (maxfield == 0) then
  call out_io (s_error$, r_name, 'ZERO MAXIMUM FIELD IN ELEMENT: ' // key_name(ele%key), &
             '----')
  if (global_com%exit_on_error) call err_exit
end if



end subroutine write_gpt_field_grid_file_2D


!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+ 
! Subroutine write_gpt_field_grid_file_3D (base_filename, ele, maxfield, ref_time,
!                                                                      dz, err)
!
!   Writes 3-D field map files for gpt. The format is:
!
!   E-fields:
!   'x', 'y', 'z', 'ExRe', 'EyRe', 'EzRe', 'ExIm ', 'EyIm ', 'EzIm '
!   H-fields
!   'x', 'y', 'z', 'HxRe', 'HyRe', 'HzRe', 'HxIm ', 'HyIm ', 'HzIm '
!
!   where the fields oscillate as exp(+i \omega t)
!
!   Note: similar to write_gpt_field_grid_file
!
! Input:
!   base_filename  -- character(*): Base filename. Files will be written as:
!                                   base_filename_E_ASCII.gpt, _H_ASCII.gpt
!                                   If set to '', no files will be written
!   ele            -- ele_struct: element to make map
!   dz             -- real(rp), optional: z step size in m. Default: 0.001 m
!
! Output:          
!   maxfield       -- Real(rp): absolute maximum on-axis field found for element field scaling
!   ref_time       -- real(rp): time that the field was evaluated at
!   err            -- Logical, optional: Set True if, say a file could not be opened.
!-

subroutine write_gpt_field_grid_file_3D (base_filename, ele, maxfield, ref_time, &
                                         dz, err)

implicit none


character(*) :: base_filename
integer      :: iu
type (ele_struct) :: ele
type (lat_param_struct) :: param

logical, optional :: err

character(40)  :: r_name = 'write_gpt_field_grid_file'
character(10)   ::  rfmt 

type (coord_struct) :: orb
type(em_field_struct) :: field_re, field_im
type (grid_field_pt1_struct), allocatable :: pt(:,:,:)
type (grid_field_pt1_struct) :: ref_field
real(rp)        :: maxfield, ref_time
real(rp) :: x_step, x_min, x_max
real(rp) :: y_step, y_min, y_max
real(rp) :: z_step, z_min, z_max, z0
real(rp) :: freq, x, y, z, phase_ref
real(rp) :: gap, edge_range
real(rp), optional :: dz
complex ::  phasor_rotation, test_field_value

integer :: nx, ny, nz, iz, ix, iy, ifield, ix_center, iy_center

logical loc_ref_frame, write_file

!

if (present(err)) err = .true.

if (base_filename == '') then
  write_file = .false.
else
  write_file = .true.
endif

loc_ref_frame = .true. 

param = ele%branch%param

! Format for numbers
rfmt = 'es13.5'

! Output will be relative to z0
z0 = ele%value(L$)/2

z_step = real_option(0.001_rp, dz)

z_min = 0.0_rp
z_max = ele%value(L$)

! Special bounds for curvilinear coordinates
if (ele%key == sbend$) then
  if (ele%value(g$) /= 0) then
    z_max = ele%value(rho$)*sin(ele%value(angle$)/2)
    z_min = -z_max
  endif
endif


x_max =  .02_rp
x_step = z_step  ! Same as z TODO: generalize
x_min = -x_max
ix_center = nint(x_max/x_step) + 1
 
y_max =  .02_rp
y_step = z_step  ! Same as z 
y_min = -x_max
iy_center = nint(y_max/y_step) + 1
 
nz = ceiling((z_max-z_min)/z_step) + 1
nx = ceiling((x_max-x_min)/x_step) + 1
ny = ceiling((y_max-y_min)/y_step) + 1

! Reference time, will be adjusted by oscillating elements below  
ref_time = 0 

!-----------
! LCavity, RFCavity, E_GUN
!-----------

! Get frequency
select case (ele%key)
case (lcavity$, rfcavity$, e_gun$) 
  if ( associated(ele%grid_field)) then
    freq = ele%value(rf_frequency$) * ele%grid_field(1)%harmonic
  else
    ! No grid present
    freq = ele%value(rf_frequency$) 
  endif
case default
  freq = 0
end select

! Allocate all sample points.
allocate (pt(1:nx, 1:ny, 1:nz))

maxfield = 0

! 3D loop to load pt array, and to get the maximum field
do iz=1, nz
do iy=1, ny
! Get field along x line
do ix=1, nx
  orb%vec(1) = x_min + (ix-1)*x_step
  orb%vec(3) = y_min + (iy-1)*y_step
  z = z_min + (iz-1)*z_step
  if (ele%key == sbend$) then
    call sbend_field_at(orb%vec(1), z, field_re)
  else
    ! Calculate field at \omegat*t=0 and \omega*t = \pi/2 to get real and imaginary parts
    call em_field_calc (ele, param, z, orb, loc_ref_frame, field_re, rf_time = 0.0_rp)
  endif
  ! if frequency is zero, zero out field_im
  if(freq == 0) then
    field_im%E=0
    field_im%B=0
  else 
    call em_field_calc (ele, param, z, orb, loc_ref_frame, field_im, rf_time = 0.25/freq)
  endif
  pt(ix, iy, iz)%E(:) = cmplx(field_re%E(:), field_im%E(:))
  pt(ix, iy, iz)%B(:) = cmplx(field_re%B(:), field_im%B(:))

enddo
enddo
enddo
  
  
! Get reference field: The largest on-axis field
do iz = 1, nz
  test_field_value = gpt_max_field_reference(pt(ix_center, iy_center, iz), ele)
  if(abs(test_field_value) > maxfield) then
    ref_field = pt(ix_center, iy_center, iz)
    maxfield = abs(test_field_value)
  end if 
enddo

if (freq == 0) then
  ! Fields should be purely real
  ! restore the sign
  maxfield = sign(maxfield, real(gpt_max_field_reference(ref_field, ele), rp))
  phasor_rotation = cmplx(1.0_rp, 0.0_rp )
else
! Calculate complex rotation number to rotate Ez onto the real axis
  phase_ref = atan2( aimag(ref_field%E(3) ), real(ref_field%E(3) ) )
  phasor_rotation = cmplx(cos(phase_ref), -sin(phase_ref) )
  ref_time = -phase_ref /(twopi*freq)
endif

! Put scaling in phasor rotation
if (maxfield /= 0) phasor_rotation = (1/maxfield)*phasor_rotation


! Write to file according to element

if (write_file) then
  select case(ele%key)
  
  case(e_gun$)
    call write_pts('E')
    if (freq /= 0) call write_pts('H')
  
  case(lcavity$, rfcavity$)
    call write_pts('E')
    call write_pts('H')
    
  case(sbend$)
    ! Curvilinear, remove z0 offset, because z_max and z_min already include it
    z0 = 0
    call write_pts('B')
    
  case(solenoid$)  
    call write_pts('B')
    
  case(quadrupole$)  
    call write_pts('B')    
    
  case default
    print *, 'not coded'
  end select
endif

! Cleanup
deallocate(pt)

contains

!--- write to file: according to component (1 for E, 2 for H)
subroutine write_pts(component)
implicit none
character(40) :: output_name
character(1) :: component
integer :: iu, ios

!


write(output_name, '(a)') trim(base_filename)//'_'//component


iu =lunget()
open(iu, file = trim(output_name)//'_ASCII.gpt', iostat = ios)

! Note: GPT fields oscillate as exp(+i wt), so the imaginary parts need a minus sign.

select case(component)
case('E')
! E field
  
  write(iu, '(9a13)') 'x', 'y', 'z', 'ExRe', 'EyRe', 'EzRe', 'ExIm ', 'EyIm ', 'EzIm '
  do iz=1, nz
  do iy=1, ny
  do ix=1, nx-1
    write(iu, '(9'//rfmt//')') x_min + (ix-1)*x_step, y_min + (iy-1)*y_step, z_min + (iz-1)*z_step - z0, &
                               real( pt(ix, iy, iz)%E(:) * phasor_rotation), &
                              -aimag( pt(ix, iy, iz)%E(:) * phasor_rotation)                        
  enddo
  enddo
  enddo
  
case('H')
  ! H field
  write(iu, '(9a13)') 'x', 'y', 'z', 'HxRe', 'HyRe', 'HzRe', 'HxIm ', 'HyIm ', 'HzIm '
  do iz=1, nz
  do iy=1, ny
  do ix=1, nx-1
    write(iu, '(9'//rfmt//')') x_min + (ix-1)*x_step, y_min + (iy-1)*y_step, z_min + (iz-1)*z_step - z0, &
                               real( pt(ix, iy, iz)%B(:) * phasor_rotation)/mu_0_vac, &
                               -aimag( pt(ix, iy, iz)%B(:) * phasor_rotation)/mu_0_vac         
  enddo
  enddo
  enddo
  
case('B')
  ! B field
  write(iu, '(6a13)') 'x', 'y', 'z', 'Bx', 'By', 'Bz'
  do iz=1, nz
  do iy=1, ny
  do ix=1, nx-1
    write(iu, '(6'//rfmt//')') x_min + (ix-1)*x_step, y_min + (iy-1)*y_step, z_min + (iz-1)*z_step - z0, &
                               real( pt(ix, iy, iz)%B(:) * phasor_rotation)     
  enddo
  enddo
  enddo
  
end select

close(iu)

end subroutine



subroutine sbend_field_at(x, z, field)
type(em_field_struct) :: field
real(rp) :: x, z, s

!s is relative to the origin of the Cartesian system. Add z0 to get s relative
!to the element
call convert_local_cartesian_to_local_curvilinear(x, z, ele%value(g$),  orb%vec(1), s)

if (s+z0>ele%value(L$) .or. s+z0 < 0) then
  field%e = 0
  field%b = 0
  return
endif  
  
call em_field_calc (ele, param, s + z0, orb, loc_ref_frame, field, rf_time = 0.0_rp)

call rotate_field_zx(field, -s*ele%value(g$))
end subroutine


end subroutine




!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+ 
!  Returns the relevant field value for the field scaling calculation
!-
function gpt_max_field_reference(pt0, ele) result(field_value)
implicit none
type (grid_field_pt1_struct) pt0 
type (ele_struct) :: ele
real(rp) :: field_value
character(40), parameter :: r_name = 'gpt_max_field_reference'
!
select case (ele%key)
case(e_gun$, lcavity$, rfcavity$)
  ! Ez
  field_value = pt0%E(3) ! Note: pt0 is complex
case(sbend$)
  ! By
  field_value = pt0%B(2)
case(solenoid$)
  ! Bz
  field_value = pt0%B(3)
  
case default
  call out_io(s_error$, r_name, 'Max field reference not specified for element type '// key_name(ele%key))

end select

end function


subroutine rotate_field_zx(field, theta)
type (em_field_struct) :: field
real(rp) :: theta, temp, ca, sa
ca = cos(theta)
sa = sin(theta)

temp       = field%E(3)*ca - field%E(1)*sa
field%E(1) = field%E(3)*sa + field%E(1)*ca
field%E(3) = temp

temp       = field%B(3)*ca - field%B(1)*sa
field%B(1) = field%B(3)*sa + field%B(1)*ca
field%B(3) = temp
end subroutine

subroutine convert_local_curvilinear_to_local_cartesian(x, s, xout, zout)
real(rp) :: x, s, g, rho, xout, zout, theta
if (g == 0) then
  xout = x
  zout = s
  return
endif
rho = 1/g
theta = s/rho
xout = (x + rho)*cos(theta) + rho
zout = (x+rho)*sin(theta)
end subroutine

subroutine convert_local_cartesian_to_local_curvilinear(x, z, g, xout, sout)
real(rp) :: x, z, g, rho, xout, sout, theta
if (g == 0) then
  xout = x
  sout = z
  return
endif
rho = 1/g
theta = atan2(z, abs(x + rho)) 
sout = theta*abs(rho)
if (abs(theta) < pi/2) then
  xout = (x + rho)/cos(theta) - rho
else
  xout = z/sin(theta) - rho
endif

end subroutine






end module gpt_interface_mod