module mod_initflow
  use mpi
  use decomp_2d
  use mod_common_mpi, only: ierr,myid
  use mod_param     , only: dims,pi,dx,dy,dz,lx,ly,lz,uref,lref,is_wallturb,bforce
  use mod_types
  implicit none
  private
  public initflow,add_noise
  contains
  subroutine initflow(inivel,ng,lo,zclzi,dzclzi,dzflzi,visc,u,v,w,p)
    !
    ! computes initial conditions for the velocity field
    !
    implicit none
    character(len=3), intent(in) :: inivel
    integer , intent(in), dimension(3) :: ng,lo
    real(rp), intent(in), dimension(0:) :: zclzi,dzclzi,dzflzi
    real(rp), intent(in) :: visc
    real(rp), dimension(0:,0:,0:), intent(out) :: u,v,w,p
    real(rp), allocatable, dimension(:) :: u1d
    !real(rp), allocatable, dimension(:,:) :: u2d
    integer :: i,j,k
    real(rp) :: q
    logical :: is_noise,is_mean,is_pair
    real(rp) :: xc,yc,zc,xf,yf,zf
    real(rp) :: reb,retau
    real(rp) :: ubulk
    integer, dimension(3) :: n
    !
    n(:) = shape(p) - 2*1
    allocate(u1d(n(3)))
    is_noise = .false.
    is_mean  = .false.
    is_pair  = .false.
    ubulk = uref
    select case(trim(inivel))
    case('cou')
      call couette(   n(3),zclzi,ubulk,u1d)
    case('poi')
      call poiseuille(n(3),zclzi,ubulk,u1d)
      is_mean=.true.
    case('tbl')
      call temporal_bl(n(3),zclzi*lz,1._rp,visc,uref,u1d)
    case('iop') ! reversed 'poi'
      call poiseuille(n(3),zclzi,ubulk,u1d)
      u1d(:) = u1d(:) - ubulk
    case('zer')
      u1d(:) = 0.
    case('uni')
      u1d(:) = uref
    case('log')
      call log_profile(n(3),zclzi,ubulk*lref/visc,u1d)
      is_noise = .true.
      is_mean = .true.
    case('hcl')
      deallocate(u1d)
      allocate(u1d(2*n(3)))
      call log_profile(2*n(3),zclzi,ubulk*lref/visc,u1d)
      is_noise = .true.
      is_mean=.true.
    case('hcp')
      deallocate(u1d)
      allocate(u1d(2*n(3)))
      call poiseuille(2*n(3),zclzi,ubulk,u1d)
      is_mean = .true.
    case('tgv')
      do k=1,n(3)
        zc = zclzi(k)*2.*pi
        do j=1,n(2)
          yc = (j+lo(2)-1-.5)*dy/ly*2.*pi
          yf = (j+lo(2)-1-.0)*dy/ly*2.*pi
          do i=1,n(1)
            xc = (i+lo(1)-1-.5)*dx/lx*2.*pi
            xf = (i+lo(1)-1-.0)*dx/lx*2.*pi
            u(i,j,k) =  sin(xf)*cos(yc)*cos(zc)*uref
            v(i,j,k) = -cos(xc)*sin(yf)*cos(zc)*uref
            w(i,j,k) = 0.
            p(i,j,k) = 0.!(cos(2.*xc)+cos(2.*yc))*(cos(2.*zc)+2.)/16.*uref**2
          end do
        end do
      end do
    case('tgw')
      do k=1,n(3)
        zc = zclzi(k)
        do j=1,n(2)
          yc = (j+lo(2)-1-.5)*dy
          yf = (j+lo(2)-1-.0)*dy
          do i=1,n(1)
            xc = (i+lo(1)-1-.5)*dx
            xf = (i+lo(1)-1-.0)*dx
            u(i,j,k) =  cos(xf)*sin(yc)*uref
            v(i,j,k) = -sin(xc)*cos(yf)*uref
            w(i,j,k) = 0.
            p(i,j,k) = -(cos(2.*xc)+cos(2.*yc))/4.*uref**2
          end do
        end do
      end do
    case('pdc')
      if(is_wallturb) then ! turbulent flow
        retau = uref*lref/visc
        reb   = (retau/.09)**(1./.88)
        ubulk = (reb/2.)/retau*uref
      else                 ! laminar flow
        ubulk = (bforce(1)*lref**2/(3.*visc))
      end if
      call poiseuille(n(3),zclzi,ubulk,u1d)
      is_mean=.true.
    case default
      if(myid == 0) print*, 'ERROR: invalid name for initial velocity field'
      if(myid == 0) print*, ''
      if(myid == 0) print*, '*** Simulation aborted due to errors in the case file ***'
      if(myid == 0) print*, '    check INFO_INPUT.md'
      call decomp_2d_finalize
      call MPI_FINALIZE(ierr)
      error stop
    end select
    if(.not.any(inivel == ['tgv','tgw'])) then
      do k=1,n(3)
        do j=1,n(2)
          do i=1,n(1)
            u(i,j,k) = u1d(k)
            v(i,j,k) = 0.
            w(i,j,k) = 0.
            p(i,j,k) = 0.
          end do
        end do
      end do
    end if
    if(is_noise) then
      call add_noise(ng,lo,123,.05_rp,u(1:n(1),1:n(2),1:n(3)))
      call add_noise(ng,lo,456,.05_rp,v(1:n(1),1:n(2),1:n(3)))
      call add_noise(ng,lo,789,.05_rp,w(1:n(1),1:n(2),1:n(3)))
    end if
    if(is_mean) then
      call set_mean(n,ubulk,dzflzi*(dx/lx)*(dy/ly),u(1:n(1),1:n(2),1:n(3)))
    end if
    if(is_wallturb) is_pair = .true.
    if(is_pair) then
      !
      ! initialize a streamwise vortex pair for a fast transition
      ! to turbulence in a pressure-driven channel:
      !        psi(x,y,z)  = f(z)*g(x,y), with
      !        f(z)        = (1-z**2)**2, and
      !        g(x,y)      = y*exp[-(16x**2-4y**2)]
      ! (x,y,z) --> (streamwise, spanwise, wall-normal) directions
      !
      ! see Henningson and Kim, JFM 1991
      !
      do k=1,n(3)
        zc = 2.*zclzi(k) - 1. ! z rescaled to be between -1 and +1
        zf = 2.*(zclzi(k) + .5*dzflzi(k)) - 1.
        do j=1,n(2)
          yc = ((lo(2)-1+j-0.5)*dy-.5*ly)*2./lz
          yf = ((lo(2)-1+j-0.0)*dy-.5*ly)*2./lz
          do i=1,n(1)
            xc = ((lo(1)-1+i-0.5)*dx-.5*lx)*2./lz
            xf = ((lo(1)-1+i-0.0)*dx-.5*lx)*2./lz
            !u(i,j,k) = u1d(k)
            v(i,j,k) = -1. * gxy(yf,xc)*dfz(zc) * ubulk * 1.5
            w(i,j,k) =  1. * fz(zf)*dgxy(yc,xc) * ubulk * 1.5
            p(i,j,k) = 0.
          end do
        end do
      end do
      !
      ! alternatively, using a Taylor-Green vortex
      ! for the cross-stream velocity components
      ! (commented below)
      !
      !do k=1,n(3)
      !  zc = (zclzi(k)              )*2.*pi
      !  zf = (zclzi(k)+0.5*dzclzi(k))*2.*pi
      !  do j=1,n(2)
      !    yc = (j+lo(2)-1-.5)*dy/ly*2.*pi
      !    yf = (j+lo(2)-1-.0)*dy/ly*2.*pi
      !    do i=1,n(1)
      !      xc = (i+lo(1)-1-.5)*dx/lx*2.*pi
      !      xf = (i+lo(1)-1-.0)*dx/lx*2.*pi
      !      !u(i,j,k) = u1d(k)
      !      v(i,j,k) =  sin(xc)*cos(yf)*cos(zc)*ubulk
      !      w(i,j,k) = -cos(xc)*sin(yc)*cos(zf)*ubulk
      !      p(i,j,k) = 0.!(cos(2.*xc)+cos(2.*yc))*(cos(2.*zc)+2.)/16.
      !    end do
      !  end do
      !end do
    end if
    deallocate(u1d)
  end subroutine initflow
  !
  subroutine add_noise(ng,lo,iseed,norm,p)
    implicit none
    integer , intent(in), dimension(3) :: ng,lo
    integer , intent(in) :: iseed
    real(rp), intent(in) :: norm
    real(rp), intent(inout), dimension(:,:,:) :: p
    integer(4), allocatable, dimension(:) :: seed
    real(rp) :: rn
    integer, dimension(3) :: n
    integer :: i,j,k,ii,jj,kk
    !
    n(:) = shape(p)
    allocate(seed(64))
    seed(:) = iseed
    call random_seed( put = seed )
    do k=1,ng(3)
      kk = k-(lo(3)-1)
      do j=1,ng(2)
        jj = j-(lo(2)-1)
        do i=1,ng(1)
          ii = i-(lo(1)-1)
          call random_number(rn)
          if(ii >= 1.and.ii <= n(1) .and. &
             jj >= 1.and.jj <= n(2) .and. &
             kk >= 1.and.kk <= n(3) ) then
             p(ii,jj,kk) = p(ii,jj,kk) + 2.*(rn-.5)*norm
          end if
        end do
      end do
    end do
  end subroutine add_noise
  !
  subroutine set_mean(n,mean,grid_vol_ratio,p)
  implicit none
  integer , intent(in), dimension(3) :: n
  real(rp), intent(in), dimension(0:) :: grid_vol_ratio
  real(rp), intent(in) :: mean
  real(rp), intent(inout), dimension(:,:,:) :: p
  real(rp) :: meanold
  integer :: i,j,k
  meanold = 0.
  !$OMP PARALLEL DO DEFAULT(none) &
  !$OMP SHARED(n,p,grid_vol_ratio) &
  !$OMP PRIVATE(i,j,k) &
  !$OMP REDUCTION(+:meanold)
  do k=1,n(3)
    do j=1,n(2)
      do i=1,n(1)
        meanold = meanold + p(i,j,k)*grid_vol_ratio(k)
      end do
    end do
  end do
  !$OMP END PARALLEL DO
  call mpi_allreduce(MPI_IN_PLACE,meanold,1,MPI_REAL_RP,MPI_SUM,MPI_COMM_WORLD,ierr)
  !
  if(meanold.ne.0.) then
    !$OMP WORKSHARE
    p(:,:,:) = p(:,:,:)/meanold*mean
    !$OMP END WORKSHARE
  end if
  end subroutine set_mean
  !
  subroutine couette(n,zc,norm,p)
    !
    ! plane couette profile normalized by the wall velocity difference
    !
    implicit none
    integer , intent(in)   :: n
    real(rp), intent(in), dimension(0:) :: zc
    real(rp), intent(in)   :: norm
    real(rp), intent(out), dimension(n) :: p
    integer :: k
    real(rp) :: z
    do k=1,n
      z    = zc(k)
      p(k) = .5*(1.-2.*z)*norm
    end do
  end subroutine couette
  !
  subroutine poiseuille(n,zc,norm,p)
    implicit none
    integer , intent(in)   :: n
    real(rp), intent(in), dimension(0:) :: zc
    real(rp), intent(in)   :: norm
    real(rp), intent(out), dimension(n) :: p
    integer :: k
    real(rp) :: z
    !
    ! plane poiseuille profile normalized by the bulk velocity
    !
    do k=1,n
      z    = zc(k)
      p(k) = 6.*z*(1.-z)*norm
    end do
  end subroutine poiseuille
  !
  subroutine temporal_bl(n,zc,d,nu,norm,p)
    implicit none
    integer , intent(in )   :: n
    real(rp), intent(in ), dimension(0:) :: zc
    real(rp), intent(in )   :: d,nu,norm
    real(rp), intent(out), dimension(n) :: p
    integer  :: k
    real(rp) :: theta
    !
    ! temporal boudary layer profile
    ! with thickness d, viscosity nu, and wall velocity norm (at z=0)
    !
    theta = 54.*nu/norm
    do k=1,n
      p(k)=(0.5+(0.5)*tanh((d/(2.*theta))*(1.-zc(k)/d)))*norm
    end do
  end subroutine temporal_bl
  !
  subroutine log_profile(n,zc,reb,p)
    implicit none
    integer , intent(in)   :: n
    real(rp), intent(in), dimension(0:) :: zc
    real(rp), intent(in)   :: reb
    real(rp), intent(out), dimension(n) :: p
    integer :: k
    real(rp) :: z,retau ! z/lz and bulk Reynolds number
    retau = 0.09*reb**(0.88) ! from Pope's book
    do k=1,n
      z    = zc(k)*2.*retau
      if(z >= retau) z = 2.*retau-z
      p(k) = 2.5*log(z) + 5.5
      if (z <= 11.6) p(k)=z
    end do
  end subroutine log_profile
  !
  ! functions to initialize the streamwise vortex pair
  ! (explained above)
  !
  function fz(zc)
  real(rp), intent(in) :: zc
  real(rp) :: fz
    fz = ((1.-zc**2)**2)
  end function
  !
  function dfz(zc)
  real(rp), intent(in) :: zc
  real(rp) :: dfz
    dfz = -4.*zc*(1.-zc**2)
  end function
  !
  function gxy(xc,yc)
  real(rp), intent(in) :: xc,yc
  real(rp) :: gxy
    gxy = yc*exp(-4.*(4.*xc**2+yc**2))
  end function
  !
  function dgxy(xc,yc)
  real(rp), intent(in) :: xc,yc
  real(rp) :: dgxy
    dgxy = exp(-4.*(4.*xc**2+yc**2))*(1.-8.*yc**2)
  end function
