module geometry2D
  use SEP
  use util
  use lloyd_mod

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

  integer, private :: nref

  logical, private :: kin,verbose

  !! GEOMETRY
  real,    allocatable, private :: ray(:,:,:)
  real,    allocatable,dimension(:), private  :: g11,g13,g33
  real,    allocatable,dimension(:), private  :: iG11,iG13,iG33
  real,    allocatable,dimension(:), private  :: m11,m13,m33
  real,    allocatable,dimension(:), private  :: gdetr,n1,n3

  integer, allocatable,dimension(:  ), private  :: allmask
  real,    allocatable,dimension(:,:), private  :: allfields

  !! For Reference Calculation
  integer, allocatable,private :: iregion(:),nstart(:)
  real,    allocatable,private :: slice(:,:),scales(:),have_zeros(:)
  real,    allocatable,private :: vslice(:),allrefs(:,:)
contains

  !----------------------------------------------------------------  
  subroutine geometry_init(nz_in,nx_in,dz_in,dx_in,kin_in,nref_in,verbose_in)
    integer          ::    nz_in,nx_in,                   nref_in
    real             ::                dz_in,dx_in
    logical          ::                            kin_in,        verbose_in

    nz=nz_in;  nx=nx_in; 
    dz=dz_in;  dx=dx_in; 
    nref=nref_in;    kin=kin_in
    verbose=verbose_in;

    !! For Calculating Geometry
    allocate(ray(3,nx,2) )
    allocate(g11(nx),g13(nx),g33(nx))
    allocate(iG11(nx),iG13(nx),iG33(nx))
    allocate(m11(nx),m13(nx),m33(nx))
    allocate(gdetr(nx),n1(nx),n3(nx),vslice(nx) )
    g11=0.;g13=0.; g33=0.;
    iG11=0.;iG13=0.;iG33=0.;
    m11=0.; m13=0.; m33=0.;
    gdetr=0.;n1=0.;n3=0.; vslice=0.

    allocate( allfields(nx,6), allrefs(nref,6), allmask(nx) )
    allfields=0.; allrefs=0.;allmask=0;

    !! For Calculating Reference Parameters
    allocate(have_zeros(6));have_zeros=0.;
    allocate(scales(6)); scales=0.
    allocate(slice(nx,6)); slice=0.
    allocate(iregion(nx)); iregion=0;
    allocate(nstart(6));nstart=0;

    call init_lloyd()

  end subroutine geometry_init

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

subroutine geometry_calc(rays,Vel,fields,refs,mask)
  !! Main subroutine dealing with geometry
  integer :: iz,inn,mask(:,:)
  integer, allocatable :: masky(:,:)
  real    :: rays(:,:,:),Vel(:,:),fields(:,:,:),refs(:,:,:)
  allocate( masky(nx,nz) );    masky=0

  do iz=2,nz-1
     ray    = rays(iz-1:iz+1,:,:)
     vslice = Vel(iz,:)
     call calc_g11_2d(ray,g11);
     write(0,*) 'g11',iz, maxval(g11),minval(g11)
     call calc_g13_2d(ray,g13);   
     write(0,*) 'g13',iz, maxval(g13),minval(g13)
     call calc_g33_2d(ray,g33);
     write(0,*) 'g33',iz, maxval(g33),minval(g33)
     call calc_determinant(g11,g13,g33,gdetr)
     write(0,*) 'gdet',iz, maxval(gdetr),minval(gdetr)

     call calc_inverse(g11,g13,g33,iG11,iG13,iG33,gdetr)
     m11=gdetr*iG11;
     m13=gdetr*iG13;
     m33=gdetr*iG33;
     call calc_n_2d(m11,m13,n1)
     call calc_n_2d(m13,m33,n3)

     !! Non-stationary coefficient fields
     allfields(:,1) = iG13/iG33
     allfields(:,2) = n3/(2.*m33)
     allfields(:,3) = sqrt(1./(vslice**2*iG33))
     allfields(:,4) = sqrt(iG11/iG33 - (iG13/iG33)**2)
     allfields(:,5) = n1/m33 - m13*n3/(m33**2)
     allfields(:,6) = n3/m33

     !! Set kinematic stuff to zero
     if (kin) then
        allfields(:,5)=0.
     end if

     !! Compute Reference Fields and Masks
     call get_references(allfields,allrefs,allmask)
     fields(:,:,iz) = allfields
     masky (:,  iz) = allmask
     refs  (:,:,iz) = allrefs

  end do


  !! Treat ends
  fields(:,:,1 )=fields(:,:,2); fields(:,:,nz)=fields(:,:,nz-1)
  refs  (:,:,1 )=refs  (:,:,2); refs  (:,:,nz)=refs  (:,:,nz-1)
  masky (:,1   )=masky (:,2  ); masky (:,nz  )=masky (:,nz-1  )
  mask=masky

  write(0,*) 'Computed reference Parameters and Geometry'
  call coeff_report(fields,refs)

