// Implementation of a variation of the HYPRE method for the steady
// Navier--Stokes problem from http://dx.doi.org/10.1137/25M1736104

method hypre_navier_stokes {
  parameter pi = 3.141592653589793238462643383279502884197169399375105820974944592307816406286198

  // Viscosity
  parameter nu = 0.025

  // Reynolds number
  parameter reynolds = pow(nu, -1.0)

  // Kovasznay solution on (-0.5, 1.5) x (0, 2)
  parameter lambda =
    0.5 * reynolds
    - sqrt(0.25 * pow(reynolds, 2.0) + 4.0 * pow(pi, 2.0))

  parameter exact_pressure_average =
    0.5 * lambda * (exp(4.0 * lambda) - 1.0)
    - (exp(3.0 * lambda) - exp(-lambda)) * pow(8.0 * lambda, -1.0)

  function exact_velocity(vector X) -> vector
  = vector(
            1.0 - exp(lambda * X[0]) * cos(2.0 * pi * X[1]),
            lambda * pow(2.0 * pi, -1.0)
            * exp(lambda * X[0]) * sin(2.0 * pi * X[1])
          )

  function exact_pressure(vector X) -> scalar =
    -0.5 * exp(2.0 * lambda * X[0])
    + 0.5 * lambda * (exp(4.0 * lambda) - 1.0)

  function exact_lagrange_multiplier(vector X) -> scalar = 0.0

  function source(vector X) -> vector = vector(0.0, 0.0)

  // Velocity space
  space Uh {
    element RaviartThomasPoly(k+1, vector)
    edge Poly(k, vector)
  }

  // Pressure space
  space Ph {
    element Poly(k, scalar)
    edge Poly(k, scalar)
  }

  // Lagrange multiplier used to prescribe the pressure average
  space Lh {
    domain Poly(0, scalar)
  }

  // Velocity-pressure-multiplier Cartesian product space
  product space Xh = Uh times Ph times Lh

  // Velocity interpolant
  interpolant interpolant_exact_velocity on Uh {
    on element T:
      dof(T) = raviart_thomas_interpolate(
                                           exact_velocity,
                                           RaviartThomasPoly(k+1)
                                         )
    on edge E:
      dof(E) = l2_project(exact_velocity, Poly(k, vector))
  }

  // Pressure interpolant
  interpolant interpolant_exact_pressure on Ph {
    on element T:
      dof(T) = l2_projection(exact_pressure, Poly(k, scalar))
    on edge E:
      dof(E) = l2_project(exact_pressure, Poly(k, scalar))
  }

  // Exact value of the pressure-average Lagrange multiplier
  interpolant interpolant_exact_lagrange_multiplier on Lh {
    on domain Omega:
      dof(Omega) = l2_project(exact_lagrange_multiplier, Poly(0, scalar))
  }

  // Operators
  operator raviart_thomas_interpolator_polykpo
    : Poly(k+1, vector)(u) -> RaviartThomasPoly(k+1) on element T {
    forall v in Poly(k-1, vector):
      int(T) raviart_thomas_interpolator_polykpo(u) dot v =
      int(T) u dot v
    forall edge E: forall q in Poly(k, scalar):
      int(E) (raviart_thomas_interpolator_polykpo(u) dot normal) * q =
      int(E) (u dot normal) * q
  }

  operator velocity_gradient : Uh(u) -> Poly(k, matrix) on element T {
    forall tau in Poly(k, matrix):
      int(T) velocity_gradient(u) dot tau =
      - int(T) dof(u, T) dot div(tau)
      + int(dT) dof(u, E) dot (tau dot normal)
  }

  operator velocity_divergence : Uh(u) -> Poly(k, scalar) on element T {
    velocity_divergence(u) = trace(velocity_gradient(u))
  }

  operator velocity_potential : Uh(u) -> Poly(k+1, vector) on element T {
    forall w in Poly(k+1, vector):
      int(T) grad(velocity_potential(u)) dot grad(w) =
      int(T) velocity_gradient(u) dot grad(w)

    if (k = 0)
      constraint int(T) velocity_potential(u) =
      0.5 * int(dT) orthogonal_distance(T, E) * dof(u, E)
    else
      constraint int(T) velocity_potential(u) = int(T) dof(u, T)
  }

  operator velocity_field : Uh(u) -> RaviartThomasPoly(k+1) on element T {
    velocity_field(u) = dof(u, T)
  }

  operator pressure_potential : Ph(p) -> Poly(k+1, scalar) on element T {
    forall q in Poly(k+1, scalar):
      int(T) grad(pressure_potential(p)) dot grad(q) =
      -int(T) dof(p, T) * div(grad(q))
      + int(dT) dof(p, E) * (grad(q) dot normal)
    constraint int(T) pressure_potential(p) = int(T) dof(p, T)
  }

  operator element_difference : Uh(u) -> RaviartThomasPoly(k+1) on element T {
    forall w in RaviartThomasPoly(k+1):
      int(T) element_difference(u) dot w =
      int(T) (
               raviart_thomas_interpolator_polykpo(velocity_potential(u))
               - dof(u, T)
             ) dot w
  }

  operator edge_difference : Uh(u) -> Poly(k, vector) on edge E of element T {
    forall w in Poly(k, vector):
      int(E) edge_difference(u) dot w =
      int(E) (velocity_potential(u) - dof(u, E)) dot w
  }

  operator edge_jump : Uh(u) -> Poly(k, vector) on edge E of element T {
    forall w in Poly(k, vector):
      int(E) edge_jump(u) dot w =
      int(E) (dof(u, E) - dof(u, T)) dot w
  }

  operator pressure_gradient : Ph(p) -> RaviartThomasPoly(k+1) on element T {
    forall w in RaviartThomasPoly(k+1):
      int(T) pressure_gradient(p) dot w =
      - int(T) dof(p, T) * div(w)
      + int(dT) dof(p, E) * (w dot normal)
  }

  operator pressure_field : Ph(p) -> Poly(k) on element T {
    pressure_field(p) = dof(p, T)
  }

  // Linear forms
  bilinear form viscous_term : Uh(trial u) times Uh(test v) {
    sum_elements(
                  nu * int(T) velocity_gradient(u) dot velocity_gradient(v)
                  + nu * n_element_edges(T) * pow(element_area(T), -1.0) * int(T) element_difference(u) dot element_difference(v)
                  // + nu * pow(diameter(T), -1.0) * int(dT) edge_difference_1(u) dot edge_difference_1(v)
                  + nu * pow(diameter(T), -1.0) * int(dT) edge_difference(u) dot edge_difference(v)
                )
  }

  bilinear form pressure_velocity : Ph(trial p) times Uh(test v) {
    sum_elements(int(T) pressure_gradient(p) dot dof(v, T))
  }

  bilinear form velocity_pressure : Uh(trial u) times Ph(test q) {
    - sum_elements(int(T) dof(u, T) dot pressure_gradient(q))
    + sum_boundary_edges(int(E) (dof(u, E) dot normal) * dof(q, E))
  }

  bilinear form pressure_average : Ph(trial p) times Lh(test m) {
    sum_elements(int(T) dof(p, T) * dof(m, Omega))
  }

  bilinear form average_pressure : Lh(trial l) times Ph(test q) {
    - sum_elements(int(T) dof(l, Omega) * dof(q, T))
  }

  linear form momentum_balance_load : Uh(test v) {
    sum_elements(int(T) source dot dof(v, T))
  }

  linear form pressure_average_load : Lh(test m) {
    sum_elements(int(T) exact_pressure_average * dof(m, Omega))
  }

  residual form convective_residual : Uh(state current) times Uh(test v) on element T {
      int(T) (grad(dof(current, T)) dot dof(current, T)) dot dof(v, T)
      + 0.5 * sum_edges(
                         int(E) (dof(current, T) dot normal)
                         * (dof(current, E) - dof(current, T))
                         dot (dof(v, E) + dof(v, T))
                       )
      + 0.5 * sum_edges(
                         int(E) abs(dof(current, E) dot normal)
                         * (dof(current, E) - dof(current, T))
                         dot (dof(v, E) - dof(v, T))
                       )
    }

  jacobian form convective_jacobian
    : Uh(state current) times Uh(trial du) times Uh(test v) on element T {
      int(T)  (grad(dof(current, T)) dot dof(du, T)) dot dof(v, T)
      + 0.5 * sum_edges(
                         int(E) (dof(du, T) dot normal)
                         * (dof(current, E) - dof(current, T))
                         dot (dof(v, E) + dof(v, T))
                       )
      + int(T) (grad(dof(du, T)) dot dof(current, T)) dot dof(v, T)
      + 0.5 * sum_edges(
                         int(E) (dof(current, T) dot normal)
                         * (dof(du, E) - dof(du, T))
                         dot (dof(v, E) + dof(v, T))
                       )
      + 0.5 * sum_edges(
                         int(E) abs(dof(current, E) dot normal)
                         * (dof(du, E) - dof(du, T))
                         dot (dof(v, E) - dof(v, T))
                       )
    }

  jacobian form convective_fixed_point
    : Uh(state current) times Uh(trial du) times Uh(test v) on element T {
      int(T) (grad(dof(du, T)) dot dof(current, T)) dot dof(v, T)
      + 0.5 * sum_edges(
                         int(E) (dof(current, T) dot normal)
                         * (dof(du, E) - dof(du, T))
                         dot (dof(v, E) + dof(v, T))
                       )
      + 0.5 * sum_edges(
                         int(E) abs(dof(current, E) dot normal)
                         * (dof(du, E) - dof(du, T))
                         dot (dof(v, E) - dof(v, T))
                       )
    }

  // Dirichlet boundary conditions
  boundary conditions dirichlet_boundary_conditions on Uh {
    on edge E:
      dof(E) = l2_project(exact_velocity, Poly(k, vector))
  }

  // Nonlinear problem
  nonlinear problem hypre_navier_stokes_problem on Xh {
    residual {
      viscous_term
      + pressure_velocity
      + velocity_pressure
      + pressure_average
      + average_pressure
      + convective_residual
      - momentum_balance_load
      - pressure_average_load
    }
    jacobian {
      viscous_term
      + pressure_velocity
      + velocity_pressure
      + pressure_average
      + average_pressure
      + convective_jacobian
    }
    fixed_point {
      viscous_term
      + pressure_velocity
      + velocity_pressure
      + pressure_average
      + average_pressure
      + convective_fixed_point
    }
    boundary conditions dirichlet_boundary_conditions
    export { velocity_potential, velocity_field, pressure_potential, pressure_field }
    compute errors using interpolant_exact_velocity {
      velocity_h1_norm, velocity_l2_norm
    }
    compute errors using interpolant_exact_pressure {
      pressure_l2_norm
    }
  }

  // Error functionals
  function velocity_h1_norm : Uh(v) -> scalar {
    sqrt(
          sum_elements(
                        int(T) squared_norm(grad(dof(v, T)))
                        + pow(diameter(T), -1.0) * int(dT) squared_norm(edge_jump(v))
                      )
        )
  }

  function velocity_l2_norm : Uh(v) -> scalar {
    sqrt(
          sum_elements(
                        int(T) squared_norm(dof(v, T))
                        + diameter(T) * int(dT) squared_norm(edge_jump(v))
                      )
        )
  }

  function divergence_l2_norm : Uh(v) -> scalar {
    sqrt(sum_elements(int(T) pow(velocity_divergence(v), 2.0)))
  }

  function pressure_l2_norm : Ph(q) -> scalar {
    sqrt(sum_elements(int(T) squared_norm(dof(q, T))))
  }

  function pressure_average_norm : Ph(q) -> scalar {
    sum_elements(int(T) dof(q, T))
  }
}