module wave_instationary_module
   implicit none
   save
   private
   public:: wave_instationary

contains

   subroutine wave_instationary(s,par)

      use params,               only: parameters
      use spaceparams
      use roelvink_module,      only: roelvink, baldock, janssen_battjes
      use wave_functions_module,only: compute_wave_direction_velocities, slope2d,breakerdelay, advecqy, advecthetaho, &
      &                               advecxho, advecqx, advecyho, compute_wave_forces, compute_stokes_drift
      use xmpi_module
      use mnemmodule
      use interp,               only: Linear_interp
      use paramsconst

      implicit none

      type(spacepars), target     :: s
      type(parameters)            :: par

      integer                     :: i
      integer                     :: j
      integer                     :: itheta
      integer                     :: dummy



      real*8 , dimension(:,:)  ,allocatable,save  :: dhdx,dhdy,dudx,dudy,dvdx,dvdy,ustw,Erfl
      real*8 , dimension(:,:)  ,allocatable,save  :: km,kmx,kmy,xwadvec,ywadvec,sinh2kh !,wm
      real*8 , dimension(:,:,:),allocatable,save  :: xadvec,yadvec,thetaadvec,dd,drr,dder
      real*8 , dimension(:,:,:),allocatable,save  :: xradvec,yradvec,thetaradvec
      real*8 , dimension(:,:)  ,allocatable,save  :: dkmxdx,dkmxdy,dkmydx,dkmydy,cgxm,cgym,arg,fac
      real*8 , dimension(:,:)  ,allocatable,save  :: uorb,hhwlocal
      real*8 , dimension(:)    ,allocatable,save  :: wcrestpos
      logical, dimension(:,:)  ,allocatable,save  :: gammax_correct, mask_ee
      real*8                                      :: coffshore



      if (.not. allocated(drr)) then
         allocate(drr         (s%nx+1,s%ny+1,s%ntheta))
         allocate(xadvec      (s%nx+1,s%ny+1,s%ntheta))
         allocate(yadvec      (s%nx+1,s%ny+1,s%ntheta))
         allocate(thetaadvec  (s%nx+1,s%ny+1,s%ntheta))
         allocate(xradvec     (s%nx+1,s%ny+1,s%ntheta))
         allocate(yradvec     (s%nx+1,s%ny+1,s%ntheta))
         allocate(thetaradvec (s%nx+1,s%ny+1,s%ntheta))
         allocate(dd          (s%nx+1,s%ny+1,s%ntheta))
         allocate(dder        (s%nx+1,s%ny+1,s%ntheta))

         allocate(dhdx        (s%nx+1,s%ny+1))
         allocate(dhdy        (s%nx+1,s%ny+1))
         allocate(dudx        (s%nx+1,s%ny+1))
         allocate(dudy        (s%nx+1,s%ny+1))
         allocate(dvdx        (s%nx+1,s%ny+1))
         allocate(dvdy        (s%nx+1,s%ny+1))
         allocate(km          (s%nx+1,s%ny+1))
         allocate(kmx         (s%nx+1,s%ny+1))
         allocate(kmy         (s%nx+1,s%ny+1))
         !       allocate(wm          (nx+1,ny+1))
         allocate(ustw        (s%nx+1,s%ny+1))
         allocate(Erfl        (s%nx+1,s%ny+1)) ! wwvv not used
         allocate(xwadvec     (s%nx+1,s%ny+1))
         allocate(ywadvec     (s%nx+1,s%ny+1))
         allocate(sinh2kh     (s%nx+1,s%ny+1))
         allocate(dkmxdx      (s%nx+1,s%ny+1))
         allocate(dkmxdy      (s%nx+1,s%ny+1))
         allocate(dkmydx      (s%nx+1,s%ny+1))
         allocate(dkmydy      (s%nx+1,s%ny+1))
         allocate(cgxm        (s%nx+1,s%ny+1))
         allocate(cgym        (s%nx+1,s%ny+1))
         allocate(arg         (s%nx+1,s%ny+1))
         allocate(fac         (s%nx+1,s%ny+1))
         allocate(uorb        (s%nx+1,s%ny+1))
         allocate(hhwlocal    (s%nx+1,s%ny+1))
         allocate(wcrestpos   (s%nx+1))
         allocate(gammax_correct(s%nx+1,s%ny+1))
         allocate(mask_ee     (s%nx+1,s%ny+1))


         ! wwvv todo: I think these iniailization are superfluous
         drr         = 0.d0
         xadvec      = 0.d0
         yadvec      = 0.d0
         thetaadvec  = 0.d0
         xradvec     = 0.d0
         yradvec     = 0.d0
         thetaradvec = 0.d0
         dd          = 0.d0
         dder        = 0.d0
         dhdx        = 0.d0
         dhdy        = 0.d0
         dudx        = 0.d0
         dudy        = 0.d0
         dvdx        = 0.d0
         dvdy        = 0.d0
         km          = 0.d0
         kmx         = 0.d0
         kmy         = 0.d0
         !     wm          = 0.d0
         ustw        = 0.d0
         Erfl        = 0.d0
         xwadvec     = 0.d0
         ywadvec     = 0.d0
         sinh2kh     = 0.d0
         dkmxdx      = 0.d0
         dkmxdy      = 0.d0
         dkmydx      = 0.d0
         dkmydy      = 0.d0
         cgxm        = 0.d0
         cgym        = 0.d0
         arg         = 0.d0
         fac         = 0.d0
         uorb        = 0.d0
         hhwlocal    = s%hh
         gammax_correct = .false.
         mask_ee     = .false.
         s%Fx          = 0.d0 ! in spacepars
         s%Fy          = 0.d0 ! in spacepars
      endif

      mask_ee (imin_ee:imax_ee,jmin_ee:jmax_ee)   = .true.  ! Mask tu use in where loops, MPI ok

      ! the local water depth to use in this subroutine generally depend on whether we're using wci or not
      ! both options include par%delta*s%H effect
      ! dispersion routine in wave_timestep already accounts for these differences
      if (par%wci==1) then
         hhwlocal = s%hhwcins
      else
         hhwlocal = s%hhw
      endif

      if (par%single_dir==1) then
         s%costh(:,:,1)=cos(s%thetamean-s%alfaz)
         s%sinth(:,:,1)=sin(s%thetamean-s%alfaz)

      elseif (par%snells==0) then
         s%thetamean=(sum(s%ee*s%thet,3)/s%ntheta)/(max(sum(s%ee,3),0.00001d0)/s%ntheta)
      elseif(par%snells==1) then  !Dano: Snellius
         ! Check for borderline cases where critical c/c(1,1) is reached....