end module mod_initflow


module poisson3d_assembly
use types, only: dp
use linalg, only: inv
use utils, only: assert, stop_error
use constants, only: pi
use sparse, only: coo2csr_canonical
implicit none
private
public assemble_3d, integral, func2quad, func_xyz, assemble_3d_precalc, &
    assemble_3d_coo, assemble_3d_csr, assemble_3d_coo_rhs, assemble_3d_coo_A, &
    local_overlap_matrix, assemble_3d_coo_rhs_spectral, &
    assemble_3d_coo_A_spectral, assemble_3d_coo_A_subdomain, &
    assemble_3d_coo_A_subdomain_spectral

interface
    real(dp) function func_xyz(x, y, z)
    import :: dp
    implicit none
    real(dp), intent(in) :: x, y, z
    end function
end interface

contains

function func2quad(nodes, elems, xiq, func) result(fq)
! Return an array of function 'func' values at quadrature points
real(dp), intent(in):: nodes(:, :), xiq(:)
procedure(func_xyz) :: func
integer, intent(in):: elems(:, :)
integer :: Ne, e, iqx, iqy, iqz
real(dp), dimension(size(xiq), size(xiq), size(xiq), size(elems, 2)) :: fq
real(dp), dimension(size(xiq)) :: x, y, z, xp, yp, zp
real(dp) :: lx, ly, lz
real(dp) :: jacx, jacy, jacz, jac_det
lx = nodes(1, elems(7, 1)) - nodes(1, elems(1, 1)) ! Element sizes
ly = nodes(2, elems(7, 1)) - nodes(2, elems(1, 1))
lz = nodes(3, elems(7, 1)) - nodes(3, elems(1, 1))
jacx = lx/2
jacy = ly/2
jacz = lz/2
jac_det = abs(jacx*jacy*jacz)
xp = (xiq + 1) * jacx
yp = (xiq + 1) * jacy
zp = (xiq + 1) * jacz
Ne = size(elems, 2)
do e = 1, Ne
    x = xp + nodes(1, elems(1, e))
    y = yp + nodes(2, elems(1, e))
    z = zp + nodes(3, elems(1, e))
    do iqz = 1, size(xiq)
    do iqy = 1, size(xiq)
    do iqx = 1, size(xiq)
        fq(iqx, iqy, iqz, e) = func(x(iqx), y(iqy), z(iqz))
    end do
    end do
    end do
end do
end function

subroutine assemble_3d(xin, nodes, elems, ib, xiq, wtq, phihq, dphihq, &
        rhsq, matBp, matBj, matBx, rhs, verbose)
! Assemble Poisson equation on a 3D hexahedral uniform mesh
! It solves:
!   \nabla^2 V(x, y, z) = -f(x, y, z)
! where f(x, y, z) is specified in "rhsq". Typically, one would use f = 4*pi*n,
! where "n" is the positive particle density.
real(dp), intent(in):: xin(:), nodes(:, :), xiq(:), wtq(:, :, :), &
    phihq(:, :), dphihq(:, :), rhsq(:, :, :, :)
