WENO.F90

! Copyright (C) 2019 Timothy A. Smith
!
! Permission is hereby granted, free of charge, to any person obtaining a copy
! of this software and associated documentation files (the "Software"), to deal
! in the Software without restriction, including without limitation the rights
! to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
! copies of the Software, and to permit persons to whom the Software is
! furnished to do so, subject to the following conditions:
!
! The above copyright notice and this permission notice shall be included in
! all copies or substantial portions of the Software.
!
! THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
! IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
! FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
! AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
! LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
! OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
! THE SOFTWARE.

subroutine compute_flux_derivatives(nvars, ndim, nx, ny, nz, &
    states, fluxes, metrics, metric_jacobians, flux_derivatives)
  implicit none

  integer, intent(in) :: nvars, ndim, nx, ny, nz
  real, intent(in) :: states(nvars, -2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(in) :: fluxes(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(in) :: metrics(ndim, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(in) :: metric_jacobians(-2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(out) :: flux_derivatives(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)

  real flux_derivatives_generalized(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real weno_flux_tmp(nvars, 0:nx+1, 0:ny+1, 0:nz+1)

  real generalized_fluxes(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real generalized_states_frozen(nvars, -2:3)
  real generalized_fluxes_frozen(nvars, ndim, -2:3)
  real metric_solve_tmp(ndim)
  real R(nvars, nvars), R_inv(nvars, nvars), lambda_roe(nvars)
  real lambda_pointwise(nvars, -2:3)
  real wavespeeds(nvars)
  integer i, j, k
  real delta_xi, delta_eta, delta_zeta
  real characteristic_fluxes_pos(nvars, -2:3)
  real characteristic_fluxes_neg(nvars, -2:3)
  real metrics_frozen(ndim, ndim)
  real combined_frozen_metrics(ndim)
  integer v, d

  ! grid spacing in generalized coordinates
  delta_xi = 1.0
  delta_eta = 1.0
  delta_zeta = 1.0

  !$loopy begin tagged: to_generalized
  !call convert_to_generalized(nvars, ndim, nx, ny, nz, &
  !  fluxes, metrics, metric_jacobians, generalized_fluxes)

  ! FIXME: Manually inlined for now
  do k=-2,nz+3
    do j=-2,ny+3
      do i=-2,nx+3
        do v=1,nvars
          call mult_mat_vec(ndim, ndim, 1.0/metric_jacobians(i,j,k), metrics(:,:,i,j,k), &
            fluxes(v,:,i,j,k), generalized_fluxes(v,:,i,j,k))
        end do
      end do
    end do
  end do
  !$loopy end tagged: to_generalized

  ! for the next three loops, note carefully that the inner loop indices have a different range

  !$loopy begin tagged: flux_x_compute
  do k=1,nz
    do j=1,ny
      do i=0,nx
        call convert_to_generalized_frozen(nvars, ndim, &
                                           states(:, i-2:i+3, j, k), &
                                           fluxes(:, :, i-2:i+3, j, k), &
                                           metrics(:, :, i:i+1, j, k), &
                                           metric_jacobians(i:i+1, j, k), &
                                           metrics_frozen, &
                                           combined_frozen_metrics, &
                                           generalized_states_frozen, &
                                           generalized_fluxes_frozen)

        call pointwise_eigenvalues(nvars, 1, states(:, i-2:i+3, j, k), lambda_pointwise)
        call roe_eigensystem(nvars, ndim, 1, states(:, i:i+1, j, k), &
                              metrics_frozen, R, R_inv, lambda_roe)

        call lax_wavespeeds(nvars, lambda_pointwise, lambda_roe, wavespeeds)

        call split_characteristic_fluxes(nvars, &
                                         generalized_states_frozen, &
                                         generalized_fluxes_frozen(:, 1, :), &
                                         R_inv, &
                                         wavespeeds, &
                                         characteristic_fluxes_pos, &
                                         characteristic_fluxes_neg)

        call weno_flux(nvars, &
                        generalized_fluxes(:, 1, i-2:i+3, j, k), &
                        characteristic_fluxes_pos, characteristic_fluxes_neg, &
                        combined_frozen_metrics(1), R, weno_flux_tmp(:, i, j, k))
      end do
    end do
  end do
  !$loopy end tagged: flux_x_compute

  !$loopy begin tagged: flux_x_diff
  do k=1,nz
    do j=1,ny
      do i=1,nx
        do v=1,nvars
          flux_derivatives_generalized(v, 1, i, j, k) = &
            (weno_flux_tmp(v, i, j, k) - weno_flux_tmp(v, i-1, j, k)) / delta_xi
        end do
      end do
    end do
  end do
  !$loopy end tagged: flux_x_diff

  !$loopy begin tagged: flux_y_compute
  do k=1,nz
    do i=1,nx
      do j=0,ny
        call convert_to_generalized_frozen(nvars, ndim, &
                                           states(:, i, j-2:j+3, k), &
                                           fluxes(:, :, i, j-2:j+3, k), &
                                           metrics(:, :, i, j:j+1, k), &
                                           metric_jacobians(i, j:j+1, k), &
                                           metrics_frozen, &
                                           combined_frozen_metrics, &
                                           generalized_states_frozen, &
                                           generalized_fluxes_frozen)

        call pointwise_eigenvalues(nvars, 2, states(:, i, j-2:j+3, k), lambda_pointwise)
        call roe_eigensystem(nvars, ndim, 2, states(:, i, j:j+1, k), &
                              metrics_frozen, R, R_inv, lambda_roe)

        call lax_wavespeeds(nvars, lambda_pointwise, lambda_roe, wavespeeds)

        call split_characteristic_fluxes(nvars, &
                                         generalized_states_frozen, &
                                         generalized_fluxes_frozen(:, 2, :), &
                                         R_inv, &
                                         wavespeeds, &
                                         characteristic_fluxes_pos, &
                                         characteristic_fluxes_neg)

        call weno_flux(nvars, &
                        generalized_fluxes(:, 2, i, j-2:j+3, k), &
                        characteristic_fluxes_pos, characteristic_fluxes_neg, &
                        combined_frozen_metrics(2), R, weno_flux_tmp(:, i, j, k))

      end do
    end do
  end do
  !$loopy end tagged: flux_y_compute

  !$loopy begin tagged: flux_y_diff
  do k=1,nz
    do j=1,ny
      do i=1,nx
        do v=1,nvars
          flux_derivatives_generalized(v, 2, i, j, k) = &
            (weno_flux_tmp(v, i, j, k) - weno_flux_tmp(v, i, j-1, k)) / delta_eta
        end do
      end do
    end do
  end do
  !$loopy end tagged: flux_y_diff

  !$loopy begin tagged: flux_z_compute
  do j=1,ny
    do i=1,nx
      do k=0,nz
        call convert_to_generalized_frozen(nvars, ndim, &
                                           states(:, i, j, k-2:k+3), &
                                           fluxes(:, :, i, j, k-2:k+3), &
                                           metrics(:, :, i, j, k:k+1), &
                                           metric_jacobians(i, j, k:k+1), &
                                           metrics_frozen, &
                                           combined_frozen_metrics, &
                                           generalized_states_frozen, &
                                           generalized_fluxes_frozen)

        call pointwise_eigenvalues(nvars, 3, states(:, i, j, k-2:k+3), lambda_pointwise)
        call roe_eigensystem(nvars, ndim, 3, states(:, i, j, k:k+1), &
                              metrics_frozen, R, R_inv, lambda_roe)

        call lax_wavespeeds(nvars, lambda_pointwise, lambda_roe, wavespeeds)

        call split_characteristic_fluxes(nvars, &
                                         generalized_states_frozen, &
                                         generalized_fluxes_frozen(:, 3, :), &
                                         R_inv, &
                                         wavespeeds, &
                                         characteristic_fluxes_pos, &
                                         characteristic_fluxes_neg)

        call weno_flux(nvars, &
                        generalized_fluxes(:, 3, i, j, k-2:k+3), &
                        characteristic_fluxes_pos, characteristic_fluxes_neg, &
                        combined_frozen_metrics(3), R, weno_flux_tmp(:, i, j, k))
      end do
    end do
  end do
  !$loopy end tagged: flux_z_compute

  !$loopy begin tagged: flux_z_diff
  do k=1,nz
    do j=1,ny
      do i=1,nx
        do v=1,nvars
          flux_derivatives_generalized(v, 3, i, j, k) = &
            (weno_flux_tmp(v, i, j, k) - weno_flux_tmp(v, i, j, k-1)) / delta_zeta
        end do
      end do
    end do
  end do
  !$loopy end tagged: flux_z_diff

  !$loopy begin tagged: from_generalized
  !call convert_from_generalized(nvars, ndim, nx, ny, nz, &
  !                              flux_derivatives_generalized, &
  !                              metrics, &
  !                              metric_jacobians, &
  !                              flux_derivatives)
  ! FIXME: Manually inlined for now
  do k=-2,nz+3
    do j=-2,ny+3
      do i=-2,nx+3
        do v=1,nvars
          call solve3x3(metrics(:,:,i,j,k), flux_derivatives_generalized(v,:,i,j,k), &
            metric_solve_tmp)
          do d=1,ndim
            flux_derivatives(v,d,i,j,k) = metric_solve_tmp(d)*metric_jacobians(i,j,k)
          end do
        end do
      end do
    end do
  end do
  !$loopy end tagged: from_generalized
end subroutine

subroutine pointwise_eigenvalues(nvars, d, states, lambda_pointwise)
  implicit none

  integer, intent(in) :: nvars
  integer, intent(in) :: d
  real, intent(in) :: states(nvars, -2:3)
  real, intent(out) :: lambda_pointwise(nvars, -2:3)

  real u(3), c, p, rho, ke
  integer v, k, i

  do k=-2,3
    rho = states(1,k)
    ke = 0.0
    do i=1,3
      u(i) = states(i+1,k)/rho
      ke = ke + u(i)**2
    end do
    p = (1.4 - 1)*(states(nvars,k) - 0.5*rho*ke)
    c = sqrt(1.4*p/rho)

    do v=1,nvars-2
      lambda_pointwise(v,k) = u(d)
    end do
    lambda_pointwise(nvars-1,k) = u(d) + c
    lambda_pointwise(nvars,k) = u(d) - c
  end do
end subroutine

subroutine roe_eigensystem(nvars, ndim, d, states, metrics_frozen, R, R_inv, lambda_roe)
  implicit none

  integer, intent(in) :: nvars
  integer, intent(in) :: ndim
  integer, intent(in) :: d
  real, intent(in) :: states(nvars, 2)
  real, intent(in) :: metrics_frozen(ndim, ndim)
  real, intent(out) :: R(nvars, nvars)
  real, intent(out) :: R_inv(nvars, nvars)
  real, intent(out) :: lambda_roe(nvars)

  integer ik, il, im
  real metric_norm(ndim)
  real p, ke
  real u_orig(ndim, 2), H_orig(2)
  real sqrt_rho(2), sqrt_rho_sum
  real u(ndim), rho, c, H, q
  real u_tilde(ndim), k(ndim), l(ndim), m(ndim)
  integer i, j
  real b1, b2, alpha, beta

  ik = d
  il = d+1
  im = d+2
  if (il > 3) il = il - 3
  if (im > 3) im = im - 3

  do j=1,2
    do i=1,ndim
      u_orig(i,j) = states(i+1,j)/states(1,j)
    end do
  end do

  do j=1,2
    ke = 0.0
    do i=1,ndim
      ke = ke + u_orig(i,j)**2
    end do

    p = (1.4 - 1.0)*(states(nvars,j) - 0.5*states(1,j)*ke)
    H_orig(j) = (states(nvars,j) + p)/states(1,j)
  end do

  do i=1,ndim
    metric_norm(i) = 0.0
    do j=1,ndim
      metric_norm(i) = metric_norm(i) + metrics_frozen(i,j)**2
    end do
    metric_norm(i) = sqrt(metric_norm(i))
  end do

  do i=1,ndim
    k(i) = metrics_frozen(ik,i)/metric_norm(ik)
    l(i) = metrics_frozen(il,i)/metric_norm(il)
    m(i) = metrics_frozen(im,i)/metric_norm(im)
  end do

  sqrt_rho(1) = sqrt(states(1,1))
  sqrt_rho(2) = sqrt(states(1,2))
  sqrt_rho_sum = sqrt_rho(1) + sqrt_rho(2)

  do i=1,ndim
    u(i) = (u_orig(i,1)*sqrt_rho(1) + u_orig(i,2)*sqrt_rho(2))/sqrt_rho_sum
  end do
  rho = sqrt_rho(1)*sqrt_rho(2)
  H = (H_orig(1)*sqrt_rho(1) + H_orig(2)*sqrt_rho(2))/sqrt_rho_sum

  q = 0.0
  do i=1,ndim
    q = q + u(i)**2
  end do

  c = sqrt((1.4 - 1.0)*(H - 0.5*q))

  b1 = (1.4 - 1.0)/(c**2)
  b2 = 1.0 + b1*q - b1*H

  u_tilde(1) = 0.0
  do i=1,ndim
    u_tilde(1) = u_tilde(1) + k(i)*u(i)
  end do

  u_tilde(2) = 0.0
  do i=1,ndim
    u_tilde(2) = u_tilde(2) + l(i)*u(i)
  end do

  u_tilde(3) = 0.0
  do i=1,ndim
    u_tilde(3) = u_tilde(3) + m(i)*u(i)
  end do

  alpha = rho/(2.0*c)
  beta = 1.0/(2.0*alpha)

  lambda_roe(1) = u_tilde(1)*metric_norm(ik)
  lambda_roe(2) = u_tilde(1)*metric_norm(ik)
  lambda_roe(3) = u_tilde(1)*metric_norm(ik)
  lambda_roe(4) = u_tilde(1)*metric_norm(ik) + c*metric_norm(ik)
  lambda_roe(5) = u_tilde(1)*metric_norm(ik) - c*metric_norm(ik)

  R(1,1) = 1.0
  R(2,1) = u(1)
  R(3,1) = u(2)
  R(4,1) = u(3)
  R(5,1) = H - 1.0/b1

  R(1,2) = 0.0
  R(2,2) = rho*l(1)
  R(3,2) = rho*l(2)
  R(4,2) = rho*l(3)
  R(5,2) = rho*u_tilde(2)

  R(1,3) = 0.0
  R(2,3) = rho*m(1)
  R(3,3) = rho*m(2)
  R(4,3) = rho*m(3)
  R(5,3) = rho*u_tilde(3)

  R(1,4) = alpha
  R(2,4) = alpha*(u(1) + c*k(1))
  R(3,4) = alpha*(u(2) + c*k(2))
  R(4,4) = alpha*(u(3) + c*k(3))
  R(5,4) = alpha*(H + c*u_tilde(1))

  R(1,5) = alpha
  R(2,5) = alpha*(u(1) - c*k(1))
  R(3,5) = alpha*(u(2) - c*k(2))
  R(4,5) = alpha*(u(3) - c*k(3))
  R(5,5) = alpha*(H - c*u_tilde(1))

  R_inv(1,1) = 1.0 - b2
  R_inv(1,2) = b1*u(1)
  R_inv(1,3) = b1*u(2)
  R_inv(1,4) = b1*u(3)
  R_inv(1,5) = -b1

  R_inv(2,1) = -u_tilde(2)/rho
  R_inv(2,2) = l(1)/rho
  R_inv(2,3) = l(2)/rho
  R_inv(2,4) = l(3)/rho
  R_inv(2,5) = 0.0

  R_inv(3,1) = -u_tilde(3)/rho
  R_inv(3,2) = m(1)/rho
  R_inv(3,3) = m(2)/rho
  R_inv(3,4) = m(3)/rho
  R_inv(3,5) = 0.0

  R_inv(4,1) = beta*(b2 - u_tilde(1)/c)
  R_inv(4,2) = -beta*(b1*u(1) - k(1)/c)
  R_inv(4,3) = -beta*(b1*u(2) - k(2)/c)
  R_inv(4,4) = -beta*(b1*u(3) - k(3)/c)
  R_inv(4,5) = beta*b1

  R_inv(5,1) = beta*(b2 + u_tilde(1)/c)
  R_inv(5,2) = -beta*(b1*u(1) + k(1)/c)
  R_inv(5,3) = -beta*(b1*u(2) + k(2)/c)
  R_inv(5,4) = -beta*(b1*u(3) + k(3)/c)
  R_inv(5,5) = beta*b1
end subroutine

subroutine lax_wavespeeds(nvars, lambda_pointwise, lambda_roe, wavespeeds)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: lambda_pointwise(nvars, -2:3)
  real, intent(in) :: lambda_roe(nvars)
  real, intent(out) :: wavespeeds(nvars)

  real kappa
  integer v
  integer k

  kappa = 1.1

  do v=1,nvars
    wavespeeds(v) = abs(lambda_roe(v))
    do k=-2,3
      wavespeeds(v) = max(wavespeeds(v), abs(lambda_pointwise(v,k)))
    end do
    wavespeeds(v) = kappa*wavespeeds(v)
  end do
end subroutine

subroutine convert_to_generalized_frozen( &
            nvars, ndim, &
            states, &
            fluxes, &
            metrics, &
            metric_jacobians, &
            metrics_frozen, &
            combined_frozen_metrics, &
            generalized_states_frozen, &
            generalized_fluxes_frozen)
  integer, intent(in) :: nvars, ndim
  real, intent(in) :: states(nvars, -2:3)
  real, intent(in) :: fluxes(nvars, ndim, -2:3)
  real, intent(in) :: metrics(ndim, ndim, 2)
  real, intent(in) :: metric_jacobians(2)
  real, intent(out) :: metrics_frozen(ndim, ndim)
  real, intent(out) :: combined_frozen_metrics(ndim)
  real, intent(out) :: generalized_states_frozen(nvars, -2:3)
  real, intent(out) :: generalized_fluxes_frozen(nvars, ndim, -2:3)

  real jacobian_frozen

  integer c, k, v
  integer i, j

  do j=1,ndim
    do i=1,ndim
      metrics_frozen(i,j) = 0.0
      do c=1,2
        metrics_frozen(i,j) = metrics_frozen(i,j) + 0.5*metrics(i,j,c)/metric_jacobians(c)
      end do
    end do
  end do
  jacobian_frozen = 0.5*(metric_jacobians(1) + metric_jacobians(2))

  do i=1,ndim
    combined_frozen_metrics(i) = 0.0
    do j=1,ndim
      combined_frozen_metrics(i) = combined_frozen_metrics(i) + metrics_frozen(i,j)**2
    end do
  end do

  do k=-2,3
    do v=1,nvars
      generalized_states_frozen(v,k) = states(v,k)/jacobian_frozen
    end do
  end do

  do k=-2,3
    do v=1,nvars
      call mult_mat_vec(ndim, ndim, 1.0, metrics_frozen, &
        fluxes(v,:,k), generalized_fluxes_frozen(v,:,k))
    end do
  end do
end subroutine

#if 0
subroutine convert_to_generalized(nvars, ndim, nx, ny, nz, &
                                  fluxes, &
                                  metrics, &
                                  metric_jacobians, &
                                  generalized_fluxes)
  implicit none

  integer, intent(in) :: nvars, ndim, nx, ny, nz
  real, intent(in) :: fluxes(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(in) :: metrics(ndim, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(in) :: metric_jacobians(-2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(out) :: generalized_fluxes(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)

  integer i, j, k, v

  do k=-2,nz+3
    do j=-2,ny+3
      do i=-2,nx+3
        do v=1,nvars
          call mult_mat_vec(ndim, ndim, 1.0/metric_jacobians(i,j,k), metrics(:,:,i,j,k), &
            fluxes(v,:,i,j,k), generalized_fluxes(v,:,i,j,k))
        end do
      end do
    end do
  end do
end subroutine

subroutine convert_from_generalized(nvars, ndim, nx, ny, nz, &
                                    flux_derivatives_generalized, &
                                    metrics, &
                                    metric_jacobians, &
                                    flux_derivatives)
  implicit none

  integer, intent(in) :: nvars, ndim, nx, ny, nz
  real, intent(in) :: flux_derivatives_generalized(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(in) :: metrics(ndim, ndim, -2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(in) :: metric_jacobians(-2:nx+3, -2:ny+3, -2:nz+3)
  real, intent(out) :: flux_derivatives(nvars, ndim, -2:nx+3, -2:ny+3, -2:nz+3)

  integer i, j, k, v
  integer d

  do k=-2,nz+3
    do j=-2,ny+3
      do i=-2,nx+3
        do v=1,nvars
          call solve3x3(metrics(:,:,i,j,k), flux_derivatives_generalized(v,:,i,j,k), &
            flux_derivatives(v,:,i,j,k))
          do d=1,ndim
            flux_derivatives(v,d,i,j,k) = flux_derivatives(v,d,i,j,k)*metric_jacobians(i,j,k)
          end do
        end do
      end do
    end do
  end do
end subroutine
#endif

subroutine solve3x3(a, b, x)
  implicit none

  real, intent(in) :: a(3,3), b(3)
  real, intent(out) :: x(3)

  real temp(3,3)
  real det_a(1), det_temp(1)
  integer k
  integer i,j

  call determinant3x3(a, det_a)
  do k=1,3
    do j=1,3
      do i=1,3
        temp(i,j) = a(i,j)
        temp(i,k) = b(k)
      end do
    end do
    call determinant3x3(temp, det_temp)
    x(k) = det_temp(1)/det_a(1)
  end do
end subroutine

subroutine determinant3x3(a, det)
  implicit none

  real, intent(in) :: a(3,3)
  real, intent(out) :: det(1)

  det(1) = a(1,1)*(a(2,2)*a(3,3) - a(2,3)*a(3,2)) &
        + a(1,2)*(a(2,3)*a(3,1) - a(2,1)*a(3,3)) &
        + a(1,3)*(a(2,1)*a(3,2) - a(2,2)*a(3,1))
end subroutine

subroutine split_characteristic_fluxes(nvars, &
                                       generalized_states_frozen, &
                                       generalized_fluxes_frozen, &
                                       R_inv, &
                                       wavespeeds, &
                                       characteristic_fluxes_pos, &
                                       characteristic_fluxes_neg)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: generalized_states_frozen(nvars, -2:3)
  real, intent(in) :: generalized_fluxes_frozen(nvars, -2:3)
  real, intent(in) :: R_inv(nvars, nvars)
  real, intent(in) :: wavespeeds(nvars)
  real, intent(out) :: characteristic_fluxes_pos(nvars, -2:3)
  real, intent(out) :: characteristic_fluxes_neg(nvars, -2:3)

  integer k, m
  integer l

  do k=-2,3
    do m=1,nvars
      characteristic_fluxes_pos(m,k) = 0.0
      characteristic_fluxes_neg(m,k) = 0.0
      do l=1,nvars
        characteristic_fluxes_pos(m,k) = characteristic_fluxes_pos(m,k) &
          + 0.5*R_inv(m,l) &
          *(generalized_fluxes_frozen(l,k) + wavespeeds(m)*generalized_states_frozen(l,k))
        characteristic_fluxes_neg(m,k) = characteristic_fluxes_neg(m,k) &
          + 0.5*R_inv(m,l) &
          *(generalized_fluxes_frozen(l,k) - wavespeeds(m)*generalized_states_frozen(l,k))
      end do
    end do
  end do
end subroutine

subroutine weno_flux(nvars, generalized_fluxes, characteristic_fluxes_pos, &
    characteristic_fluxes_neg, combined_frozen_metrics, R, flux)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: generalized_fluxes(nvars, -2:3)
  real, intent(in) :: characteristic_fluxes_pos(nvars, -2:3)
  real, intent(in) :: characteristic_fluxes_neg(nvars, -2:3)
  real, intent(in) :: combined_frozen_metrics
  real, intent(in) :: R(nvars, nvars)
  real, intent(out) :: flux(nvars)

  real consistent(nvars)
  real dissipation_pos(nvars)
  real dissipation_neg(nvars)
  integer v

  call consistent_part(nvars, generalized_fluxes, consistent)
  call dissipation_part_pos(nvars, characteristic_fluxes_pos, &
    combined_frozen_metrics, R, dissipation_pos)
  call dissipation_part_neg(nvars, characteristic_fluxes_neg, &
    combined_frozen_metrics, R, dissipation_neg)

  do v=1,nvars
    flux(v) = consistent(v) + dissipation_pos(v) + dissipation_neg(v)
  end do
end subroutine

subroutine consistent_part(nvars, generalized_fluxes, consistent)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: generalized_fluxes(nvars, -2:3)
  real, intent(out) :: consistent(nvars)

  integer v

  do v=1,nvars
    consistent(v) = generalized_fluxes(v,-2) - 8*generalized_fluxes(v,-1) &
      + 37*generalized_fluxes(v,0) + 37*generalized_fluxes(v,1) &
      - 8*generalized_fluxes(v,2) + generalized_fluxes(v,3)
  end do
end subroutine

subroutine dissipation_part_pos(nvars, characteristic_fluxes, &
  combined_frozen_metrics, R, dissipation_pos)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(in) :: combined_frozen_metrics
  real, intent(in) :: R(nvars, nvars)
  real, intent(out) :: dissipation_pos(nvars)

  real combined_fluxes(nvars)

  call weno_combination_pos(nvars, characteristic_fluxes, combined_frozen_metrics, combined_fluxes)

  call mult_mat_vec(nvars, nvars, -1.0/60, R, combined_fluxes, dissipation_pos)
end subroutine

subroutine weno_combination_pos(nvars, characteristic_fluxes, &
  combined_frozen_metrics, combined_fluxes)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(in) :: combined_frozen_metrics
  real, intent(out) :: combined_fluxes(nvars)

  real w(nvars, 3)
  real flux_differences(nvars, 3)
  integer v

  call weno_weights_pos(nvars, characteristic_fluxes, combined_frozen_metrics, w)
  call flux_differences_pos(nvars, characteristic_fluxes, flux_differences)

  do v=1,nvars
    combined_fluxes(v) = (20*w(v,1) - 1)*flux_differences(v,1) &
      - (10*(w(v,1) + w(v,2)) - 5)*flux_differences(v,2) &
      + flux_differences(v,3)
  end do
end subroutine

subroutine weno_weights_pos(nvars, characteristic_fluxes, combined_frozen_metrics, w)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(in) :: combined_frozen_metrics
  real, intent(out) :: w(nvars, 3)

  real :: C(3) = (/0.1, 0.6, 0.3/)
  real :: weights1(0:2) = (/1, -4, 3/)
  real :: weights2(0:2) = (/-1, 0, 1/)
  real :: weights3(0:2) = (/-3, 4, 1/)
  real :: weightsc(0:2) = (/1, -2, 1/)

  real IS(3), alpha(3), eps

  real w1sum(1), w2sum(1), w3sum(1), c1sum(1), c2sum(1), c3sum(1)

  integer p, i, j
  real sum_alpha(1)

  eps = 1e-6*combined_frozen_metrics
  p = 2

  do i=1,nvars
    call weighted_sum(3, weights1, characteristic_fluxes(i,-2:0), w1sum)
    call weighted_sum(3, weightsc, characteristic_fluxes(i,-2:0), c1sum)
    call weighted_sum(3, weights2, characteristic_fluxes(i,-1:1), w2sum)
    call weighted_sum(3, weightsc, characteristic_fluxes(i,-1:1), c2sum)
    call weighted_sum(3, weights3, characteristic_fluxes(i,0:2), w3sum)
    call weighted_sum(3, weightsc, characteristic_fluxes(i,0:2), c3sum)

    IS(1) = (1.0/4)*w1sum(1)**2 + (13.0/12)*c1sum(1)**2
    IS(2) = (1.0/4)*w2sum(1)**2 + (13.0/12)*c2sum(1)**2
    IS(3) = (1.0/4)*w3sum(1)**2 + (13.0/12)*c3sum(1)**2

    do j=1,3
      alpha(j) = C(j)/(IS(j) + eps)**p
    end do
    call sum_array(3, alpha, sum_alpha)
    do j=1,3
      w(i,j) = alpha(j)/sum_alpha(1)
    end do
  end do
end subroutine

subroutine flux_differences_pos(nvars, characteristic_fluxes, flux_differences)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(out) :: flux_differences(nvars, 3)

  integer i, v

  do i=1,3
    do v=1,nvars
      flux_differences(v,i) = -characteristic_fluxes(v,-3+i) + 3*characteristic_fluxes(v,-2+i) &
                                - 3*characteristic_fluxes(v,-1+i) + characteristic_fluxes(v,i)
    end do
  end do
end subroutine

subroutine dissipation_part_neg(nvars, characteristic_fluxes, &
  combined_frozen_metrics, R, dissipation_neg)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(in) :: combined_frozen_metrics
  real, intent(in) :: R(nvars, nvars)
  real, intent(out) :: dissipation_neg(nvars)

  real combined_fluxes(nvars)

  call weno_combination_neg(nvars, characteristic_fluxes, combined_frozen_metrics, combined_fluxes)

  call mult_mat_vec(nvars, nvars, 1.0/60, R, combined_fluxes, dissipation_neg)
end subroutine

subroutine weno_combination_neg(nvars, characteristic_fluxes, &
  combined_frozen_metrics, combined_fluxes)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(in) :: combined_frozen_metrics
  real, intent(out) :: combined_fluxes(nvars)

  real w(nvars, 3)
  real flux_differences(nvars, 3)
  integer v

  call weno_weights_neg(nvars, characteristic_fluxes, combined_frozen_metrics, w)
  call flux_differences_neg(nvars, characteristic_fluxes, flux_differences)

  do v=1,nvars
    combined_fluxes(v) = (20*w(v,1) - 1)*flux_differences(v,1) &
      - (10*(w(v,1) + w(v,2)) - 5)*flux_differences(v,2) &
      + flux_differences(v,3)
  end do
end subroutine

subroutine weno_weights_neg(nvars, characteristic_fluxes, combined_frozen_metrics, w)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(in) :: combined_frozen_metrics
  real, intent(out) :: w(nvars, 3)

  real :: C(3) = (/0.1, 0.6, 0.3/)
  real :: weights1(0:2) = (/1, -4, 3/)
  real :: weights2(0:2) = (/-1, 0, 1/)
  real :: weights3(0:2) = (/-3, 4, 1/)
  real :: weightsc(0:2) = (/1, -2, 1/)

  real IS(3), alpha(3), eps
  real w1sum(1), w2sum(1), w3sum(1), c1sum(1), c2sum(1), c3sum(1)

  integer p, i, j
  real sum_alpha(1)

  eps = 1e-6*combined_frozen_metrics
  p = 2

  do i=1,nvars
    call weighted_sum(3, weights1, characteristic_fluxes(i,3:1:-1), w1sum)
    call weighted_sum(3, weightsc, characteristic_fluxes(i,3:1:-1), c1sum)
    call weighted_sum(3, weights2, characteristic_fluxes(i,2:0:-1), w2sum)
    call weighted_sum(3, weightsc, characteristic_fluxes(i,2:0:-1), c2sum)
    call weighted_sum(3, weights3, characteristic_fluxes(i,1:-1:-1), w3sum)
    call weighted_sum(3, weightsc, characteristic_fluxes(i,1:-1:-1), c3sum)

    IS(1) = (1.0/4)*w1sum(1)**2 + (13.0/12)*c1sum(1)**2
    IS(2) = (1.0/4)*w2sum(1)**2 + (13.0/12)*c2sum(1)**2
    IS(3) = (1.0/4)*w3sum(1)**2 + (13.0/12)*c3sum(1)**2

    do j=1,3
      alpha(j) = C(j)/(IS(j) + eps)**p
    end do
    call sum_array(3, alpha, sum_alpha)
    do j=1,3
      w(i,j) = alpha(j)/sum_alpha(1)
    end do
  end do
end subroutine

subroutine flux_differences_neg(nvars, characteristic_fluxes, flux_differences)
  implicit none

  integer, intent(in) :: nvars
  real, intent(in) :: characteristic_fluxes(nvars, -2:3)
  real, intent(out) :: flux_differences(nvars, 3)

  integer i, v

  do i=1,3
    do v=1,nvars
      flux_differences(v,i) = -characteristic_fluxes(v,1-i) + 3*characteristic_fluxes(v,2-i) &
                                - 3*characteristic_fluxes(v,3-i) + characteristic_fluxes(v,4-i)
    end do
  end do
end subroutine

subroutine mult_mat_vec(m, n, alpha, a, b, c)
  implicit none

  integer, intent(in) :: m, n
  real, intent(in) :: alpha
  real, intent(in) :: a(m, n)
  real, intent(in) :: b(n)
  real, intent(out) :: c(m)

  integer i, j
  real accumulator

  do i=1,m
    accumulator = 0.0
    do j=1,n
      accumulator = accumulator + alpha*a(i,j)*b(j)
    end do
    c(i) = accumulator
  end do
end subroutine

subroutine sum_array(n, a, a_sum)
  implicit none

  integer, intent(in) :: n
  real, intent(in) :: a(n)
  real, intent(out) :: a_sum(1)

  integer i

  a_sum(1) = 0.0
  do i=1,n
    a_sum(1) = a_sum(1) + a(i)
  end do
end subroutine

subroutine weighted_sum(n, w, a, a_sum)
  implicit none

  integer, intent(in) :: n
  real, intent(in) :: a(n), w(n)
  real, intent(out) :: a_sum(1)

  integer i

  a_sum(1) = 0.0
  do i=1,n
    a_sum(1) = a_sum(1) + w(i)*a(i)
  end do
end subroutine

!$loopy begin
!
! prg = lp.parse_fortran(lp.c_preprocess(SOURCE), FILENAME)
! prg = lp.fix_parameters(prg, ndim=3, nvars=5, _remove=False)
!
! cfd = prg["compute_flux_derivatives"]
!
! cfd = lp.assume(cfd, "nx > 0 and ny > 0 and nz > 0")
!
! cfd = lp.set_temporary_scope(cfd, "flux_derivatives_generalized",
!         lp.AddressSpace.GLOBAL)
! cfd = lp.set_temporary_scope(cfd, "generalized_fluxes",
!         lp.AddressSpace.GLOBAL)
! cfd = lp.set_temporary_scope(cfd, "weno_flux_tmp",
!         lp.AddressSpace.GLOBAL)
!
! prg = prg.with_kernel(cfd)
!
! RESULT = prg
!
!$loopy end