! 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 states

  use parameters
  use init_functions
  use topography
  use read_files
  use sources

  implicit double precision (a-h,o-z)

  double precision :: temps,delt,sig,sigy,init_vol
  double precision, dimension(:,:), allocatable :: rho,u,ut,etadiss
  double precision, dimension(:,:), allocatable :: du_dx,du_dy,dut_dx,dut_dy
  double precision, dimension(:,:), allocatable :: dtu,dtut
  double precision, dimension(:,:), allocatable :: ug,ud,us,un
  double precision, dimension(:,:), allocatable :: utg,utd,uts,utn
  double precision, dimension(:,:), allocatable :: rhog,rhod,rhos,rhon
  double precision, dimension(:,:), allocatable :: rhosave,usave,utsave
  double precision, dimension(:,:), allocatable :: qgsave,qdsave,qssave,qnsave
  double precision, dimension(:,:), allocatable :: qtgsave,qtdsave,qtssave,qtnsave
  double precision, dimension(:,:), allocatable :: etadisssave
  double precision, dimension(:,:), allocatable :: rhomax,unormmax
  double precision, dimension(:,:), allocatable :: momentummax,ekmax
  double precision, dimension(:,:), allocatable :: pimpactmax,fimpactmax
  
contains
  
  !-----------------------------------------------------------------

  !Allocate tables
  subroutine states_allocate_tables
    allocate (rho(0:nx+1,0:ny+1),ut(0:nx+1,0:ny+1),u(0:nx+1,0:ny+1))
    allocate (du_dx(0:nx+1,0:ny+1),du_dy(0:nx+1,0:ny+1), &
         & dut_dx(0:nx+1,0:ny+1),dut_dy(0:nx+1,0:ny+1))
    allocate (etadiss(0:nx+1,0:ny+1),etadisssave(0:nx+1,0:ny+1))
    allocate (dtu(0:nx+1,0:ny+1),dtut(0:nx+1,0:ny+1))
    allocate (rhog(0:nx+1,0:ny+1),rhod(0:nx+1,0:ny+1),rhos(0:nx+1,0:ny+1),rhon(0:nx+1,0:ny+1))
    allocate (ug(0:nx+1,0:ny+1),ud(0:nx+1,0:ny+1),us(0:nx+1,0:ny+1),un(0:nx+1,0:ny+1))
    allocate (utg(0:nx+1,0:ny+1),utd(0:nx+1,0:ny+1),uts(0:nx+1,0:ny+1),utn(0:nx+1,0:ny+1))
    allocate (qgsave(0:nx+1,0:ny+1),qdsave(0:nx+1,0:ny+1),qssave(0:nx+1,0:ny+1),qnsave(0:nx+1,0:ny+1))
    allocate (qtgsave(0:nx+1,0:ny+1),qtdsave(0:nx+1,0:ny+1),qtssave(0:nx+1,0:ny+1),qtnsave(0:nx+1,0:ny+1))
    allocate (rhosave(0:nx+1,0:ny+1),usave(0:nx+1,0:ny+1),utsave(0:nx+1,0:ny+1))    
    allocate (rhomax(0:nx+1,0:ny+1),unormmax(0:nx+1,0:ny+1))
    allocate (momentummax(0:nx+1,0:ny+1),ekmax(0:nx+1,0:ny+1))
    allocate (pimpactmax(0:nx+1,0:ny+1),fimpactmax(0:nx+1,0:ny+1))
    rho=0.
    ut=0.
    u=0.
  end subroutine states_allocate_tables
  

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

  !Initialize states
  subroutine states_init

    call states_mass_init
    call states_velocity_init

    !Compute initial volume
    init_vol=delx*dely*sum((/((rho(i,j)/thet(i,j),i=1,nx),j=1,ny)/))
    print*,'Initial volume : ',init_vol

    !Adjust Friction coefficient if Lucas power law is used
    if (icomp.eq.11) then
       delta1=atan(adelta*init_vol**(gammadelta))
       icomp=1
       print*,'delta1=adelta*init_vol**(gammadelta) = ',delta1
    endif
    
  end subroutine states_init

  !-----------------------------------------------------
  
  !Set initial mass
  subroutine states_mass_init
    !Read initial mass
    if (ipr.eq.0) then
       
		!From file
		call read_file(file_m_init,rho)
		if (hinit_vert.eq.1) then
			!Convert vertical thickness to thickness normal to the topography
			!Note that thicknesses in output will be normal to the topography
			print*, 'Initial thickness is given as vertical and multiplied by cos(slope)'
			rho=rho*thet
		endif

    else
       !From embeded functions
       do j = 1, ny
			do i = 1, nx
				rho(i,j) = rho0(grillex(i),grilley(j))
          enddo
       enddo
    endif

    !Add source term if any
    call sources_update_thickness(rho,0.d0,0.d0)

  end subroutine states_mass_init

  !---------------------------------------------------------------
  
  !set initial velocity
  subroutine states_velocity_init

    if (initu.eq.0) then
       
       !From file
       call read_file(file_u_init,u)

    end if

    if (initut.eq.0) then
       
       !From file
       call read_file(file_ut_init,ut)

    end if

    !Add source term if any
    call sources_update_velocity(u,ut,0.d0)

  end subroutine states_velocity_init

  !--------------------------------------------------------------
  !Add source terms if any, during the simulation
  subroutine states_add_sources
    call sources_update_thickness(rho,temps,delt)
    call sources_update_velocity(u,ut,temps)
  end subroutine states_add_sources

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

  !Compute gradients
  subroutine states_gradients
    du_dx(1:nx,1:ny)=(u(2:nx+1,:)-u(0:nx-1,:))*.5/delx
    dut_dx(1:nx,1:ny)=(ut(2:nx+1,:)-ut(0:nx-1,:))*.5/delx
    du_dy(1:nx,1:ny)=(u(:,2:ny+1)-u(:,0:ny-1))*.5/dely
    dut_dy(1:nx,1:ny)=(ut(:,2:ny+1)-ut(:,0:ny-1))*.5/dely
  end subroutine states_gradients

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

  !Save states
  subroutine states_save
    do j=1,ny
       do i=1,nx
          rhosave(i,j)=rho(i,j)
          usave(i,j)=u(i,j)
          utsave(i,j)=ut(i,j)
       enddo
    enddo
  end subroutine states_save

  !---------------------------------------------------------------
  
  !Minimum values for rho, u and ut
  subroutine states_min0
    do j=1,ny
       do i=1,nx
          if (rho(i,j).le.eps0) then
             rho(i,j)=0.d0
             u(i,j)=0.d0
             ut(i,j)=0.d0
          end if
          if (abs(u(i,j)).le.eps0) u(i,j)=0.d0
          if (abs(ut(i,j)).le.eps0) ut(i,j)=0.d0
       enddo
    enddo
  end subroutine states_min0

  !--------------------------------------------------------------
  !Increase small thickness for velocity estimation
  function states_smooth_velocity(rho)

    implicit double precision (a-h,o-z)

    select case (smooth_velocity)
    case(0)
       states_smooth_velocity=rho
    case(1)
       states_smooth_velocity=max(rho,eps_smooth_velocity)
    case(2)
       states_smooth_velocity=rho+eps_smooth_velocity
    end select

  end function states_smooth_velocity

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

  !Save flux
  subroutine states_flux_save
    do j=0,ny+1
       do i=0,nx+1
          qgsave(i,j)=rhog(i,j)*ug(i,j)
          qtgsave(i,j)=rhog(i,j)*utg(i,j)
          qdsave(i,j)=rhod(i,j)*ud(i,j)
          qtdsave(i,j)=rhod(i,j)*utd(i,j)
          qssave(i,j)=rhos(i,j)*us(i,j)
          qtssave(i,j)=rhos(i,j)*uts(i,j)
          qnsave(i,j)=rhon(i,j)*un(i,j)
          qtnsave(i,j)=rhon(i,j)*utn(i,j)
       enddo
    enddo
  end subroutine states_flux_save

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

  !Save etadiss
  subroutine states_etadiss_save
    do j=1,ny
       do i=1,nx
          etadisssave(i,j)=etadiss(i,j)
       end do
    end do
  end subroutine states_etadiss_save

  !----------------------------------------------------------------
  
  !Update states if 2nd temporal order
  subroutine states_update_mtordre2

    do j=1,ny
       do i=1,nx
          r=.5d0*(rho(i,j)+rhosave(i,j))
          if (r.gt.0.d0) then
             unew=.5d0*(rho(i,j)*u(i,j)+rhosave(i,j)*usave(i,j))/r
             utnew=.5d0*(rho(i,j)*ut(i,j)+rhosave(i,j)*utsave(i,j))/r
             u(i,j)=unew
             ut(i,j)=utnew
          else
             u(i,j)=.5d0*(u(i,j)+usave(i,j))
             ut(i,j)=.5d0*(ut(i,j)+utsave(i,j))
          endif
          rho(i,j)=r
          !etadiss(i,j)=.5d0*(etadiss(i,j)+etadisssave(i,j))&
          !     & +(r*(u(i,j)**2+ut(i,j)**2+(p(i,j)*u(i,j)+pt(i,j)*ut(i,j))**2)&
          !     & /(1.d0+p(i,j)**2+pt(i,j)**2)/2.d0+ckappa*r**gamma&
          !     & *ota(i,j))/thet(i,j)/delt
          etadiss(i,j)=.5d0*(etadiss(i,j)+etadisssave(i,j))&
               & +(r*(u(i,j)**2+ut(i,j)**2+(p(i,j)*u(i,j)+pt(i,j)*ut(i,j))**2)&
               & /(1.d0+p(i,j)**2+pt(i,j)**2)/2.d0+ckappa*r**gamma&
               & *ota(i,j))/thet(i,j)/delt
          if (max_velocity.gt.0) then
             call states_correct_vel(u(i,j),ut(i,j),p(i,j),pt(i,j),unew,utnew)
             u(i,j)=unew
             ut(i,j)=utnew
          endif
       enddo
    enddo

  end subroutine states_update_mtordre2

  subroutine states_update_max

    do j=1,ny
       do i=1,nx
          r=rho(i,j)
          !Update maximum thickness
          if (r.gt.rhomax(i,j)) then
             rhomax(i,j)=r
          endif
          unorm2=(u(i,j)**2+ut(i,j)**2+(p(i,j)*u(i,j)+pt(i,j)*ut(i,j))**2)/&
               &(1.d0+p(i,j)**2+pt(i,j)**2)
          !Update maximum velocity
          if ((unorm2**0.5.gt.unormmax(i,j)).and.(r.gt.epsh_unormmax)) then
             unormmax(i,j)=unorm2**0.5
          end if
          !Update maximum knetic energy
          tmp=r*unorm2*0.5
          if (tmp.gt.ekmax(i,j)) then
             ekmax(i,j)=tmp
          endif
          !Update maximum momentum
          tmp=r*unorm2**0.5
          if (tmp.gt.momentummax(i,j)) then
             momentummax(i,j)=tmp
          endif
          !Update maximum impact pressure (divided by density)
          tmp=coeff_p_stat*grav*r*0.5
          if (r.gt.epsh_unormmax) then
             tmp=tmp+coeff_p_dyn*unorm2
          end if
          if (tmp.gt.pimpactmax(i,j)) then
             pimpactmax(i,j)=tmp
          end if
          !Update mamximum impact force (divided by density)
          tmp=coeff_p_stat*0.5*grav*r**2+coeff_p_dyn*unorm2*h
          if (tmp.gt.fimpactmax(i,j)) then
             fimpactmax(i,j)=tmp
          end if
       enddo
    end do
    
  end subroutine states_update_max

  subroutine states_correct_vel(u_in,ut_in,p_in,pt_in,u_out,ut_out)
    double precision,intent(in) :: u_in,ut_in,p_in,pt_in
    double precision,intent(out) :: u_out,ut_out
    unorm=sqrt(u_in**2+ut_in**2+(p_in*u_in+pt_in*ut_in)**2/&
         &(1.d0+p_in**2+pt_in**2))
    coef=max_velocity/(unorm+1.d-5)
    if (coef.lt.1) then
       u_out=coef*u_in
       ut_out=coef*ut_in
    else
       u_out=u_in
       ut_out=ut_in
    endif
  end subroutine states_correct_vel


end module states