integer, intent(in):: elems(:, :), ib(:, :, :, :)
real(dp), intent(out):: rhs(:)
integer, allocatable, intent(out) :: matBp(:), matBj(:)
real(dp), allocatable, intent(out) :: matBx(:)
logical, intent(in), optional :: verbose
integer :: Ne, p, Nq
real(dp), dimension(size(xiq), size(xiq), size(xiq), &
    size(xin), size(xin), size(xin)) :: phi_v
real(dp) :: lx, ly, lz
real(dp) :: jac_det
real(dp), dimension(size(xiq), size(xiq), size(xiq), &
    size(xiq), size(xiq), size(xiq)) :: Am_loc
logical :: verbose_
verbose_ = .false.
if (present(verbose)) verbose_ = verbose

Ne = size(elems, 2)
p = size(xin) - 1
Nq = size(xiq)
! Element sizes
lx = nodes(1, elems(7, 1)) - nodes(1, elems(1, 1))
ly = nodes(2, elems(7, 1)) - nodes(2, elems(1, 1))
lz = nodes(3, elems(7, 1)) - nodes(3, elems(1, 1))
call assemble_3d_precalc(p, Nq, lx, ly, lz, wtq, phihq, &
        dphihq, jac_det, Am_loc, phi_v)
if (verbose_) then
    print *, "Assembly..."
end if
call assemble_3d_csr(Ne, p, rhsq, jac_det, wtq, ib, Am_loc, phi_v, &
        matBp, matBj, matBx, rhs)
if (verbose_) then
    print *, "CSR Matrix:"
    print *, "    dimension:", size(matBp)-1
    print *, "    number of nonzeros:", size(matBx)
    print "('     density:', f7.2, '%')", size(matBx) * 100._dp / &
        (size(matBp)-1._dp)**2
end if
end subroutine

subroutine assemble_3d_precalc(p, Nq, lx, ly, lz, wtq, phihq, &
        dphihq, jac_det, Am_loc, phi_v)
integer, intent(in) :: p, Nq
real(dp), intent(in) :: lx, ly, lz
real(dp), intent(in):: wtq(:, :, :), phihq(:, :), dphihq(:, :)
integer :: iqx, iqy, iqz
real(dp), dimension(Nq, Nq, Nq, p+1, p+1, p+1) :: phi_dx, phi_dy, phi_dz
integer :: ax, ay, az, bx, by, bz
real(dp) :: jacx, jacy, jacz
real(dp), intent(out) :: jac_det
real(dp), intent(out), dimension(:, :, :, :, :, :) :: Am_loc, phi_v
! Precalculate basis functions:
print *, "Precalculate basis functions"
!$omp parallel default(none) shared(p, Nq, phihq, dphihq, phi_v, phi_dx, phi_dy, phi_dz) private(ax, ay, az, iqx, iqy, iqz)
!$omp do
do az = 1, p+1
do ay = 1, p+1
do ax = 1, p+1
    do iqz = 1, Nq
    do iqy = 1, Nq
    do iqx = 1, Nq
        phi_v (iqx, iqy, iqz, ax, ay, az) = &
             phihq(iqx, ax) *  phihq(iqy, ay) *  phihq(iqz, az)
        phi_dx(iqx, iqy, iqz, ax, ay, az) = &
            dphihq(iqx, ax) *  phihq(iqy, ay) *  phihq(iqz, az)
        phi_dy(iqx, iqy, iqz, ax, ay, az) = &
             phihq(iqx, ax) * dphihq(iqy, ay) *  phihq(iqz, az)
        phi_dz(iqx, iqy, iqz, ax, ay, az) = &
             phihq(iqx, ax) *  phihq(iqy, ay) * dphihq(iqz, az)
    end do
    end do
    end do
end do
end do
end do
!$omp end do
!$omp end parallel
jacx = lx/2
jacy = ly/2
jacz = lz/2
jac_det = abs(jacx*jacy*jacz)
phi_dx = phi_dx / jacx
phi_dy = phi_dy / jacy
phi_dz = phi_dz / jacz
! Precalculate element matrix:
print *, "Precalculate element matrix"
!$omp parallel default(none) shared(p, phi_dx, phi_dy, phi_dz, jac_det, wtq, Am_loc) private(ax, ay, az, bx, by, bz)
!$omp do
do bz = 1, p+1
do by = 1, p+1
do bx = 1, p+1
    do az = 1, p+1
    do ay = 1, p+1
    do ax = 1, p+1
        Am_loc(ax, ay, az, bx, by, bz) = sum(( &
            phi_dx(:, :, :, ax, ay, az)*phi_dx(:, :, :, bx, by, bz) + &
            phi_dy(:, :, :, ax, ay, az)*phi_dy(:, :, :, bx, by, bz) + &
            phi_dz(:, :, :, ax, ay, az)*phi_dz(:, :, :, bx, by, bz)) &
            * jac_det * wtq)
    end do
    end do
    end do
end do
end do
end do
!$omp end do
!$omp end parallel
end subroutine

subroutine local_overlap_matrix(wtq, jac_det, phi_v, Am_loc)
real(dp), intent(in):: wtq(:, :, :)
real(dp), intent(in) :: jac_det
real(dp), intent(in), dimension(:, :, :, :, :, :) :: phi_v
real(dp), intent(out), dimension(:, :, :, :, :, :) :: Am_loc
integer :: p, ax, ay, az, bx, by, bz
p = size(phi_v, 6)-1
! Precalculate element matrix:
!$omp parallel default(none) shared(p, jac_det, wtq, Am_loc, phi_v) private(ax, ay, az, bx, by, bz)
!$omp do
do bz = 1, p+1
do by = 1, p+1
do bx = 1, p+1
    do az = 1, p+1
    do ay = 1, p+1
    do ax = 1, p+1
        Am_loc(ax, ay, az, bx, by, bz) = sum( &
            phi_v(:, :, :, ax, ay, az) * phi_v(:, :, :, bx, by, bz) &
            * jac_det * wtq)
    end do
    end do
    end do
end do
end do
end do
!$omp end do
!$omp end parallel
end subroutine

subroutine assemble_3d_coo_A(Ne, p, ib, Am_loc, matAi, matAj, matAx, idx)
! Assembles the matrix A
! Assumes ib is never 0!
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: Am_loc(:, :, :, :, :, :)
! Allocate the following three arrays to Ne*(p+1)**6 elements
integer, intent(out) :: matAi(:), matAj(:)
real(dp), intent(out) :: matAx(:)
integer, intent(out) :: idx
integer :: e, i, j
integer :: ax, ay, az, bx, by, bz
call assert(all(ib > 0))
idx = 0
do e = 1, Ne
    do bz = 1, p+1
    do by = 1, p+1
    do bx = 1, p+1
        j = ib(bx, by, bz, e)
        do az = 1, p+1
        do ay = 1, p+1
        do ax = 1, p+1
            i = ib(ax, ay, az, e)
            if (j > i) cycle
            idx = idx + 1
            matAi(idx) = i
            matAj(idx) = j
            matAx(idx) = Am_loc(ax, ay, az, bx, by, bz)
            if (abs(matAx(idx)) < 1e-12_dp) then
                idx = idx - 1
                cycle
            end if
            if (i /= j) then
                ! Symmetric contribution
                idx = idx + 1
                matAi(idx) = j
                matAj(idx) = i
                matAx(idx) = Am_loc(ax, ay, az, bx, by, bz)
            end if
        end do
        end do
        end do
    end do
    end do
    end do
end do
!call assert(idx == Ne*(p+1)**6)
end subroutine

subroutine assemble_3d_coo_A_subdomain_spectral(Ne, p, ib, dphihq, &
        lx, ly, lz, wtq, elem_mask, matAi, matAj, matAx, idx, &
        jac_det, global_to_local)
! Assembles the matrix A, assuming spectral elements (Guass-Lobatto integration
! points and Nq=p+1).
! Assumes ib is never 0!
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: lx, ly, lz
real(dp), intent(in):: wtq(:), dphihq(:, :)
integer, intent(in) :: elem_mask(:)
integer, intent(out) :: matAi(:), matAj(:)
real(dp), intent(out) :: matAx(:)
integer, intent(out) :: idx
real(dp), intent(out) :: jac_det
integer, intent(out) :: global_to_local(:)
integer :: e, i, j, counter
integer :: ax, ay, az, bx, by, bz
real(dp) :: jacx, jacy, jacz
real(dp) :: t1, t2
real(dp) :: a_loc(p+1, p+1)
call assert(all(ib > 0))