end subroutine geometry_calc

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

  subroutine calc_g11_2d(ray,g11)
    real,intent(in ) :: ray(:,:,:)
    real,intent(out) :: g11(:)
    g11(2:nx-1) = &
    ( ( ray(2,3:nx,1)-ray(2,1:nx-2,1) ) / (2.*dx) )**2 + &
    ( ( ray(2,3:nx,2)-ray(2,1:nx-2,2) ) / (2.*dx) )**2 
    g11(1 ) = g11(   2) 
    g11(nx) = g11(nx-1)
  end subroutine calc_g11_2d

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

  subroutine calc_g33_2d(ray,g33)
    real,intent(in ) :: ray(:,:,:)
    real,intent(out) :: g33(:)
    g33(2:nx-1) = &
    ( ( ray(3,2:nx-1,1)-ray(1,2:nx-1,1) ) / (2.*dz) )**2 + &
    ( ( ray(3,2:nx-1,2)-ray(1,2:nx-1,2) ) / (2.*dz) )**2 
    g33(1 ) = g33(   2) 
    g33(nx) = g33(nx-1)
  end subroutine calc_g33_2d

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

   subroutine calc_g13_2d(ray,g13)
     real,intent(in ) :: ray(:,:,:)
     real,intent(out) :: g13(:)
     g13(2:nx-1) = &
     ( ray(2,3:nx  ,1)-ray(2,1:nx-2,1) ) * &
     ( ray(3,2:nx-1,1)-ray(1,2:nx-1,1) ) / (4.*dx*dz) + &
     ( ray(2,3:nx  ,2)-ray(2,1:nx-2,2) ) * &
     ( ray(3,2:nx-1,2)-ray(1,2:nx-1,2) ) / (4.*dx*dz) 
     g13(1 ) = g13(   2) 
     g13(nx) = g13(nx-1)
   end subroutine calc_g13_2d
 
  !----------------------------------------------------------------

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

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

  end subroutine calc_determinant

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

  subroutine calc_inverse(g11,g13,g33,iG11,iG13,iG33,gdetr)
    real,dimension(:),intent(in ) :: g11,g13,g33,gdetr
    real,dimension(:),intent(out) :: iG11,iG13,iG33
    iG11 =  g33 / gdetr**2
    iG13 = -g13 / gdetr**2
    iG33 =  g11 / gdetr**2
  end subroutine calc_inverse

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

  subroutine calc_n_2d(m11,m13,n)
    real, intent(in) :: m11(:),m13(:)
    real, intent(out):: n(:)
    n(2:nx-1) = ( m11(3:nx) - m11(1:nx-2) )/ (2.*dx)
    n(1 ) = n(   2) 
    n(nx) = n(nx-1)
    end subroutine calc_n_2d

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

  subroutine get_references(allfields,allrefs,allmask)
    integer :: iref,ii
    real   ,intent(in ) :: allfields(:,:)
    real   ,intent(out) :: allrefs(:,:)
    integer,intent(out) :: allmask(:)

    !! allfields (in)  - Input fields to be used (nx,ny,6)
    !! scales (tmp)    - Scaling function (6)
    !! slice (tmp)     - 1D array used to get refs (nx*ny,6)
    !! iregion (tmp)   - What region has the best attribute set (nx*ny)
    !! nstart (tmp)    - Number of regions (nref)
    !! have_zeros(tmp) - zero mask(6)
    !! allmask(out)    - Output mask (nx)
    !! allrefs (out)   - Reference values in (nref,6)
    !! Global for reference call

    nstart = (/ 1,1,1,1,1,1 /); 

    slice=allfields

    do ii=1,6 !! Get new scaling function
       scales(ii) = sum( abs( slice(:,ii) ) )  / real(nx)
    end do

    where (scales .eq. 0.) 
       have_zeros = 0.
    elsewhere
       have_zeros = 1. 
    end where

    do ii=1,6 !! Rescale Data
       if ( have_zeros(ii) .ne. 0.) then
          slice(:,ii)=slice(:,ii)/scales(ii)
       else
          slice(:,ii)=0. 
       end if
    end do

    !! Call ND Lloyd's algorithm
    iref=lloyd_go(slice,allrefs,nref,iregion,nstart) 

    do ii=1,6  !! rescaling 
       if ( have_zeros(ii) .ne. 0.) then
          allrefs(:,ii)=allrefs(:,ii)*scales(ii)
       else
          allrefs(:,ii)=0. 
       end if
    end do

    if(iref < nref)  allrefs(iref+1:nref,:)=-1.

    allmask=iregion !! Mask
  end subroutine get_references

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

  subroutine coeff_report(fields,refs)
    real  ::  fields(:,:,:),refs(:,:,:)

    write(0,*) "COEFFICIENT RANGES"
    write(0,*) "a1 min/max", minval(fields(2:,1,:)),maxval(fields(2:,1,:))
    write(0,*) "a2 min/max", minval(fields(2:,2,:)),maxval(fields(2:,2,:))
    write(0,*) "a3 min/max", minval(fields(2:,3,:)),maxval(fields(2:,3,:))
    write(0,*) "a4 min/max", minval(fields(2:,4,:)),maxval(fields(2:,4,:))
    write(0,*) "a5 min/max", minval(fields(2:,5,:)),maxval(fields(2:,5,:))
    write(0,*) "a6 min/max", minval(fields(2:,6,:)),maxval(fields(2:,6,:))

    write(0,*) " "
    write(0,*) "REFERENCE PARAMETERS"
    write(0,*) "b1 min/max", minval(refs(:,1,:)),maxval(refs(:,1,:))
    write(0,*) "b2 min/max", minval(refs(:,2,:)),maxval(refs(:,2,:))
    write(0,*) "b3 min/max", minval(refs(:,3,:)),maxval(refs(:,3,:))
    write(0,*) "b4 min/max", minval(refs(:,4,:)),maxval(refs(:,4,:))
    write(0,*) "b5 min/max", minval(refs(:,5,:)),maxval(refs(:,5,:))
    write(0,*) "b6 min/max", minval(refs(:,6,:)),maxval(refs(:,6,:))
 

  end subroutine coeff_report

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

  subroutine geometry_close()

    deallocate(g11,g13,g33,iG11,iG13,iG33,m11,m13)
    deallocate(m33,gdetr,n1,n3,allfields,allrefs,allmask)
    deallocate(have_zeros,scales,slice,iregion,nstart)

  end subroutine geometry_close

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

end module geometry2D
