module geometry2D_TOPO
  !! ---------------------------------------------
  !!
  !! Geometry calculation module for RWE2D
  !! Calculates:
  !!		2) Coefficient fields
  !!
  !! Written Jeff Shragge May 16, 2006
  !!			email: jshragge@gmail.com
  !!
  !! ---------------------------------------------
  use SEP
  use util_         !! Regular SEP Utilities

  implicit none
  integer, private :: nz,nx
  real,    private :: dz,dx
  real   , private :: eps
  logical, private :: kin,verbose

  !! GEOMETRY
  real,    allocatable, private :: ray(:,:,:)
  real,    allocatable,dimension(:), private  :: g11,g12,g13,g22,g23,g33
  real,    allocatable,dimension(:), private  :: iG11,iG12,iG13,iG22,iG23,iG33
  real,    allocatable,dimension(:), private  :: m11,m12,m13,m22,m23,m33
  real,    allocatable,dimension(:), private  :: gdetr,n1,n2,n3

  real,    allocatable,dimension(:,:), private  :: gfields


contains

  !----------------------------------------------------------------  
  !!
  !! Geometry initialization routine
  !!
  subroutine geometry_init(nz_in,nx_in,dz_in,dx_in,verbose_in)
    integer          ::    nz_in,nx_in
    real             ::                dz_in,dx_in
    logical          ::                            verbose_in

    nz=nz_in;  nx=nx_in
    dz=dz_in;  dx=dx_in
    verbose=verbose_in;    eps=0.0001;

    !! For Calculating Geometry
    allocate(ray(3,nx,2) )
    allocate(g11(nx),g12(nx),g13(nx),g22(nx),g23(nx),g33(nx))
    allocate(iG11(nx),iG12(nx),iG13(nx),iG22(nx),iG23(nx),iG33(nx))
    allocate(m11(nx),m12(nx),m13(nx),m22(nx),m23(nx),m33(nx))
    allocate(gdetr(nx),n1(nx),n2(nx),n3(nx) )
    g11=0.;g12=0.;g13=0.;g22=0.;g23=0.; g33=0.;
    iG11=0.;iG12=0.;iG13=0.;iG22=0.;iG23=0.;iG33=0.;
    m11=0.;m12=0.;m13=0.;m22=0.;m23=0.;m33=0.;
    gdetr=0.;n1=0.;n2=0.;n3=0.; 

    allocate( gfields(nx,4) )
    gfields=0.;

  end subroutine geometry_init

  subroutine geometry_deallocate()
    deallocate (ray,g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,m11,m12,m13,m22,m23,m33)
    deallocate (gdetr,n1,n2,n3)
  end subroutine geometry_deallocate
  !!----------------------------------------------------------------  
  !! 
  !! Geometry calcuation routine
  !!
  !! Outputs:
  !!		1) fields - coefficient fields
  !!		4) gnorm  - normalization array
  subroutine geometry_calc(rays,Vel,fields,gnorm)
    !! Main subroutine dealing with geometry
    integer :: iz,ix,if
    real    :: rays(:,:,:),Vel(:,:),fields(:,:,:),gnorm(:,:)
    real, allocatable:: m33old(:)

    allocate ( m33old(nx) )

    !! 2D Constants
    m12=0.
    m23=0. 
    g12=0.; 
    g22=1.; 
    g23=0.; 
    n2=0.

    m33old =0.
    iG12 = 0. 
    iG22 = 1. 
    iG23 = 0. 

    do iz=2,nz-1 !! For all extrapolation steps

       !! Metric tensor calculation
       call calc_g11_2d(rays(iz-1:iz+1,:,:),g11);
       call calc_g13_2d(rays(iz-1:iz+1,:,:),g13); 
       if ( maxval(abs(g13)) .lt. 0.01) then
          g13=0.
       end if
       call calc_g33_2d(rays(iz-1:iz+1,:,:),g33);
       call calc_determinant(g11,g13,g33,gdetr)

       !! Inverse metric tensor calculation
       iG11 = g33 / gdetr**2
       iG13 =-g13 / gdetr**2
       iG33 = g11 / gdetr**2

       !! Weighted metric tensor calculation
       m11=gdetr*iG11
       m13=gdetr*iG13
       m22=gdetr
       m33=gdetr*iG33

       !! Stretch/rotate parameter calculation
       call calc_n1(m11,n1)
       call calc_n3(m33,m33old,n3)

       m33old=m33
       if (iz .eq. 2) n3=0.

       !! Non-stationary coefficient fields
       do ix=1,nx
          gfields(ix,1) = n3(ix)/(2.*m33(ix))
       end do
       do ix=1,nx
          gfields(ix,2) = sqrt(1./iG33(ix))
       end do
       do ix=1,nx
          gfields(ix,3) = sqrt(iG11(ix))
       end do
       do ix=1,nx
          gfields(ix,4) = n3(ix)/(2.*m33(ix))
       end do
       !! Store all fields in arrays
       do ix=1,nx
          gnorm (iz,ix ) = gdetr(ix)
       end do
       do ix=1,nx
          do if=1,4
             fields(ix,if,iz) = gfields(ix,if)
          end do
       end do
    end do

    !! Treat ends
    do ix=1,nx
       gnorm(1 ,ix)=gnorm(2   ,ix)
    end do
    do ix=1,nx
       gnorm(nz,ix)=gnorm(nz-1,ix)
    end do

    do ix=1,nx
       do if=1,4
          fields(ix,if,1 )=fields(ix,if,2)
       end do
    end do

    do ix=1,nx
       do if=1,4
          fields(ix,if,nz)=fields(ix,if,nz-1)
       end do
    end do

    write(0,*) 'Computed reference Parameters and Geometry'
    call coeff_report(fields)
    deallocate (m33old)
  end subroutine geometry_calc

  !----------------------------------------------------------------

  subroutine calc_g11_2d(ray,g11)
    integer :: ix
    real,intent(in ) :: ray(:,:,:)
    real,intent(out) :: g11(:)

    do ix=2,nx-1
       g11(ix) = &
       ( ( ray(2,ix+1,1)-ray(2,ix-1,1) ) / (2.*dx) )**2 + &
       ( ( ray(2,ix+1,2)-ray(2,ix-1,2) ) / (2.*dx) )**2 
    end do
    g11(1 ) = g11(   2) 
    g11(nx) = g11(nx-1)

  end subroutine calc_g11_2d

  !----------------------------------------------------------------

  subroutine calc_g33_2d(ray,g33)
    integer :: ix
    real,intent(in ) :: ray(:,:,:)
    real,intent(out) :: g33(:)

    do ix=1,nx
       g33(ix) = &
       ( ( ray(3,ix,1)-ray(1,ix,1) ) / (2.*dz) )**2 + &
       ( ( ray(3,ix,2)-ray(1,ix,2) ) / (2.*dz) )**2 
    end do

  end subroutine calc_g33_2d

  !----------------------------------------------------------------

  subroutine calc_g13_2d(ray,g13)
    integer :: ix
    real,intent(in ) :: ray(:,:,:)
    real,intent(out) :: g13(:)

    do ix=2,nx-1
       g13(ix) = &
       ( ray(2,ix+1,1)-ray(2,ix-1,1) ) * &
       ( ray(3,ix,1  )-ray(1,ix,1  ) ) / (4.*dx*dz) + &
       ( ray(2,ix+1,2)-ray(2,ix-1,2) ) * &
       ( ray(3,ix,2  )-ray(1,ix,2  ) ) / (4.*dx*dz) 
    end do
    g13(1 ) = g13(   2) 
    g13(nx) = g13(nx-1)

  end subroutine calc_g13_2d

  !----------------------------------------------------------------

  subroutine calc_determinant(g11,g13,g33,gdetr)
    integer :: ix
    real, dimension(:),intent(in ) :: g11,g13,g33
    real, dimension(:),intent(out) :: gdetr

    do ix=1,nx
       gdetr(ix) = sqrt( abs(g11(ix)*g33(ix)-g13(ix)**2) )	
    end do
  end subroutine calc_determinant

  !----------------------------------------------------------------

  subroutine calc_inverse(g11,g13,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
    integer :: ix
    real,dimension(:) :: g11,g13,g33,gdetr
    real,dimension(:) :: iG11,iG12,iG13,iG22,iG23,iG33

    do ix=1,nx
       iG11(ix) = g33(ix) / gdetr(ix)**2
    end do
    do ix=1,nx
       iG12(ix) = 0. 
    end do
    do ix=1,nx
       iG13(ix) =-g13(ix) / gdetr(ix)**2
    end do
    do ix=1,nx
       iG22(ix) = 1. 
    end do
    do ix=1,nx
       iG23(ix) = 0. 
    end do
    do ix=1,nx
       iG33(ix) = g11(ix) / gdetr(ix)**2
    end do

  end subroutine calc_inverse

  !----------------------------------------------------------------

  subroutine calc_n1(m11,n)
    integer :: ix
    real, intent(in) :: m11(:)
    real, intent(out):: n(:)

    do ix=2,nx-1
       n(ix) = ( m11(ix+1) - m11(ix-1) )/ (2.*dx) 
    end do

    n(1 ) = n(   2) 
    n(nx) = n(nx-1)
  end subroutine calc_n1


  subroutine calc_n3(m33,m33old,n)
    integer :: ix
    real, intent(in) :: m33(:),m33old(:)
    real, intent(out):: n(:)

    do ix=1,nx
       n(ix) = ( m33(ix) - m33old(ix) )/ dz 
    end do
  end subroutine calc_n3

  !----------------------------------------------------------------

  subroutine coeff_report(fields)
    real  ::  fields(:,:,:)
    write(0,*) "COEFFICIENT RANGES"
    write(0,*) "a3  min/max", minval(fields(2:,1,:)),maxval(fields(2:,1,:))
    write(0,*) "a4  min/max", minval(fields(2:,2,:)),maxval(fields(2:,2,:))
    write(0,*) "a5  min/max", minval(fields(2:,3,:)),maxval(fields(2:,3,:))
    write(0,*) "a10 min/max", minval(fields(2:,4,:)),maxval(fields(2:,4,:))
    write(0,*) " "

  end subroutine coeff_report
  !----------------------------------------------------------------

  subroutine geom_report( )
    write(0,*) 'METRIC COEFFICIENTS'
    write(0,*) "g11 min/max", minval(g11(2:)),maxval(g11(2:))
    write(0,*) "g13 min/max", minval(g13(2:)),maxval(g13(2:))
    write(0,*) "g33 min/max", minval(g33(2:)),maxval(g33(2:))
    write(0,*) "gdetr min/max", minval(gdetr(2:)),maxval(gdetr(2:))
    write(0,*) " "
    write(0,*) 'INVERSE METRIC COEFFICIENTS'
    write(0,*) "G11 min/max", minval(iG11(2:)),maxval(iG11(2:))
    write(0,*) "G13 min/max", minval(iG13(2:)),maxval(iG13(2:))
    write(0,*) "G22 min/max", minval(iG22(2:)),maxval(iG22(2:))
    write(0,*) "G33 min/max", minval(iG33(2:)),maxval(iG33(2:))
    write(0,*) " "

  end subroutine geom_report

  !----------------------------------------------------------------

  subroutine geometry_close()

    ! deallocate(g11,g12,g13,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,m11,m12,m13)
    ! deallocate(m22,m23,m33,gdetr,n1,n2,n3,gfields,grefs,gmask)
    ! deallocate(have_zeros,scaly,slice,iregion,nstart,vslice)

  end subroutine geometry_close

  !----------------------------------------------------------------  

end module geometry2D_TOPO