call cpu_time(t1)

jacx = lx/2
jacy = ly/2
jacz = lz/2
jac_det = abs(jacx*jacy*jacz)

do bx = 1, p+1
do ax = 1, p+1
    a_loc(ax, bx) = sum(dphihq(:, ax) * dphihq(:, bx) * wtq) * jac_det
end do
end do

! Make sure there are no zeros in the local matrix
call assert(all(abs(a_loc) > 1e-12_dp))

idx = 0
global_to_local = 0
do e = 1, Ne
    if (elem_mask(e) == 1) then
        do az = 1, p+1
        do ay = 1, p+1
        do ax = 1, p+1
            i = ib(ax, ay, az, e)

            by = ay
            bz = az
            do bx = 1, p+1
                j = ib(bx, by, bz, e)
                if (j < i) cycle
                idx = idx + 1
                matAi(idx) = i
                matAj(idx) = j
                matAx(idx) = a_loc(ax, bx) / jacx**2 * wtq(ay)*wtq(az)
                global_to_local(i) = 1
            end do

            bx = ax
            bz = az
            do by = 1, p+1
                j = ib(bx, by, bz, e)
                if (j < i) cycle
                idx = idx + 1
                matAi(idx) = i
                matAj(idx) = j
                matAx(idx) = a_loc(ay, by) / jacy**2 * wtq(ax)*wtq(az)
                global_to_local(i) = 1
            end do

            bx = ax
            by = ay
            do bz = 1, p+1
                j = ib(bx, by, bz, e)
                if (j < i) cycle
                idx = idx + 1
                matAi(idx) = i
                matAj(idx) = j
                matAx(idx) = a_loc(az, bz) / jacz**2 * wtq(ax)*wtq(ay)
                global_to_local(i) = 1
            end do

        end do
        end do
        end do
    end if
end do
call cpu_time(t2)
print *, "Assembly time:", t2-t1, "s"
counter = 0
do i = 1, size(global_to_local)
    if (global_to_local(i) /= 0) then
        counter = counter + 1
        global_to_local(i) = counter
    end if
end do
end subroutine

subroutine assemble_3d_coo_A_spectral(Ne, p, ib, dphihq, lx, ly, lz, wtq, matAi, matAj, matAx, idx, jac_det)
! Assembles the matrix A, assuming spectral elements (Guass-Lobatto integration
! points and Nq=p+1).
! Assumes ib is never 0!
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: lx, ly, lz
real(dp), intent(in):: wtq(:), dphihq(:, :)
integer, intent(out) :: matAi(:), matAj(:)
real(dp), intent(out) :: matAx(:)
integer, intent(out) :: idx
real(dp), intent(out) :: jac_det
integer :: e, i, j
integer :: ax, ay, az, bx, by, bz
real(dp) :: jacx, jacy, jacz
real(dp) :: t1, t2
real(dp) :: a_loc(p+1, p+1)
call assert(all(ib > 0))

call cpu_time(t1)

jacx = lx/2
jacy = ly/2
jacz = lz/2
jac_det = abs(jacx*jacy*jacz)

do bx = 1, p+1
do ax = 1, p+1
    a_loc(ax, bx) = sum(dphihq(:, ax) * dphihq(:, bx) * wtq) * jac_det
end do
end do

! Make sure there are no zeros in the local matrix
call assert(all(abs(a_loc) > 1e-12_dp))

idx = 0
do e = 1, Ne
    do az = 1, p+1
    do ay = 1, p+1
    do ax = 1, p+1
        i = ib(ax, ay, az, e)

        by = ay
        bz = az
        do bx = 1, p+1
            j = ib(bx, by, bz, e)
            if (j > i) cycle
            idx = idx + 1
            matAi(idx) = i
            matAj(idx) = j
            matAx(idx) = a_loc(ax, bx) / jacx**2 * wtq(ay)*wtq(az)
            if (i /= j) then
                ! Symmetric contribution
                idx = idx + 1
                matAi(idx) = j
                matAj(idx) = i
                matAx(idx) = matAx(idx-1)
            end if
        end do

        bx = ax
        bz = az
        do by = 1, p+1
            j = ib(bx, by, bz, e)
            if (j > i) cycle
            idx = idx + 1
            matAi(idx) = i
            matAj(idx) = j
            matAx(idx) = a_loc(ay, by) / jacy**2 * wtq(ax)*wtq(az)
            if (i /= j) then
                ! Symmetric contribution
                idx = idx + 1
                matAi(idx) = j
                matAj(idx) = i
                matAx(idx) = matAx(idx-1)
            end if
        end do

        bx = ax
        by = ay
        do bz = 1, p+1
            j = ib(bx, by, bz, e)
            if (j > i) cycle
            idx = idx + 1
            matAi(idx) = i
            matAj(idx) = j
            matAx(idx) = a_loc(az, bz) / jacz**2 * wtq(ax)*wtq(ay)
            if (i /= j) then
                ! Symmetric contribution
                idx = idx + 1
                matAi(idx) = j
                matAj(idx) = i
                matAx(idx) = matAx(idx-1)
            end if
        end do

    end do
    end do
    end do
end do
call cpu_time(t2)
print *, "Assembly time:", t2-t1, "s"
end subroutine

subroutine assemble_3d_coo(Ne, p, rhsq, jac_det, wtq, ib, Am_loc, phi_v, &
        matAi, matAj, matAx, rhs, idx)
! The actual, low level assembly
real(dp), intent(in):: wtq(:, :, :), rhsq(:, :, :, :)
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: phi_v(:, :, :, :, :, :)
real(dp), intent(in) :: jac_det
real(dp), intent(in) :: Am_loc(:, :, :, :, :, :)
integer, intent(out) :: matAi(:), matAj(:)
real(dp), intent(out) :: matAx(:)
real(dp), intent(out):: rhs(:)
integer, intent(out) :: idx
real(dp), dimension(size(wtq, 1), size(wtq, 2), size(wtq, 3)) :: fq
integer :: e, i, j
integer :: ax, ay, az, bx, by, bz
rhs = 0
idx = 0
do e = 1, Ne
    fq = rhsq(:, :, :, e) * jac_det * wtq
    do bz = 1, p+1
    do by = 1, p+1
    do bx = 1, p+1
        j = ib(bx, by, bz, e)
        if (j==0) cycle
        do az = 1, p+1
        do ay = 1, p+1
        do ax = 1, p+1
            i = ib(ax, ay, az, e)
            if (i == 0) cycle
            if (j > i) cycle
            idx = idx + 1
            matAi(idx) = i
            matAj(idx) = j
            matAx(idx) = Am_loc(ax, ay, az, bx, by, bz)
            if (i /= j) then
                ! Symmetric contribution
                idx = idx + 1
                matAi(idx) = j
                matAj(idx) = i
                matAx(idx) = Am_loc(ax, ay, az, bx, by, bz)
            end if
        end do
        end do
        end do
        rhs(j) = rhs(j) + sum(phi_v(:, :, :, bx, by, bz) * fq)
    end do
    end do
    end do
end do
end subroutine

subroutine assemble_3d_coo_A_subdomain(Ne, p, rhsq, jac_det, wtq, ib, Am_loc, &
        elem_mask, matAi, matAj, matAx, idx, global_to_local)
! The actual, low level assembly
real(dp), intent(in):: wtq(:, :, :), rhsq(:, :, :, :)
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: jac_det
real(dp), intent(in) :: Am_loc(:, :, :, :, :, :)
integer, intent(in) :: elem_mask(:)
integer, intent(out) :: matAi(:), matAj(:)
real(dp), intent(out) :: matAx(:)
integer, intent(out) :: idx
integer, intent(out) :: global_to_local(:)
real(dp), dimension(size(wtq, 1), size(wtq, 2), size(wtq, 3)) :: fq
integer :: e, i, j, counter
integer :: ax, ay, az, bx, by, bz
idx = 0
global_to_local = 0
do e = 1, Ne
    if (elem_mask(e) == 1) then
        fq = rhsq(:, :, :, e) * jac_det * wtq
        do bz = 1, p+1
        do by = 1, p+1
        do bx = 1, p+1
            j = ib(bx, by, bz, e)
            call assert(j /= 0)
            do az = 1, p+1
            do ay = 1, p+1
            do ax = 1, p+1
                i = ib(ax, ay, az, e)
                if (i == 0) cycle
                call assert(i /= 0)
                if (j < i) cycle
                idx = idx + 1
                matAi(idx) = i
                matAj(idx) = j
                matAx(idx) = Am_loc(ax, ay, az, bx, by, bz)
                global_to_local(i) = 1
            end do
            end do
            end do
        end do
        end do
        end do
    end if
