module RWE3D_geometry
  use sep
  use util_

  implicit none
  integer, private :: nz,nx,ny
  real,    private :: dz,dx,dy

  integer, private :: nref

  real   , private :: eps
  logical, private :: kin,verbose

  !! GEOMETRY
  real,    allocatable,dimension(:,:), private ::  g11, g12, g13, g22, g23, g33
  real,    allocatable,dimension(:,:), private :: iG11,iG12,iG13,iG22,iG23,iG33
  real,    allocatable,dimension(:,:), private :: gdetr
contains

  !----------------------------------------------------------------  
  subroutine RWE3D_geometry_init(aE1,aE2,aE3)
    type(myaxis)     :: aE1,aE2,aE3

    nx=aE1%n
    ny=aE2%n
    nz=aE3%n

    dx=aE1%d
    dy=aE2%d
    dz=aE3%d

    !! For Calculating Geometry
    allocate( g11(nx,ny), g12(nx,ny), g13(nx,ny), g22(nx,ny), g23(nx,ny), g33(nx,ny))
    allocate(iG11(nx,ny),iG12(nx,ny),iG13(nx,ny),iG22(nx,ny),iG23(nx,ny),iG33(nx,ny))
    allocate(gdetr(nx,ny) )

     g11=0.; g12=0.; g13=0.; g22=0.; g23=0.; g33=0.;
    iG11=0.;iG12=0.;iG13=0.;iG22=0.;iG23=0.;iG33=0.;
    gdetr=0.
  end subroutine RWE3D_geometry_init

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

  subroutine RWE3D_geometry_calc(raysin,fields,gnorm)
    !! Main subroutine dealing with geometry
    integer :: iz,ix,iy,inn
    real    :: raysin(:,:,:,:),fields(:,:,:,:),gnorm(:,:,:)

    do iz=2,nz-1
       if (ny .eq. 1 ) then
          call calc_g11_2dx(raysin(:,:,iz-1:iz+1,:),g11)
          call calc_g13_2dx(raysin(:,:,iz-1:iz+1,:),g13)   
          call calc_g33_2dx(raysin(:,:,iz-1:iz+1,:),g33)
       else if (nx .eq. 1) then
          call calc_g22_2dy(raysin(:,:,iz-1:iz+1,:),g22)
          call calc_g23_2dy(raysin(:,:,iz-1:iz+1,:),g23)  
          call calc_g33_2dy(raysin(:,:,iz-1:iz+1,:),g33)
       else
          call calc_g11(raysin(:,:,iz-1:iz+1,:),g11)
          call calc_g12(raysin(:,:,iz-1:iz+1,:),g12)   
          call calc_g13(raysin(:,:,iz-1:iz+1,:),g13)   
          call calc_g22(raysin(:,:,iz-1:iz+1,:),g22)
          call calc_g23(raysin(:,:,iz-1:iz+1,:),g23)  
          call calc_g33(raysin(:,:,iz-1:iz+1,:),g33)
       end if

       where (abs(g12) .lt. 10e-4) g12=0.
       where (abs(g13) .lt. 10e-4) g13=0.
       where (abs(g23) .lt. 10e-4) g23=0.

       if (ny .eq. 1) then 
          call calc_determinant_2dx(g11,g12,g13,g22,g23,g33,gdetr)
       else if (nx .eq. 1) then
          call calc_determinant_2dy(g11,g12,g13,g22,g23,g33,gdetr)
       else
          call calc_determinant    (g11,g12,g13,g22,g23,g33,gdetr)
       end if

       if (ny .eq. 1) then
          call calc_inverse_2dx(g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
       else if (nx .eq. 1) then
          call calc_inverse_2dy(g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
       else      
          call calc_inverse    (g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
       end if


       !! Non-stationary coefficient fields 
       do iy=1,ny
          do ix=1,nx
             fields(ix,iy,iz,1) = sqrt(         1./iG33(ix,iy))
          end do
       end do

       do iy=1,ny
          do ix=1,nx
             fields(ix,iy,iz,2) = sqrt(iG11(ix,iy)/iG33(ix,iy))
          end do
       end do

       do iy=1,ny
          do ix=1,nx
             fields(ix,iy,iz,3) = sqrt(iG22(ix,iy)/iG33(ix,iy))
          end do
       end do

       do iy=1,ny
          do ix=1,nx
             gnorm(ix,iy,iz)=gdetr(ix,iy)
          end do
       end do

     end do

    !! Treat ends
    fields(:,:, 1,:)=fields(:,:,   2,:) 
    fields(:,:,nz,:)=fields(:,:,nz-1,:)
    


  end subroutine RWE3D_geometry_calc


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

  subroutine calc_g11(ray,g11)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g11(:,:)

    g11(2:nx-1,2:ny-1) = &
    ( ( ray(3:nx,2:ny-1,2,1)-ray(1:nx-2,2:ny-1,2,1) ) / (2.*dx) )**2 + &
    ( ( ray(3:nx,2:ny-1,2,2)-ray(1:nx-2,2:ny-1,2,2) ) / (2.*dx) )**2 + &
    ( ( ray(3:nx,2:ny-1,2,3)-ray(1:nx-2,2:ny-1,2,3) ) / (2.*dx) )**2 
    g11(1 ,:) = g11(   2,:) 
    g11(nx,:) = g11(nx-1,:)
    g11(:,1 ) = g11(:,2   ) 
    g11(:,ny) = g11(:,ny-1)

  end subroutine calc_g11
  !----------------------------------------------------------------

  subroutine calc_g11_2dx(ray,g11)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g11(:,:)

    g11(2:nx-1,1) = &
    ( ( ray(3:nx,1,2,1)-ray(1:nx-2,1,2,1) ) / (2.*dx) )**2 + &
    ( ( ray(3:nx,1,2,3)-ray(1:nx-2,1,2,3) ) / (2.*dx) )**2 
    g11(1 ,1) = g11(   2,1) 
    g11(nx,1) = g11(nx-1,1)

  end subroutine calc_g11_2dx

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

  subroutine calc_g22(ray,g22)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g22(:,:)


    g22(2:nx-1,2:ny-1) = &
    ( ( ray(2:nx-1,3:ny,2,1)-ray(2:nx-1,1:ny-2,2,1) ) / (2.*dy) )**2 + &
    ( ( ray(2:nx-1,3:ny,2,2)-ray(2:nx-1,1:ny-2,2,2) ) / (2.*dy) )**2 + &
    ( ( ray(2:nx-1,3:ny,2,3)-ray(2:nx-1,1:ny-2,2,3) ) / (2.*dy) )**2 

    g22(1 ,:) = g22(   2,:) 
    g22(nx,:) = g22(nx-1,:)
    g22(:,1 ) = g22(:,2   ) 
    g22(:,ny) = g22(:,ny-1)

  end subroutine calc_g22

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

  subroutine calc_g22_2dy(ray,g22)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g22(:,:)

    g22(1,2:ny-1) = &
    ( ( ray(1,3:ny,2,2)-ray(1,1:ny-2,2,2) ) / (2.*dy) )**2 + &
    ( ( ray(1,3:ny,2,3)-ray(1,1:ny-2,2,3) ) / (2.*dy) )**2 

    g22(1,1 ) = g22(1,2   ) 
    g22(1,ny) = g22(1,ny-1)

  end subroutine calc_g22_2dy

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

  subroutine calc_g33(ray,g33)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g33(:,:)

       g33(2:nx-1,2:ny-1) = &
       ( ( ray(2:nx-1,2:ny-1,3,1)-ray(2:nx-1,2:ny-1,1,1) ) / (2.*dz) )**2 + &
       ( ( ray(2:nx-1,2:ny-1,3,2)-ray(2:nx-1,2:ny-1,1,2) ) / (2.*dz) )**2 + &
       ( ( ray(2:nx-1,2:ny-1,3,3)-ray(2:nx-1,2:ny-1,1,3) ) / (2.*dz) )**2 

       g33(1 ,:) = g33(   2,:) 
       g33(nx,:) = g33(nx-1,:)
       g33(:,1 ) = g33(:,2   ) 
       g33(:,ny) = g33(:,ny-1)

  end subroutine calc_g33
  !----------------------------------------------------------------

  subroutine calc_g33_2dx(ray,g33)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g33(:,:)

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

       g33(1 ,1) = g33(   2,1) 
       g33(nx,1) = g33(nx-1,1)

  end subroutine calc_g33_2dx
  !----------------------------------------------------------------

  subroutine calc_g33_2dy(ray,g33)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g33(:,:)

       g33(1,2:ny-1) = &
       ( ( ray(1,2:ny-1,3,2)-ray(1,2:ny-1,1,2) ) / (2.*dz) )**2 + &
       ( ( ray(1,2:ny-1,3,3)-ray(1,2:ny-1,1,3) ) / (2.*dz) )**2 

       g33(1,1 ) = g33(1,2   ) 
       g33(1,ny) = g33(1,ny-1)

  end subroutine calc_g33_2dy
 
  !----------------------------------------------------------------

  subroutine calc_g12(ray,g12)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g12(:,:)
  
       g12(2:nx-1,2:ny-1) = &
       ( ray(3:nx  ,2:ny-1,2,1)-ray(1:nx-2,2:ny-1,2,1) ) * &
       ( ray(2:nx-1,3:ny  ,2,1)-ray(2:nx-1,1:ny-2,2,1) ) / (4.*dx*dy) + &
       ( ray(3:nx  ,2:ny-1,2,2)-ray(1:nx-2,2:ny-1,2,2) ) * &
       ( ray(2:nx-1,3:ny  ,2,2)-ray(2:nx-1,1:ny-2,2,2) ) / (4.*dx*dy) + &
       ( ray(3:nx  ,2:ny-1,2,3)-ray(1:nx-2,2:ny-1,2,3) ) * &
       ( ray(2:nx-1,3:ny  ,2,3)-ray(2:nx-1,1:ny-2,2,3) ) / (4.*dx*dy) 

       g12(1 ,:) = g12(   2,:) 
       g12(nx,:) = g12(nx-1,:)
       g12(:,1 ) = g12(:,2   ) 
       g12(:,ny) = g12(:,ny-1)
  end subroutine calc_g12

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

  subroutine calc_g13(ray,g13)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g13(:,:)


       g13(2:nx-1,2:ny-1) = &
       ( ray(3:nx  ,2:ny-1,2,1)-ray(1:nx-2,2:ny-1,2,1) ) * &
       ( ray(2:nx-1,2:ny-1,3,1)-ray(2:nx-1,2:ny-1,1,1) ) / (4.*dx*dz) + &
       ( ray(3:nx  ,2:ny-1,2,2)-ray(1:nx-2,2:ny-1,2,2) ) * &
       ( ray(2:nx-1,2:ny-1,3,2)-ray(2:nx-1,2:ny-1,1,2) ) / (4.*dx*dz) + &
       ( ray(3:nx  ,2:ny-1,2,3)-ray(1:nx-2,2:ny-1,2,3) ) * &
       ( ray(2:nx-1,2:ny-1,3,3)-ray(2:nx-1,2:ny-1,1,3) ) / (4.*dx*dz) 

       g13(1 ,:) = g13(   2,:) 
       g13(nx,:) = g13(nx-1,:)
       g13(:,1 ) = g13(:,2   ) 
       g13(:,ny) = g13(:,ny-1)

  end subroutine calc_g13

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

  subroutine calc_g13_2dx(ray,g13)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g13(:,:)

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

       g13(1 ,1) = g13(   2,1)
       g13(nx,1) = g13(nx-1,1)

  end subroutine calc_g13_2dx

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

  subroutine calc_g23(ray,g23)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g23(:,:)

    g23(2:nx-1,2:ny-1) = &
    ( ray(2:nx-1,3:ny  ,2,1)-ray(2:nx-1,1:ny-2,2,1) ) * &
    ( ray(2:nx-1,2:ny-1,3,1)-ray(2:nx-1,2:ny-1,1,1) ) / (4.*dy*dz) + &
    ( ray(2:nx-1,3:ny  ,2,2)-ray(2:nx-1,1:ny-2,2,2) ) * &
    ( ray(2:nx-1,2:ny-1,3,2)-ray(2:nx-1,2:ny-1,1,2) ) / (4.*dy*dz) + &
    ( ray(2:nx-1,3:ny  ,2,3)-ray(2:nx-1,1:ny-2,2,3) ) * &
    ( ray(2:nx-1,2:ny-1,3,3)-ray(2:nx-1,2:ny-1,1,3) ) / (4.*dy*dz)       

    g23(1 ,:) = g23(   2,:) 
    g23(nx,:) = g23(nx-1,:)
    g23(:,1 ) = g23(:,2   ) 
    g23(:,ny) = g23(:,ny-1)

  end subroutine calc_g23

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

  subroutine calc_g23_2dy(ray,g23)
    real,intent(in ) :: ray(:,:,:,:)
    real,intent(out) :: g23(:,:)

    g23(1,2:ny-1) = &
    ( ray(1,3:ny  ,2,2)-ray(1,1:ny-2,2,2) ) * &
    ( ray(1,2:ny-1,3,2)-ray(1,2:ny-1,1,2) ) / (4.*dy*dz) + &
    ( ray(1,3:ny  ,2,3)-ray(1,1:ny-2,2,3) ) * &
    ( ray(1,2:ny-1,3,3)-ray(1,2:ny-1,1,3) ) / (4.*dy*dz)       

    g23(1,1 ) = g23(1,2   ) 
    g23(1,ny) = g23(1,ny-1)

  end subroutine calc_g23_2dy

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

  subroutine calc_determinant(g11,g12,g13,g22,g23,g33,gdetr)
    real, dimension(:,:),intent(in ) :: g11,g12,g13,g22,g23,g33
    real, dimension(:,:),intent(out) :: gdetr

    gdetr = sqrt( abs(g11*g22*g33-g11*g23**2-g12**2*g33+2.*g12*g23*g13-g22*g13**2) )	

  end subroutine calc_determinant
  !----------------------------------------------------------------

  subroutine calc_determinant_2dx(g11,g12,g13,g22,g23,g33,gdetr)
    real, dimension(:,:),intent(in ) :: g11,g12,g13,g22,g23,g33
    real, dimension(:,:),intent(out) :: gdetr

    gdetr = sqrt( abs(g11*g33 - g13**2) )	

  end subroutine calc_determinant_2dx
  !----------------------------------------------------------------

  subroutine calc_determinant_2dy(g11,g12,g13,g22,g23,g33,gdetr)
    real, dimension(:,:),intent(in ) :: g11,g12,g13,g22,g23,g33
    real, dimension(:,:),intent(out) :: gdetr

    gdetr = sqrt( abs(g22*g33 - g23**2) )	

  end subroutine calc_determinant_2dy

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

  subroutine calc_inverse(g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
    real,dimension(:,:),intent(in ) :: g11,g12,g13,g22,g23,g33,gdetr
    real,dimension(:,:),intent(out) :: iG11,iG12,iG13,iG22,iG23,iG33

    iG11 = ( g22*g33-g23*g23 ) / gdetr**2
    iG12 = ( g13*g23-g12*g33 ) / gdetr**2
    iG13 = ( g12*g23-g13*g22 ) / gdetr**2
    iG22 = ( g11*g33-g13*g13 ) / gdetr**2
    iG23 = ( g13*g12-g11*g23 ) / gdetr**2
    iG33 = ( g11*g22-g12*g12 ) / gdetr**2

  end subroutine calc_inverse
  !----------------------------------------------------------------

  subroutine calc_inverse_2dx(g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
    real,dimension(:,:),intent(in ) :: g11,g12,g13,g22,g23,g33,gdetr
    real,dimension(:,:),intent(out) :: iG11,iG12,iG13,iG22,iG23,iG33

    iG11 = g33 / gdetr**2
    iG12 = 0.
    iG13 =-g13 / gdetr**2
    iG22 = 1.
    iG23 = 0.
    iG33 = g11 / gdetr**2

  end subroutine calc_inverse_2dx
  !----------------------------------------------------------------

  subroutine calc_inverse_2dy(g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
    real,dimension(:,:),intent(in ) :: g11,g12,g13,g22,g23,g33,gdetr
    real,dimension(:,:),intent(out) :: iG11,iG12,iG13,iG22,iG23,iG33

    iG11 = 1.
    iG12 = 0.
    iG13 = 0.
    iG22 =  g33  / gdetr**2
    iG23 = -g23  / gdetr**2
    iG33 =  g22  / gdetr**2

  end subroutine calc_inverse_2dy

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

  subroutine RWE3D_geometry_close()
    deallocate(g11,g12,g13,g22,g23,g33,iG11,iG12,iG13,iG22,iG23,iG33,gdetr)
  end subroutine RWE3D_geometry_close

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

end module RWE3D_geometry