#ifdef USEMPI
         if (xmpi_istop .and. xmpi_isleft) then
            coffshore = s%c(1,1)
         else
            coffshore = -huge(0.d0)
         endif
         call xmpi_allreduce(coffshore,MPI_MAX)
#else
         coffshore = s%c(1,1)
#endif
         s%thetamean=asin(max(-1.0d0, min(1.0d0, sin(s%theta0-s%alfaz(1,1))*s%c/coffshore)))+s%alfaz(1,1)
         s%costh(:,:,1)=cos(s%thetamean-s%alfaz)
         s%sinth(:,:,1)=sin(s%thetamean-s%alfaz)
      endif

      ! Slopes of water depth
      call slope2D(hhwlocal,s%nx,s%ny,s%dsu,s%dnv,dhdx,dhdy,s%wete)
      if (par%wci==1) then
         call slope2D(s%uwcins,s%nx,s%ny,s%dsu,s%dnv,dudx,dudy,s%wete)
         call slope2D(s%vwcins,s%nx,s%ny,s%dsu,s%dnv,dvdx,dvdy,s%wete)
      else
         dudx = 0.d0
         dudy = 0.d0
         dvdx = 0.d0
         dvdy = 0.d0
      endif
      !
      ! Calculate once sinh(2kh)
      where(2*hhwlocal*s%k<=3000.d0)
         sinh2kh=sinh(min(2*s%k*hhwlocal,10.0d0))
      elsewhere
         sinh2kh = 3000.d0
      endwhere

      ! split wave velocities in wave grid directions theta
      call compute_wave_direction_velocities(s,par,1,dhdx,dhdy,dudx,dudy,dvdx,dvdy,sinh2kh)
      !
      ! transform to wave action
      !
      do itheta = 1,s%ntheta
         where(s%wete == 1)
            s%ee(:,:,itheta) = s%ee(:,:,itheta)/s%sigt(:,:,itheta)
         endwhere
      enddo
      !
      ! Upwind Euler timestep propagation
      !
      call advecxho(s%ee,s%cgx,xadvec,s%nx,s%ny,s%ntheta,s%dnu,s%dsu,s%dsdnzi,par%scheme,s%wete,par%dt,s%dsz)
      if (s%ny>0) then
         call advecyho(s%ee,s%cgy,yadvec,s%nx,s%ny,s%ntheta,s%dsv,s%dnv,s%dsdnzi,par%scheme,s%wete,par%dt,s%dnz)
      endif
      call advecthetaho(s%ee,s%ctheta,thetaadvec,s%nx,s%ny,s%ntheta,s%dtheta,par%scheme,s%wete)
      !
      do itheta = 1,s%ntheta
         where(s%wete==1)
            s%ee(:,:,itheta)=s%ee(:,:,itheta)-par%dt*(xadvec(:,:,itheta)+yadvec(:,:,itheta)+thetaadvec(:,:,itheta))
         endwhere
      enddo
      !
      ! transform back to wave energy
      !
      do itheta = 1,s%ntheta
         where(s%wete == 1)
            s%ee(:,:,itheta) = max(s%ee(:,:,itheta)*s%sigt(:,:,itheta),0.d0)
         elsewhere
            s%ee(:,:,itheta) = 0.d0
         endwhere
      enddo
      !
      ! Energy integrated over wave directions,Hrms
      !
      where(s%wete == 1)
         s%E=sum(s%ee,3)*s%dtheta
      elsewhere
         s%E = 0.d0
      endwhere
      s%H=sqrt(s%E/par%rhog8)
      !
      ! Correct for gammax in these areas
      where(s%H>par%gammax*s%hh .and. s%wete==1 .and. mask_ee)
         gammax_correct = .true.
      elsewhere
         gammax_correct = .false.
      endwhere
      !
      do itheta=1,s%ntheta
         ! note: here we do use instantaneous water depth (and excluding par%delta*s%H effect)
         where(gammax_correct)
            s%ee(:,:,itheta)=s%ee(:,:,itheta)/(s%H/(par%gammax*s%hh))**2
         endwhere
      enddo
      where(gammax_correct)
         s%H=min(s%H,par%gammax*s%hh)
      endwhere
      where(gammax_correct)
         s%E=par%rhog8*s%H**2 ! thread-safe computation of all H before recomputing E
      endwhere
      !
      ! Total dissipation
      !
      select case(par%break)
       case(BREAK_ROELVINK1,BREAK_ROELVINK2)
         call roelvink(par,s)
       case(BREAK_BALDOCK)
         call baldock(par,s)
       case(BREAK_JANSSEN)
         call janssen_battjes(par,s)
       case(BREAK_ROELVINK_DALY)  ! Dano: probably not MPI ok
         cgxm = s%c*dcos(s%thetamean)
         cgym = s%c*dsin(s%thetamean)
         call advecqx(cgxm,s%Qb,xwadvec,s%nx,s%ny,s%dsu,s%wete)
         if (s%ny>0) then
            call advecqy(cgym,s%Qb,ywadvec,s%nx,s%ny,s%dnv,s%wete)
            s%Qb  = s%Qb-par%dt*(xwadvec+ywadvec)
         else
            s%Qb  = s%Qb-par%dt*xwadvec
         endif
         call roelvink(par,s)
      end select

      ! Dissipation by bed friction
      uorb=par%px*s%H/par%Trep/sinh(min(max(s%k,0.01d0)*s%hhw,10.0d0))
      s%Df= 2d0/(3d0*par%px)*par%rho*s%fw*uorb**3
      where (s%hh>par%fwcutoff)
         s%Df = 0.d0
      end where
      !
      ! Distribution of dissipation over directions and frequencies
      !
      do itheta=1,s%ntheta
         ! Only calculate for E>0 FB
         ! First just the dissipation that is fed to the roller
         dder(:,:,itheta)=s%ee(:,:,itheta)*s%D/max(s%E,0.00001d0)
         ! Then all short wave energy dissipation, including bed friction and vegetation
         dd(:,:,itheta)=dder(:,:,itheta) + s%ee(:,:,itheta)*(s%Df+s%Dveg)/max(s%E,0.00001d0)
      enddo
      !
      ! Euler step dissipation
      !
      ! calculate roller energy balance
      !
      call advecxho(s%rr,s%cx,xradvec,s%nx,s%ny,s%ntheta,s%dnu,s%dsu,s%dsdnzi,par%scheme,s%wete,par%dt,s%dsz)
      if (s%ny>0) then
         call advecyho(s%rr,s%cy,yradvec,s%nx,s%ny,s%ntheta,s%dsv,s%dnv,s%dsdnzi,par%scheme,s%wete,par%dt,s%dnz)
      endif
      !call advectheta(rr*ctheta,thetaradvec,nx,ny,ntheta,dtheta)
      call advecthetaho(s%rr,s%ctheta,thetaradvec,s%nx,s%ny,s%ntheta,s%dtheta,par%scheme,s%wete)
      do itheta=1,s%ntheta
         where(s%wete==1 .and. mask_ee)
            s%rr(:,:,itheta)=s%rr(:,:,itheta) &
            -par%dt*(xradvec(:,:,itheta) &
            +yradvec(:,:,itheta) &
            +thetaradvec(:,:,itheta))
            s%rr(:,:,itheta)=max(s%rr(:,:,itheta),0.0d0)
         elsewhere (s%wete==1)
            s%rr(:,:,itheta) = 0.d0
         endwhere
      enddo
      !
      ! euler step roller energy dissipation (source and sink function)
      !
      do itheta=1,s%ntheta
         do j=jmin_ee,jmax_ee
            do i=imin_ee,imax_ee
               if(s%wete(i,j)==1) then
                  s%ee(i,j,itheta)=s%ee(i,j,itheta)-par%dt*dd(i,j,itheta)
                  if(par%roller==1) then
                     drr(i,j,itheta) = 2*par%g*s%BR(i,j)*max(s%rr(i,j,itheta),0.0d0)/   &
                     sqrt(s%cx(i,j,itheta)**2 +s%cy(i,j,itheta)**2)
                     s%rr(i,j,itheta)=s%rr(i,j,itheta)+par%dt*(dder(i,j,itheta)         &  ! Robert: changed from dd to dder
                     -drr(i,j,itheta))                                            ! (only from wave s%breaking,
                  else                                                           !  not vegetation or bed friction)
                     s%rr(i,j,itheta)= 0.0d0
                     drr(i,j,itheta)= 0.0d0
                  endif
                  s%ee(i,j,itheta)=max(s%ee(i,j,itheta),0.0d0)
                  s%rr(i,j,itheta)=max(s%rr(i,j,itheta),0.0d0)
               else
                  s%ee(i,j,itheta)=0.0d0
                  s%rr(i,j,itheta)=0.0d0
               endif
            enddo
         enddo
      enddo
      !
      ! Bay boundary Robert + Jaap
      !
      ! wwvv
      ! this has consequences for the parallel version,
      ! but also, if we do  nothing, there are discrepancies
      ! between ee and rr in the different processes. We need to
      ! get valid values for ee(nx+1,:,:) and rr(nx+1,:,:) from
      ! the neighbour below. We cannot postpone this until this
      ! subroutine ends, because ee and rr are used in this subroutine

      if (xmpi_isbot) then
         s%ee(s%nx+1,:,:) =s%ee(s%nx,:,:)
         s%rr(s%nx+1,:,:) =s%rr(s%nx,:,:)
      endif

      if (par%t>0.0d0) then
         if (s%ny>0) then
            if (xmpi_isleft)then ! Jaap
               if (par%lateralwave==LATERALWAVE_NEUMANN) then
                  !
                  ! Lateral boundary at y=0;
                  !
                  s%ee(1:s%nx+1,1,:)=s%ee(1:s%nx+1,2,:)
                  s%rr(1:s%nx+1,1,:)=s%rr(1:s%nx+1,2,:)
               elseif (par%lateralwave==LATERALWAVE_WAVECREST) then
                  !   wcrestpos=xz+tan(thetamean(:,2))*(yz(2)-yz(1))
                  wcrestpos=s%sdist(:,1)+tan(s%thetamean(:,2)-s%alfaz(:,2))*s%dnv(:,1)
                  do itheta=1,s%ntheta
                     do i=1,s%nx+1
                        call Linear_interp((/-huge(0.d0)   ,wcrestpos     ,huge(0.d0)/),&
                        (/s%ee(1,2,itheta),s%ee(:,2,itheta),s%ee(s%nx+1,2,itheta)/),&
                        s%nx+1,s%sdist(i,1),s%ee(i,1,itheta),dummy)
                        call Linear_interp((/-huge(0.d0)   ,wcrestpos     ,huge(0.d0)/),&
                        (/s%rr(1,2,itheta),s%rr(:,2,itheta),s%rr(s%nx+1,2,itheta)/),&
                        s%nx+1,s%sdist(i,1),s%rr(i,1,itheta),dummy)
                        s%ee(i,1,itheta)=max(s%ee(i,1,itheta),0.d0)
                        s%rr(i,1,itheta)=max(s%rr(i,1,itheta),0.d0)
                     enddo
                  enddo
               endif
            endif
            if (xmpi_isright)then
               if (par%lateralwave==LATERALWAVE_NEUMANN) then
                  !
                  ! lateral; boundary at y=ny*dy
                  !
                  s%ee(1:s%nx+1,s%ny+1,:)=s%ee(1:s%nx+1,s%ny,:)
                  s%rr(1:s%nx+1,s%ny+1,:)=s%rr(1:s%nx+1,s%ny,:)
               elseif (par%lateralwave==LATERALWAVE_WAVECREST) then
                  !  wcrestpos=xz-tan(thetamean(:,ny))*(yz(ny+1)-yz(ny))
                  wcrestpos=s%sdist(:,s%ny+1)-tan(s%thetamean(:,s%ny)-s%alfaz(:,s%ny))*s%dnv(:,s%ny)
                  do itheta=1,s%ntheta
                     do i=1,s%nx+1
                        call Linear_interp((/-huge(0.d0)   ,wcrestpos     ,huge(0.d0)/),&
                        (/s%ee(1,s%ny,itheta),s%ee(:,s%ny,itheta),s%ee(s%nx+1,s%ny,itheta)/),&
                        s%nx+1,s%sdist(i,s%ny+1),s%ee(i,s%ny+1,itheta),dummy)
                        call Linear_interp((/-huge(0.d0)   ,wcrestpos     ,huge(0.d0)/),&
                        (/s%rr(1,s%ny,itheta),s%rr(:,s%ny,itheta),s%rr(s%nx+1,s%ny,itheta)/),&
                        s%nx+1,s%sdist(i,s%ny+1),s%rr(i,s%ny+1,itheta),dummy)
                        s%ee(i,s%ny+1,itheta)=max(s%ee(i,s%ny+1,itheta),0.d0)
                        s%rr(i,s%ny+1,itheta)=max(s%rr(i,s%ny+1,itheta),0.d0)
                     enddo
                  enddo
               endif
            endif
         endif
      endif
      ! wwvv communicate ee(:,1,:)
#ifdef USEMPI
      call xmpi_shift_ee(s%ee)
      call xmpi_shift_ee(s%rr)
#endif

      !
      ! Energy integrated over wave directions,Hrms
      !
      s%E  = sum(s%ee,3)*s%dtheta
      s%R  = sum(s%rr,3)*s%dtheta
      s%DR = sum(drr,3)*s%dtheta
      s%H  = sqrt(s%E/par%rhog8)
      !Dano: recompute uorb as it is affected by the MPI shifts and time stepping
      uorb=par%px*s%H/par%Trep/sinh(min(max(s%k,0.01d0)*s%hhw,10.0d0))

      !
      ! Compute mean wave direction
      !
      if (par%snells==0 .and. par%single_dir==0) then
         where(s%wete == 1)
            s%thetamean=(sum(s%ee*s%thet,3)/s%ntheta)/(max(sum(s%ee,3),0.00001d0)/s%ntheta)
         endwhere
      endif
      !
      ! compute wave forces
      call compute_wave_forces(s)
      !
      ! orbital velocity
      s%urms=uorb/sqrt(2.d0)
      !
      ! Stokes drift and orbital velocities
      call compute_stokes_drift(s,par)
      
   end subroutine wave_instationary

end module wave_instationary_module