end do
counter = 0
do i = 1, size(global_to_local)
    if (global_to_local(i) /= 0) then
        counter = counter + 1
        global_to_local(i) = counter
    end if
end do
end subroutine

subroutine assemble_3d_coo_rhs(Ne, p, rhsq, jac_det, wtq, ib, phi_v, rhs)
real(dp), intent(in):: wtq(:, :, :), rhsq(:, :, :, :)
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: phi_v(:, :, :, :, :, :)
real(dp), intent(in) :: jac_det
real(dp), intent(out):: rhs(:)
real(dp), dimension(size(wtq, 1), size(wtq, 2), size(wtq, 3)) :: fq
integer :: e, j
integer :: bx, by, bz
rhs = 0
do e = 1, Ne
    fq = rhsq(:, :, :, e) * jac_det * wtq
    do bz = 1, p+1
    do by = 1, p+1
    do bx = 1, p+1
        j = ib(bx, by, bz, e)
        ! For periodic BC, 'j' is always nonzero. Uncomment the following line
        ! to check that:
        ! call assert(j /= 0)
        rhs(j) = rhs(j) + sum(phi_v(:, :, :, bx, by, bz) * fq)
    end do
    end do
    end do
end do
end subroutine

subroutine assemble_3d_coo_rhs_spectral(Ne, p, rhsq, jac_det, wtq, ib, rhs)
real(dp), intent(in):: wtq(:, :, :), rhsq(:, :, :, :)
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: jac_det
real(dp), intent(out):: rhs(:)
real(dp), dimension(size(wtq, 1), size(wtq, 2), size(wtq, 3)) :: fq
integer :: e, j
integer :: bx, by, bz
rhs = 0
do e = 1, Ne
    fq = rhsq(:, :, :, e) * jac_det * wtq
    do bz = 1, p+1
    do by = 1, p+1
    do bx = 1, p+1
        j = ib(bx, by, bz, e)
        ! For periodic BC, 'j' is always nonzero. Uncomment the following line
        ! to check that:
        ! call assert(j /= 0)
        rhs(j) = rhs(j) + fq(bx, by, bz)
    end do
    end do
    end do
end do
end subroutine

subroutine assemble_3d_csr(Ne, p, rhsq, jac_det, wtq, ib, Am_loc, phi_v, &
        matBp, matBj, matBx, rhs)
real(dp), intent(in):: wtq(:, :, :), rhsq(:, :, :, :)
integer, intent(in):: ib(:, :, :, :)
integer, intent(in) :: Ne, p
real(dp), intent(in) :: phi_v(:, :, :, :, :, :)
real(dp), intent(in) :: jac_det
real(dp), intent(in) :: Am_loc(:, :, :, :, :, :)
integer, intent(out), allocatable :: matBp(:), matBj(:)
real(dp), intent(out), allocatable :: matBx(:)
real(dp), intent(out):: rhs(:)
! Maximum possible size:
integer, dimension(Ne*(p+1)**6) :: matAi, matAj
real(dp),  dimension(Ne*(p+1)**6) :: matAx
! Actual size:
integer :: idx
call assemble_3d_coo(Ne, p, rhsq, jac_det, wtq, ib, Am_loc, phi_v, &
        matAi, matAj, matAx, rhs, idx)
call coo2csr_canonical(matAi(:idx), matAj(:idx), matAx(:idx), &
    matBp, matBj, matBx)
end subroutine

real(dp) function integral(nodes, elems, wtq, fq) result(r)
real(dp), intent(in) :: nodes(:, :)
integer, intent(in) :: elems(:, :)
real(dp), intent(in) :: wtq(:, :, :), fq(:, :, :, :)
real(dp) :: jacx, jacy, jacz, l(3), jac_det
integer :: e, Ne
Ne = size(elems, 2)
r = 0
do e = 1, Ne
    ! l is the diagonal vector:
    l = nodes(:, elems(7, e)) - nodes(:, elems(1, e))
    ! Assume rectangular shape:
    jacx = l(1)/2
    jacy = l(2)/2
    jacz = l(3)/2
    jac_det = abs(jacx*jacy*jacz)
    r = r + sum(fq(:, :, :, e) * jac_det * wtq)
end do
end function

end module


program display_demo
    implicit none
    integer, parameter :: max_pow = 10
    integer :: i
    real :: total

    total = 0.0
    do i = 0, max_pow
        total = total + 10.0**i
    end do

    print '(ES15.3)', total

end program display_demo


      PROGRAM xpzextr
C     driver for routine pzextr
      INTEGER NV
      PARAMETER(NV=4)
      INTEGER i,iest,j
      REAL dum,xest,yest(NV),yz(NV),dy(NV)
      do 12 i=1,10
        iest=i
        xest=1.0/float(i)
        dum=1.0-xest+xest*xest*xest
        do 11 j=1,NV
          dum=dum/(xest+1.0)
          yest(j)=dum
11      continue
        call pzextr(iest,xest,yest,yz,dy,NV)
        write(*,'(/1x,a,i2)') 'I = ',i
        write(*,'(1x,a,4f12.6)') 'Extrap. function:',(yz(j),j=1,NV)
        write(*,'(1x,a,4f12.6)') 'Estimated error: ',(dy(j),j=1,NV)
12    continue
      write(*,'(/1x,a,4f12.6)') 'Actual values:   ',1.0,1.0,1.0,1.0
      END

NR-Functions/Numerical Recipes- Example & Functions/Examples/xpzextr.for


MODULE PRECISION
    IMPLICIT NONE

    INTRINSIC KIND

    ! standard IEEE arithmetic
    !INTEGER, PARAMETER :: skind = SELECT_REAL_KIND(p=6, r=37)
    !INTEGER, PARAMETER :: dkind = SELECT_REAL_KIND(p=15, r=307)

    ! standard precision set by compiler
    INTEGER, PARAMETER :: skind = KIND(0.0E0)
    INTEGER, PARAMETER :: dkind = KIND(0.0D0)

    ! Default precision of module
    INTEGER, PARAMETER :: wp = dkind
END MODULE


! { dg-do compile }
! { dg-require-effective-target vect_double }

MODULE UPML_mod

IMPLICIT NONE

PUBLIC UPMLupdateE

PRIVATE

real(kind=8), dimension(:,:,:), allocatable :: Dx_ilow

real(kind=8), dimension(:), allocatable :: aye, aze
real(kind=8), dimension(:), allocatable :: bye, bze
real(kind=8), dimension(:), allocatable :: fxh, cxh

real(kind=8) :: epsinv
real(kind=8) :: dxinv, dyinv, dzinv

integer :: xstart, ystart, zstart, xstop, ystop, zstop

CONTAINS

SUBROUTINE UPMLupdateE(nx,ny,nz,Hx,Hy,Hz,Ex,Ey,Ez)

integer, intent(in) :: nx, ny, nz
real(kind=8), intent(inout),                                                &
         dimension(xstart:xstop+1,ystart:ystop+1,zstart:zstop+1) :: Ex, Ey, Ez
real(kind=8), intent(inout),                                                &
         allocatable :: Hx(:,:,:), Hy(:,:,:), Hz(:,:,:)

integer :: i, j, k
real(kind=8) :: Dxold, Dyold, Dzold

do k=zstart+1,zstop
  do j=ystart+1,ystop
    do i=xstart+1,0

      Dxold = Dx_ilow(i,j,k)

      Dx_ilow(i,j,k) = aye(j) * Dx_ilow(i,j,k) +                              &
                       bye(j) * ((Hz(i,j,k  )-Hz(i,j-1,k))*dyinv +            &
                                 (Hy(i,j,k-1)-Hy(i,j,k  ))*dzinv)

      Ex(i,j,k) = aze(k) * Ex(i,j,k) +                                        &
                  bze(k) * (cxh(i)*Dx_ilow(i,j,k) - fxh(i)*Dxold) * epsinv
    end do
  end do
