! Copyright 2022, Francois BOUCHUT, Anne MANGENEY, Marc PERUZZETTO, Antoine LUCAS,  Jacques SAINTE-MARIE
! This program is distributed under the terms of the Lesser GPL v3
! This file is part of SHALTOP.


module viscosity

  use parameters
  use states
  use topography

  !Valid only for topographies depending on x only
  implicit double precision (a-h,o-z)
  double precision, dimension(:,:), allocatable :: sqvisc, sqviscmod,sqviscy,sqviscymod
  double precision, dimension(:,:), allocatable :: frhoug,frhoutg,fyrhous,fyrhouts
  double precision, dimension(:,:), allocatable :: v
  integer, dimension(1) :: loc
  
contains

  subroutine viscosity_allocate

    allocate(sqvisc(0:nx,0:ny+1),sqviscmod(0:nx,0:ny+1))
    allocate(sqviscy(0:nx+1,0:ny),sqviscymod(0:nx+1,0:ny))
    allocate(frhoug(0:nx,1:ny),frhoutg(0:nx,1:ny))
    allocate(fyrhous(1:nx,0:ny),fyrhouts(1:nx,0:ny))
    allocate(v(0:nx+1,0:ny+1))
    
  end subroutine viscosity_allocate

  !--------------------------------------------------------------------------------------
  
  subroutine viscosity_bound_conditions

    !c conditions aux limites
    !c gauche
    if (limviscl.eq.-1) then
       do j=1,ny
          u(0,j)=u(nx,j)
          v(0,j)=v(nx,j)
       enddo
    else if (limviscl.eq.0) then
       do j=1,ny
          u(0,j)=2.d0*(0.d0)-u(1,j) ! Dirichlet
          v(0,j)=2.d0*(0.d0)-v(1,j) ! Dirichlet
       enddo
    else if (limviscl.eq.1) then
       do j=1,ny
          u(0,j)=u(1,j)
          v(0,j)=v(1,j)
       enddo
    else
       do j=1,ny
          u(0,j)=-u(1,j)
          v(0,j)=v(1,j)
       enddo
    endif
    !c droite
    if (limviscr.eq.-1) then
       do j=1,ny
          u(nx+1,j)=u(1,j)
          v(nx+1,j)=v(1,j)
       enddo
    else if (limviscr.eq.0) then
       do j=1,ny
          u(nx+1,j)=2.d0*(0.d0)-u(nx,j) ! Dirichlet
          v(nx+1,j)=2.d0*(0.d0)-v(nx,j) ! Dirichlet
       enddo
    else if (limviscr.eq.1) then
       do j=1,ny
          u(nx+1,j)=u(nx,j)
          v(nx+1,j)=v(nx,j)
       enddo
    else
       do j=1,ny
          u(nx+1,j)=-u(nx,j)
          v(nx+1,j)=v(nx,j)
       enddo
    endif
    !c sud
    if (limviscs.eq.-1) then
       do i=1,nx
          u(i,0)=u(i,ny)
          v(i,0)=v(i,ny)
       enddo
    else if (limviscs.eq.0) then
       do i=1,nx
          u(i,0)=2.d0*(0.d0)-u(i,1) ! Dirichlet
          v(i,0)=2.d0*(0.d0)-v(i,1) ! Dirichlet
       enddo
    else if (limviscs.eq.1) then
       do i=1,nx
          u(i,0)=u(i,1)
          v(i,0)=v(i,1)
       enddo
    else
       do i=1,nx
          u(i,0)=u(i,1)
          v(i,0)=-v(i,1)
       enddo
    endif
    !c nord
    if (limviscn.eq.-1) then
       do i=1,nx
          u(i,ny+1)=u(i,1)
          v(i,ny+1)=v(i,1)
       enddo
    else if (limviscn.eq.0) then
       do i=1,nx
          u(i,ny+1)=2.d0*(0.d0)-u(i,ny) ! Dirichlet
          v(i,ny+1)=2.d0*(0.d0)-v(i,ny) ! Dirichlet
       enddo
    else if (limviscn.eq.1) then
       do i=1,nx
          u(i,ny+1)=u(i,ny)
          v(i,ny+1)=v(i,ny)
       enddo
    else
       do i=1,nx
          u(i,ny+1)=u(i,ny)
          v(i,ny+1)=-v(i,ny)
       enddo
    endif

  end subroutine viscosity_bound_conditions

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

  subroutine viscosity_flux(time,delt)

    double precision, intent(in) :: time,delt

    !Border condition for rho
    do j=1,ny
       rho(0,j)=rho(1,j)
       rho(nx+1,j)=rho(nx,j)
    enddo
    do i=1,nx
       rho(i,0)=rho(i,1)
       rho(i,ny+1)=rho(i,ny)
    enddo
    rho(0,0)=rho(1,1)
    rho(nx+1,0)=rho(nx,1)
    rho(0,ny+1)=rho(1,ny)
    rho(nx+1,ny+1)=rho(nx,ny)

    !c calcul des coefficients de viscosite
    do j=0,ny+1
       do i=0,nx
          sqvisc(i,j)=sqrt(viscmu(min(rho(i,j),rho(i+1,j))))
          sqviscmod(i,j)=sqrt((rapvisc+betavisc)*viscmu(min(rho(i,j),rho(i+1,j))))
       enddo
    enddo
    do j=0,ny
       do i=0,nx+1
          sqviscy(i,j)=sqrt(viscmu(min(rho(i,j),rho(i,j+1))))
          sqviscymod(i,j)=sqrt((rapvisc+betavisc)*viscmu(min(rho(i,j),rho(i,j+1))))
       enddo
    enddo
    vmax=maxval((/(((1.d0+(rapvisc+betavisc)/2.d0)*sqvisc(i,j)**2/(min(rho(i,j),rho(i+1,j))+1.d-14)&
         & *(max(ota(i,j),ota(i+1,j)))**2,i=0,nx),j=1,ny)/))
    loc=maxloc((/(((1.d0+(rapvisc+betavisc)/2.d0)*sqvisc(i,j)**2/(min(rho(i,j),rho(i+1,j))+1.d-14)&
         & *(max(ota(i,j),ota(i+1,j)))**2,i=0,nx),j=1,ny)/))
    vymax=maxval((/(((1.d0+(rapvisc+betavisc)/2.d0)*sqviscy(i,j)**2/(min(rho(i,j),rho(i,j+1))+1.d-14)&
         & ,i=1,nx),j=0,ny)/))
    !tmp=maxval((/(()/))

    !c debut boucle en temps pour les sous-pas de temps paraboliques
    tempsp=time
    lp=0
    do while (tempsp.lt.time+delt)
       lp=lp+1

       !c vitesse dans la direction orthogonale a l'ecoulement
       do j=0,ny+1
          do i=0,nx+1
             v(i,j)=ut(i,j)*ota(i,j)
          enddo
       enddo

       !Boundary conditions
       call viscosity_bound_conditions

       !c calcul du pas de temps parabolique
       !tmp=0
       deltp=min(delt,cflvisc*delx**2/vmax,cflvisc*dely**2/vymax,time+delt-tempsp)
       !write(*,*) vmax
       if (tempsp+deltp.gt.time+delt) deltp=time+delt-tempsp
       sigp=deltp/delx
       sigyp=deltp/dely

       !c calcul des flux de viscosite
       u(0,0)=u(1,0)+u(0,1)-u(1,1)
       u(nx+1,0)=u(nx,0)+u(nx+1,1)-u(nx,1)
       u(0,ny+1)=u(1,ny+1)+u(0,ny)-u(1,ny)
       u(nx+1,ny+1)=u(nx,ny+1)+u(nx+1,ny)-u(nx,ny)
       v(0,0)=v(1,0)+v(0,1)-v(1,1)
       v(nx+1,0)=v(nx,0)+v(nx+1,1)-v(nx,1)
       v(0,ny+1)=v(1,ny+1)+v(0,ny)-v(1,ny)
       v(nx+1,ny+1)=v(nx,ny+1)+v(nx+1,ny)-v(nx,ny)
       do j=1,ny
          do i=0,nx
             frhoug(i,j)=-(2.d0+rapvisc)*sqvisc(i,j)**2*(u(i+1,j)-u(i,j))/delx*.5d0*(ota(i,j)+ota(i+1,j))&
                  & +betavisc*sqvisc(i,j)*(sqviscy(i,j)*(v(i,j+1)-v(i,j))&
                  & +sqviscy(i,j-1)*(v(i,j)-v(i,j-1))&
                  & +sqviscy(i+1,j)*(v(i+1,j+1)-v(i+1,j))&
                  & +sqviscy(i+1,j-1)*(v(i+1,j)-v(i+1,j-1)))/dely/4.d0&
                  & -sqviscmod(i,j)*(sqviscymod(i,j)*(v(i,j+1)-v(i,j))&
                  & +sqviscymod(i,j-1)*(v(i,j)-v(i,j-1))&
                  & +sqviscymod(i+1,j)*(v(i+1,j+1)-v(i+1,j))&
                  & +sqviscymod(i+1,j-1)*(v(i+1,j)-v(i+1,j-1)))/dely/4.d0
             frhoutg(i,j)=-sqvisc(i,j)**2*(v(i+1,j)-v(i,j))/delx*.5d0*(ota(i,j)+ota(i+1,j))&
                  & -sqvisc(i,j)*(sqviscy(i,j)*(u(i,j+1)-u(i,j))&
                  & +sqviscy(i,j-1)*(u(i,j)-u(i,j-1))&
                  & +sqviscy(i+1,j)*(u(i+1,j+1)-u(i+1,j))&
                  & +sqviscy(i+1,j-1)*(u(i+1,j)-u(i+1,j-1)))/dely/4.d0
          enddo
       enddo
       do j=0,ny
          do i=1,nx
             fyrhous(i,j)=-sqviscy(i,j)**2*(u(i,j+1)-u(i,j))/dely&
                  & -sqviscy(i,j)*(sqvisc(i,j)*(v(i+1,j)-v(i,j))*.5d0*(ota(i,j)+ota(i+1,j))&
                  & +sqvisc(i-1,j)*(v(i,j)-v(i-1,j))*.5d0*(ota(i,j)+ota(i-1,j))&
                  & +sqvisc(i,j+1)*(v(i+1,j+1)-v(i,j+1))*.5d0*(ota(i+1,j+1)+ota(i,j+1))&
                  & +sqvisc(i-1,j+1)*(v(i,j+1)-v(i-1,j+1))*.5d0*(ota(i,j+1)+ota(i-1,j+1)))/delx/4.d0
             fyrhouts(i,j)=-(2.d0+rapvisc)*sqviscy(i,j)**2*(v(i,j+1)-v(i,j))/dely&
                  & +betavisc*sqviscy(i,j)*(sqvisc(i,j)*(u(i+1,j)-u(i,j))*.5d0*(ota(i,j)+ota(i+1,j))&
                  & +sqvisc(i-1,j)*(u(i,j)-u(i-1,j))*.5d0*(ota(i,j)+ota(i-1,j))&
                  & +sqvisc(i,j+1)*(u(i+1,j+1)-u(i,j+1))*.5d0*(ota(i+1,j+1)+ota(i,j+1))&
                  & +sqvisc(i-1,j+1)*(u(i,j+1)-u(i-1,j+1))*.5d0*(ota(i,j+1)+ota(i-1,j+1)))/delx/4.d0&
                  & -sqviscymod(i,j)*(sqviscmod(i,j)*(u(i+1,j)-u(i,j))*.5d0*(ota(i,j)+ota(i+1,j))&
                  & +sqviscmod(i-1,j)*(u(i,j)-u(i-1,j))*.5d0*(ota(i,j)+ota(i-1,j))&
                  & +sqviscmod(i,j+1)*(u(i+1,j+1)-u(i,j+1))*.5d0*(ota(i+1,j+1)+ota(i,j+1))&
                  & +sqviscmod(i-1,j+1)*(u(i,j+1)-u(i-1,j+1))*.5d0*(ota(i,j+1)+ota(i-1,j+1)))/delx/4.d0
          enddo
       enddo
       !c calcul des nouvelles vitesses
       do j=1,ny
          do i=1,nx
             !Smooth velocity if r is too small
             r2=states_smooth_velocity(rho(i,j))
             unew=u(i,j)+(-sigp*ota(i,j)*(frhoug(i,j)-frhoug(i-1,j))&
                  & -sigyp*(fyrhous(i,j)-fyrhous(i,j-1)))/r2
             utnew=ut(i,j)+(-sigp*ota(i,j)*(frhoutg(i,j)-frhoutg(i-1,j))&
                  & -sigyp*(fyrhouts(i,j)-fyrhouts(i,j-1)))/r2/ota(i,j)
             u(i,j)=unew
             ut(i,j)=utnew
          enddo
       enddo

       tempsp=tempsp+deltp
    enddo ! fin des sous-pas de temps paraboliques

    if (tempsp.ne.time+delt) then
       write(*,*)'tempsp=',tempsp
       write(*,*)'temps=',time+delt
    endif
    !if (lp.gt.1) then
    !   write(*,*)'nbre de sous-pas de temps pour la viscosite :',lp
    !endif

  end subroutine viscosity_flux

  function viscmu(rho)
    viscmu=.5d0*viscnu*rho**(1.5d0) ! correspond a la formule de Rocha, Johnson, Gray 2018
  end function viscmu
  
end module viscosity
