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

  use parameters
  use read_files
  use topography

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

  double precision, dimension(:,:), allocatable :: rho_source,u_source,ut_source

contains
  
  !Allocate tables
  subroutine sources_allocate_tables
    allocate (rho_source(0:nx+1,0:ny+1),u_source(0:nx+1,0:ny+1),ut_source(0:nx+1,0:ny+1))
  end subroutine sources_allocate_tables

  !----------------------------------------------------------
  !Read source for thickness
  subroutine sources_rho_init

    if (i_rho_source.gt.0) then
       allocate (rho_source(0:nx+1,0:ny+1))
       call read_file(file_rho_source,rho_source)
       !if ipr.eq.1, a constant thickness is imposed in rho_source
       if (i_rho_source.eq.2) then
          !if ipr.eq.2, the source term is a flux and must be converted to heights
          if (source_flux.ne.ndv) then
             !source_flux is the total flux imposed on the area where rho_source>=0
             where (rho_source.ne.ndv)
                rho_source=1.
             elsewhere
                rho_source=0.
             end where
             S=delx*dely*sum((/((rho_source(i,j)/thet(i,j),i=1,nx),j=1,ny)/))
             where (rho_source.gt.0)
                rho_source=source_flux/S
             elsewhere
                rho_source=ndv
             end where
          else
             !Convert flux given in rho_source (m3/s) in thickness per second, by deviding by local
             !grid node size delx*dely/cos(theta)
             do j=1,ny
                do i=1,nx
                   if (rho_source(i,j).ne.ndv) then
                      rho_source(i,j)=rho_source(i,j)*thet(i,j)/(delx*dely)
                   end if
                enddo
             enddo
          end if
       end if
    end if 
  end subroutine sources_rho_init

  !----------------------------------------------------------
  !Read source for velocity
  subroutine sources_vel_init

    if (i_vel_source.gt.0) then
       allocate (u_source(0:nx+1,0:ny+1))
       if (file_u_source.eq.'n.a.') then
          u_source=ndv
       else 
          call read_file(file_u_source,u_source)
       endif
       if (i_vel_source.eq.1) then
          !u_source and ut_source are given directly in file_u_source and file_ut_source
          allocate (ut_source(0:nx+1,0:ny+1))
          if (file_ut_source.eq.'n.a.') then
             ut_source=ndv
          else 
             call read_file(file_ut_source,ut_source)
          end if
       elseif (i_vel_source.eq.2) then
          !velocity is uniform on the mask given in file_u_source, with direction vel_source_dir
          !and norm vel_source_norm
          allocate (ut_source(0:nx+1,0:ny+1))
          if ((vel_source_dir.ne.ndv).and.(vel_source_dir.ne.(-ndv))) then
             !velocity direction is given by vel_source_dir
             alpha=vel_source_dir*pi/180
             ca=cos(alpha)
             sa=sin(alpha)
             do j=1,ny
                do i=1,nx
                   uhoriz=vel_source_norm/sqrt(thet(i,j)**2+(sp(i,j)*ca+spt(i,j)*sa)**2)
                   if (u_source(i,j).eq.ndv) then
                      ut_source(i,j)=ndv
                   else
                      ut_source(i,j)=sa*uhoriz
                      u_source(i,j)=ca*uhoriz
                   endif
                enddo
             enddo
          else 
             if (vel_source_dir.eq.ndv) then
                sgn=-1.d0
             elseif (vel_source_dir.eq.(-ndv)) then
                sgn=1.d0
             endif
             !velocity direction is along the steepest slope. In this case vel_source_norm is
             !both the real velocity and sqrt(u_source**2+ut_source**2)
             do j=1,ny
                do i=1,nx
                   tmp=sqrt(sp(i,j)**2+spt(i,j)**2)
                   ca=sgn*sp(i,j)/tmp
                   sa=sgn*spt(i,j)/tmp
                   if (u_source(i,j).eq.ndv) then
                      ut_source(i,j)=ndv
                   else
                      ut_source(i,j)=sa*vel_source_norm
                      u_source(i,j)=ca*vel_source_norm
                   end if
                enddo
             enddo
          end if
       endif
    end if
    
  end subroutine sources_vel_init
  

  !--------------------------------------------------------
  !Udpate thicknesses with sources terms
  subroutine sources_update_thickness(rho2,time,dt)
    double precision, dimension(0:nx+1,0:ny+1), intent(inout) :: rho2
    double precision, intent(in) :: time,dt
    double precision :: coeff
    if((i_rho_source.gt.0).and.(time.ge.t_source_start).and.(time.le.t_source_end)) then
       if (i_rho_source.eq.1) then
          !Impose constant thickness
          do j=1,ny
             do i=1,nx
                if (rho_source(i,j).ne.ndv) then
                   rho2(i,j)=rho_source(i,j)
                end if
             enddo
          enddo
       elseif (i_rho_source.eq.2) then
          if (i_source_flux_var.eq.0) then
             !Impose constant flux
             coeff = 1.
          elseif (i_source_flux_var.eq.(-1)) then
             !Impose linearly decreasing flux
             coeff = (time - t_source_end)/(t_source_start - t_source_end)
          elseif (i_source_flux_var.eq.1) then
             !Impose linearly increasing flux
             coeff = (time - t_source_start)/(t_source_end - t_source_start)
          end if
          do j=1,ny
             do i=1,nx
                if (rho_source(i,j).ne.ndv) then
                   rho2(i,j)=max(0.d0,rho2(i,j)+rho_source(i,j)*coeff*dt)
                end if
             enddo
          enddo
       end if
    endif
  end subroutine sources_update_thickness

  !---------------------------------------------------------
  !Update velocities with source terms
  subroutine sources_update_velocity(velx,vely,time)
    double precision, dimension(0:nx+1,0:ny+1), intent(inout) :: velx,vely
    double precision, intent(in) :: time
    if((i_vel_source.gt.0).and.(time.ge.t_source_start).and.(time.le.t_source_end)) then
       do j=1,ny
          do i=1,nx
             if (u_source(i,j).ne.ndv) then
                velx(i,j)=u_source(i,j)
             end if
             if (ut_source(i,j).ne.ndv) then
                vely(i,j)=ut_source(i,j)
             end if
          enddo
       enddo
    end if
  end subroutine sources_update_velocity
  
end module sources