end do

END SUBROUTINE UPMLupdateE

END MODULE UPML_mod

! { dg-final { scan-tree-dump-times "vectorized 1 loops" 1 "vect"  } }
! { dg-final { cleanup-tree-dump "vect" } }
! { dg-final { cleanup-modules "upml_mod" } }


module generic_sum_m(TA, TB)
    implicit none
    private
    public :: gen_sum_req, generic_sum

    abstract interface gen_sum_req
        type(TB) operator(+)(a, b)
            type(TA), intent(in) :: a, b
        end operator(+)

        type(TB) operator(+)(a, b)
            type(TB), intent(in) :: a, b
        end operator(+)

        type(TB) function zero(a)
            type(TB), intent(in) :: a
        end subroutine

        assignment(=)(b, a)
            type(TB), intent(inout) :: b
            type(TA), intent(in) :: a
        end assignment(=)
    end interface

    interface generic_sum
        module procedure generic_sum_impl
    end interface
contains
    function generic_sum_impl(a) result(total)
        type(TA), intent(in) :: a(:)
        type(TB) :: total

        if (size(a) == 0) then
            total = zero(total)
        else if (size(a) == 1) then
            total = a(1)
        else if (size(a) == 2) then
            total = a(1) + a(2)
        else
            total = generic_sum(a( : size(a)/2)) + generic_sum(a(size(a)/2+1 : ))
        end if
    end function
end module

module generic_matmul_m(TU, TV, TW, TX, TY, TN, TP)
    use generic_sum_m(TX, TN), only: gen_sum_req_tx_tn => gen_sum_req, generic_sum
    use generic_sum_m(TY, TP), only: gen_sum_req_ty_tp => gen_sum_req, generic_sum

    implicit none
    private
    public :: gen_mm_req, generic_matmul

    abstract interface gen_mm_req
        elemental type(TX) operator(*)(u, v)
            type(TU), intent(in) :: u
            type(TV), intent(in) :: v
        end operator(*)

        elemental operator(*)(v, u) result(tv_times_tu)
            type(TV), intent(in) :: v
            type(TU), intent(in) :: u
            type(TY) :: tv_times_tu
        end operator(*)

        type(TZ) operator(+)(n, p)
            type(TN), intent(in) :: n
            type(TP), intent(in) :: p
        end operator(+)

        assignment(=)(w, z)
            type(TW), intent(inout) :: w
            type(TZ), intent(in) :: z
        end assignment(=)

        gen_sum_req_tx_tn
        gen_sum_req_ty_tp
    end interface
contains
    function generic_matmul(u, v) result(w)
        type(TU), intent(in) :: u(:, :)
        type(TV), intent(in) :: v(:, :)
        type(TW) :: w(size(u, 1), size(v, 2))

        call assert(size(u, 2) == size(v, 1))

        do concurrent(i=1:size(u, 1), j=1:size(v,2))
            w(i,j) = generic_sum(u(i,:) * v(:,j)) + generic_sum(v(:,j) * u(i,:))
        end do
    end function
end module


      PROGRAM LOOPBREAK
        INTEGER I, RNDINT

C       It doesn't matter what number you put here.
        CALL SDRAND(123)

C       Because FORTRAN 77 semantically lacks many loop structures, we
C       have to use GOTO statements to do the same thing.
   10   CONTINUE
C         Print a random number.
          I = RNDINT(0, 19)
          WRITE (*,*) I

C         If the random number is ten, break (i.e. skip to after the end
C         of the "loop").
          IF (I .EQ. 10) GOTO 20

C         Otherwise, print a second random number.
          I = RNDINT(0, 19)
          WRITE (*,*) I

C         This is the end of our "loop," meaning we jump back to the
C         beginning again.
          GOTO 10

   20   CONTINUE

        STOP
      END

C FORTRAN 77 does not have come with a random number generator, but it
C is easy enough to type "fortran 77 random number generator" into your
C preferred search engine and to copy and paste what you find. The
C following code is a slightly-modified version of:
C
C     http://www.tat.physik.uni-tuebingen.de/
C         ~kley/lehre/ftn77/tutorial/subprograms.html
      SUBROUTINE SDRAND (IRSEED)
        COMMON  /SEED/ UTSEED, IRFRST
        UTSEED = IRSEED
        IRFRST = 0
        RETURN
      END
      INTEGER FUNCTION RNDINT (IFROM, ITO)
        INTEGER IFROM, ITO
        PARAMETER (MPLIER=16807, MODLUS=2147483647,                     &
     &              MOBYMP=127773, MOMDMP=2836)
        COMMON  /SEED/ UTSEED, IRFRST
        INTEGER HVLUE, LVLUE, TESTV, NEXTN
        SAVE    NEXTN
        IF (IRFRST .EQ. 0) THEN
          NEXTN = UTSEED
          IRFRST = 1
        ENDIF
        HVLUE = NEXTN / MOBYMP
        LVLUE = MOD(NEXTN, MOBYMP)
        TESTV = MPLIER*LVLUE - MOMDMP*HVLUE
        IF (TESTV .GT. 0) THEN
          NEXTN = TESTV
        ELSE
          NEXTN = TESTV + MODLUS
        ENDIF
        IF (NEXTN .GE. 0) THEN
          RNDINT = MOD(MOD(NEXTN, MODLUS), ITO - IFROM + 1) + IFROM
        ELSE
          RNDINT = MOD(MOD(NEXTN, MODLUS), ITO - IFROM + 1) + ITO + 1
        ENDIF
        RETURN
      END


MODULE projection
!!-----------------------------------------------------------
!!
!!          to make a polar stereographic projection
!!              in the area of the north pole
!!
!!	         Created by Brice Lemaire on 05/2010.
!!
!!-----------------------------------------------------------
  USE readwrite
  !
  IMPLICIT NONE
  PUBLIC
  !
  REAL*8,DIMENSION(4) :: dmixlam, dmixphi
  REAL*8 :: dray = 6357          !rayon polaire de la Terre (km)
  REAL*8 :: PI = ACOS(-1.0)
  REAL*8 :: dlong0 = 0.          !longitude du centre de projection (origine)
  REAL*8 :: dlat0 = 90.          !latitude       ''
  INTEGER :: idisc = 0           !to check the +/- 180 discontinuity
  INTEGER :: m
  !
  CONTAINS
  !********************************************************
  !            SUBROUTINE stereo_projection      	    *
  !									        			*
  !												        *
  !		          CALLED by cfg_tools.F90	            *
  !********************************************************
  SUBROUTINE stereo_projection(kji,kjj,kjk,knorth_pole,kway)
    !
    INTEGER,INTENT(IN) :: kji,kjj,kjk
	LOGICAL,INTENT(IN) :: knorth_pole
	INTEGER,INTENT(IN) :: kway
	INTEGER :: klon, klat
	REAL*8,DIMENSION(4) :: dlx, dly
	REAL*8,DIMENSION(4) :: dllam, dlphi, dlk
	REAL*8 :: dl_lat0, dl_long0
	!
    dl_lat0 = dlat0 * PI/180.
    dl_long0 = dlong0 * PI/180.
    !
	!Either we interpolate along longitude or along latitude
	IF(kway.EQ.1)THEN
	  klon = 1
	  klat = 0
	ELSEIF(kway.EQ.2) THEN
	  klon = 0
	  klat = 1
	ENDIF
	!
	! Check the discontinuity +/- 180
	IF(kway.EQ.1)THEN
	  IF(smixgrd%glam(kji,kjj).LT.180..AND.smixgrd%glam(kji+3*nn_rhox,kjj).GT.180.) THEN
	    IF(smixgrd%glam(kji+2*nn_rhox,kjj).LT.180.) THEN
	      idisc = 1
	    ELSEIF(smixgrd%glam(kji+1*nn_rhox,kjj).GT.180.) THEN
	      idisc = 2
        ELSEIF(smixgrd%glam(kji+2*nn_rhox,kjj).GT.180.AND.smixgrd%glam(kji+1*nn_rhox,kjj).LT.180.)THEN
	      idisc = 3
	    ENDIF
	  ELSE
	  idisc = 0
	  ENDIF
	ENDIF
	!
	dllam(1) = smixgrd%glam(kji,kjj) * PI/180.
	dlphi(1) = smixgrd%gphi(kji,kjj) * PI/180.
	dlk(1)   = (2*dray) / (1 + (SIN(dl_lat0)*SIN(dlphi(1))) + (COS(dl_lat0)*COS(dlphi(1))*COS(dllam(1) - dl_long0)))
	!
	dllam(2) = smixgrd%glam(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat) * PI/180.
	dlphi(2) = smixgrd%gphi(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat) * PI/180.
	dlk(2)   = (2*dray) / (1 + (SIN(dl_lat0)*SIN(dlphi(2))) + (COS(dl_lat0)*COS(dlphi(2))*COS(dllam(2) - dl_long0)))
	!
	dllam(3) = smixgrd%glam(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat) * PI/180.
	dlphi(3) = smixgrd%gphi(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat) * PI/180.
	dlk(3)   = (2*dray) / (1 + (SIN(dl_lat0)*SIN(dlphi(3))) + (COS(dl_lat0)*COS(dlphi(3))*COS(dllam(3) - dl_long0)))
	!
	dllam(4) = smixgrd%glam(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat) * PI/180.
	dlphi(4) = smixgrd%gphi(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat) * PI/180.
	dlk(4)   = (2*dray) / (1 + (SIN(dl_lat0)*SIN(dlphi(4))) + (COS(dl_lat0)*COS(dlphi(4))*COS(dllam(4) - dl_long0)))
	!
	dlx(1) = dlk(1) * COS(dlphi(1)) * SIN(dllam(1) - dl_long0)
	dly(1) = dlk(1) * ((COS(dl_lat0) * SIN(dlphi(1))) - (SIN(dl_lat0) * COS(dlphi(1)) * COS(dllam(1) - dl_long0)))
    !
	dlx(2) = dlk(2) * COS(dlphi(2)) * SIN(dllam(2) - dl_long0)
	dly(2) = dlk(2) * ((COS(dl_lat0) * SIN(dlphi(2))) - (SIN(dl_lat0) * COS(dlphi(2)) * COS(dllam(2) - dl_long0)))
    !
	dlx(3) = dlk(3) * COS(dlphi(3)) * SIN(dllam(3) - dl_long0)
	dly(3) = dlk(3) * ((COS(dl_lat0) * SIN(dlphi(3))) - (SIN(dl_lat0) * COS(dlphi(3)) * COS(dllam(3) - dl_long0)))
    !
	dlx(4) = dlk(4) * COS(dlphi(4)) * SIN(dllam(4) - dl_long0)
	dly(4) = dlk(4) * ((COS(dl_lat0) * SIN(dlphi(4))) - (SIN(dl_lat0) * COS(dlphi(4)) * COS(dllam(4) - dl_long0)))
    !
	dmixlam(1) = smixgrd%glam(kji,kjj)
	dmixlam(2) = smixgrd%glam(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat)
    dmixlam(3) = smixgrd%glam(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat)
    dmixlam(4) = smixgrd%glam(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat)
	!
	dmixphi(1) = smixgrd%gphi(kji,kjj)
	dmixphi(2) = smixgrd%gphi(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat)
    dmixphi(3) = smixgrd%gphi(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat)
    dmixphi(4) = smixgrd%gphi(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat)
	!
	smixgrd%glam(kji,kjj)                               = dlx(1)
	smixgrd%glam(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat) = dlx(2)
	smixgrd%glam(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat) = dlx(3)
	smixgrd%glam(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat) = dlx(4)
	!
	smixgrd%gphi(kji,kjj)                               = dly(1)
	smixgrd%gphi(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat) = dly(2)
    smixgrd%gphi(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat) = dly(3)
    smixgrd%gphi(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat) = dly(4)
	!
  END SUBROUTINE stereo_projection
  !
  !
  !
  !********************************************************
  !          SUBROUTINE stereo_projection_inv    	    *
  !									        			*
  !												        *
  !		            CALLED from here	                *
  !********************************************************
  SUBROUTINE stereo_projection_inv(kji,kjj,kjk,knorth_pole,kway)
    !
    INTEGER,INTENT(IN) :: kji,kjj,kjk
	LOGICAL,INTENT(IN) :: knorth_pole
	INTEGER,INTENT(IN) :: kway
	INTEGER :: klon, klat
	REAL*8,DIMENSION(5) :: dlx, dly
	REAL*8,DIMENSION(5)  :: dllam, dlphi
    REAL*8,DIMENSION(5)  :: dlro, dlc
	REAL*8 :: dl_long0, dl_lat0
	!
    dl_lat0 = dlat0 * PI/180.
    dl_long0 = dlong0 * PI/180.
    !
 	IF(kway.EQ.1)THEN
	  klon = 1
	  klat = 0
	ELSEIF(kway.EQ.2) THEN
	  klon = 0
	  klat = 1
	ENDIF
	!
	!
	dlx(1) = smixgrd%glam(kji,kjj)
	dlx(2) = smixgrd%glam(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat)
	dlx(3) = smixgrd%glam(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat)
	dlx(4) = smixgrd%glam(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat)
	dlx(5) = smixgrd%glam(kji+(nn_rhox+kjk)*klon,kjj+(nn_rhoy+kjk)*klat)
	!
	dly(1) = smixgrd%gphi(kji,kjj)
	dly(2) = smixgrd%gphi(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat)
	dly(3) = smixgrd%gphi(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat)
	dly(4) = smixgrd%gphi(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat)
	dly(5) = smixgrd%gphi(kji+(nn_rhox+kjk)*klon,kjj+(nn_rhoy+kjk)*klat)
	!
	dlro(1) = SQRT(dlx(1)*dlx(1) + dly(1)*dly(1))
	dlro(2) = SQRT(dlx(2)*dlx(2) + dly(2)*dly(2))
	dlro(3) = SQRT(dlx(3)*dlx(3) + dly(3)*dly(3))
	dlro(4) = SQRT(dlx(4)*dlx(4) + dly(4)*dly(4))
	dlro(5) = SQRT(dlx(5)*dlx(5) + dly(5)*dly(5))
	!
	dlc(1) = 2*ATAN(dlro(1)/(2*dray))
	dlc(2) = 2*ATAN(dlro(2)/(2*dray))
	dlc(3) = 2*ATAN(dlro(3)/(2*dray))
	dlc(4) = 2*ATAN(dlro(4)/(2*dray))
	dlc(5) = 2*ATAN(dlro(5)/(2*dray))
	!
	dlphi(1) =  ASIN(COS(dlc(1))*SIN(dl_lat0) + (dly(1)*SIN(dlc(1))*COS(dl_lat0) / dlro(1)))
	dlphi(2) =  ASIN(COS(dlc(2))*SIN(dl_lat0) + (dly(2)*SIN(dlc(2))*COS(dl_lat0) / dlro(2)))
	dlphi(3) =  ASIN(COS(dlc(3))*SIN(dl_lat0) + (dly(3)*SIN(dlc(3))*COS(dl_lat0) / dlro(3)))
	dlphi(4) =  ASIN(COS(dlc(4))*SIN(dl_lat0) + (dly(4)*SIN(dlc(4))*COS(dl_lat0) / dlro(4)))
	dlphi(5) =  ASIN(COS(dlc(5))*SIN(dl_lat0) + (dly(5)*SIN(dlc(5))*COS(dl_lat0) / dlro(5)))
    !
	dlphi(:) = dlphi(:) * 180./PI
    !
	dllam(1) = dl_long0 + ATAN(dlx(1)*SIN(dlc(1) / ((dlro(1)*COS(dl_lat0)*COS(dlc(1)))-(dly(1)*SIN(dl_lat0)*SIN(dlc(1))))))
	dllam(2) = dl_long0 + ATAN(dlx(2)*SIN(dlc(2) / ((dlro(2)*COS(dl_lat0)*COS(dlc(2)))-(dly(2)*SIN(dl_lat0)*SIN(dlc(2))))))
	dllam(3) = dl_long0 + ATAN(dlx(3)*SIN(dlc(3) / ((dlro(3)*COS(dl_lat0)*COS(dlc(3)))-(dly(3)*SIN(dl_lat0)*SIN(dlc(3))))))
	dllam(4) = dl_long0 + ATAN(dlx(4)*SIN(dlc(4) / ((dlro(4)*COS(dl_lat0)*COS(dlc(4)))-(dly(4)*SIN(dl_lat0)*SIN(dlc(4))))))
	dllam(5) = dl_long0 + ATAN(dlx(5)*SIN(dlc(5) / ((dlro(5)*COS(dl_lat0)*COS(dlc(5)))-(dly(5)*SIN(dl_lat0)*SIN(dlc(5))))))
    !
	dllam(1) = dllam(1) * 180./PI
	dllam(2) = dllam(2) * 180./PI
	dllam(3) = dllam(3) * 180./PI
	dllam(4) = dllam(4) * 180./PI
    !
	dllam(5) = dllam(5) * 180./PI
	!
	IF(kway.EQ.1)THEN
	  IF(idisc.EQ.1) THEN
		dllam(5) = dllam(5) + 180.
	  ELSEIF(idisc.EQ.2) THEN
		dllam(5) = dllam(5) - 180.
	  ELSEIF(idisc.EQ.3) THEN
		 dllam(5) = dllam(5) - 180.
	  ELSE
		dllam(5) = dllam(5)
	  ENDIF
	  !
	ELSEIF(kway.EQ.2)THEN
	  IF(smixgrd%glam(kji,kjj+1*nn_rhoy).GT.0.AND.smixgrd%glam(kji,kjj+2*nn_rhoy).GT.0)THEN
	    IF(dllam(5).LT.0.)THEN
	      dllam(5) = dllam(5) + 180
		ENDIF
	  ELSEIF(smixgrd%glam(kji,kjj+1*nn_rhoy).LT.0.AND.smixgrd%glam(kji,kjj+2*nn_rhoy).LT.0)THEN
	    IF(dllam(5).GT.0.)THEN
		  dllam(5) = dllam(5) - 180
	    ENDIF
	  ELSEIF(smixgrd%glam(kji,kjj+1*nn_rhoy)*smixgrd%glam(kji,kjj+2*nn_rhoy).LT.0)THEN
	    IF(dllam(5).LT.0.)THEN
	      dllam(5) = dllam(5) + 180
	    ELSEIF(dllam(5).GT.0.)THEN
	      dllam(5) = dllam(5) - 180
	    ENDIF
	  ENDIF
	ENDIF
	!
	smixgrd%glam(kji,kjj)                                       = dmixlam(1)
	smixgrd%glam(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat)         = dmixlam(2)
    smixgrd%glam(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat)         = dmixlam(3)
    smixgrd%glam(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat)         = dmixlam(4)
	!
    smixgrd%glam(kji+(nn_rhox+kjk)*klon,kjj+(nn_rhoy+kjk)*klat) = dllam(5)
	!
	smixgrd%gphi(kji,kjj)                                       = dmixphi(1)
	smixgrd%gphi(kji+1*nn_rhox*klon,kjj+1*nn_rhoy*klat)         = dmixphi(2)
    smixgrd%gphi(kji+2*nn_rhox*klon,kjj+2*nn_rhoy*klat)         = dmixphi(3)
    smixgrd%gphi(kji+3*nn_rhox*klon,kjj+3*nn_rhoy*klat)         = dmixphi(4)
	!
    smixgrd%gphi(kji+(nn_rhox+kjk)*klon,kjj+(nn_rhoy+kjk)*klat) = dlphi(5)
	!
  END SUBROUTINE stereo_projection_inv
  !
  !
  !
END MODULE projection


      SUBROUTINE MD03BA( N, IPAR, LIPAR, FNORM, J, LDJ, E, JNORMS,
     $                   GNORM, IPVT, DWORK, LDWORK, INFO )
C
C     SLICOT RELEASE 5.7.
C
C     Copyright (c) 2002-2020 NICONET e.V.
C
C     PURPOSE
C
C     To compute the QR factorization with column pivoting of an
C     m-by-n Jacobian matrix J (m >= n), that is, J*P = Q*R, where Q is
C     a matrix with orthogonal columns, P a permutation matrix, and
C     R an upper trapezoidal matrix with diagonal elements of
C     nonincreasing magnitude, and to apply the transformation Q' on
C     the error vector e (in-situ). The 1-norm of the scaled gradient
C     is also returned.
C
C     This routine is an interface to SLICOT Library routine MD03BX,
C     for solving standard nonlinear least squares problems using SLICOT
C     routine MD03BD.
C
C     ARGUMENTS
C
C     Input/Output Parameters
C
C     N       (input) INTEGER
C             The number of columns of the Jacobian matrix J.  N >= 0.
C
C     IPAR    (input) INTEGER array, dimension (LIPAR)
C             The integer parameters describing the structure of the
C             matrix J, as follows:
C             IPAR(1) must contain the number of rows M of the Jacobian
C                     matrix J.  M >= N.
C             IPAR is provided for compatibility with SLICOT Library
C             routine MD03BD.
C
C     LIPAR   (input) INTEGER
C             The length of the array IPAR.  LIPAR >= 1.
C
C     FNORM   (input) DOUBLE PRECISION
C             The Euclidean norm of the vector e.  FNORM >= 0.
C
C     J       (input/output) DOUBLE PRECISION array, dimension (LDJ, N)
C             On entry, the leading M-by-N part of this array must
C             contain the Jacobian matrix J.
C             On exit, the leading N-by-N upper triangular part of this
C             array contains the upper triangular factor R of the
C             Jacobian matrix. Note that for efficiency of the later
C             calculations, the matrix R is delivered with the leading
C             dimension MAX(1,N), possibly much smaller than the value
C             of LDJ on entry.
C
C     LDJ     (input/output) INTEGER
C             The leading dimension of array J.
C             On entry, LDJ >= MAX(1,M).
C             On exit,  LDJ >= MAX(1,N).
C
C     E       (input/output) DOUBLE PRECISION array, dimension (M)
C             On entry, this array must contain the error vector e.
C             On exit, this array contains the updated vector Q'*e.
C
C     JNORMS  (output) DOUBLE PRECISION array, dimension (N)
C             This array contains the Euclidean norms of the columns
C             of the Jacobian matrix, considered in the initial order.
C
C     GNORM   (output) DOUBLE PRECISION
C             If FNORM > 0, the 1-norm of the scaled vector
C             J'*Q'*e/FNORM, with each element i further divided
C             by JNORMS(i) (if JNORMS(i) is nonzero).
C             If FNORM = 0, the returned value of GNORM is 0.
C
C     IPVT    (output) INTEGER array, dimension (N)
C             This array defines the permutation matrix P such that
C             J*P = Q*R. Column j of P is column IPVT(j) of the identity
C             matrix.
C
C     Workspace
C
C     DWORK   DOUBLE PRECISION array, dimension (LDWORK)
C             On exit, if INFO = 0, DWORK(1) returns the optimal value
C             of LDWORK.
C
C     LDWORK  INTEGER
C             The length of the array DWORK.
C             LDWORK >= 1,      if N = 0 or  M = 1;
C             LDWORK >= 4*N+1,  if N > 1.
C             For optimum performance LDWORK should be larger.
C
C     Error Indicator
C
C     INFO    INTEGER
C             = 0:  successful exit;
C             < 0:  if INFO = -i, the i-th argument had an illegal
C                   value.
C
C     METHOD
C
C     This routine calls SLICOT Library routine MD03BX to perform the
C     calculations.
C
C     FURTHER COMMENTS
C
C     For efficiency, the arguments are not checked. This is done in
C     the routine MD03BX (except for LIPAR).
C
C     CONTRIBUTORS
C
C     V. Sima, Research Institute for Informatics, Bucharest, Dec. 2001.
C
C     REVISIONS
C
C     -
C
C     KEYWORDS
C
C     Elementary matrix operations, Jacobian matrix, matrix algebra,
C     matrix operations.
C
C     ******************************************************************
C
C     .. Scalar Arguments ..
      INTEGER           INFO, LDJ, LDWORK, LIPAR, N
      DOUBLE PRECISION  FNORM, GNORM
C     .. Array Arguments ..
      INTEGER           IPAR(*), IPVT(*)
      DOUBLE PRECISION  DWORK(*), E(*), J(*), JNORMS(*)
C     .. External Subroutines ..
      EXTERNAL          MD03BX
C     ..
C     .. Executable Statements ..
C
      CALL MD03BX( IPAR(1), N, FNORM, J, LDJ, E, JNORMS, GNORM, IPVT,
     $             DWORK, LDWORK, INFO )
      RETURN
C
C *** Last line of MD03BA ***
      END