text:
@doc raw"""
    IdealGlmMhdEquations2D

The ideal compressible GLM-MHD equations in two space dimensions.
"""
mutable struct IdealGlmMhdEquations2D{RealT<:Real} <: AbstractIdealGlmMhdEquations{2, 9}
  gamma::RealT
  c_h::RealT # GLM cleaning speed
end

function IdealGlmMhdEquations2D(gamma; initial_c_h=convert(typeof(gamma), NaN))
  # Use `promote` to ensure that `gamma` and `initial_c_h` have the same type
  IdealGlmMhdEquations2D(promote(gamma, initial_c_h)...)
end

have_nonconservative_terms(::IdealGlmMhdEquations2D) = Val(true)
varnames(::typeof(cons2cons), ::IdealGlmMhdEquations2D) = ("rho", "rho_v1", "rho_v2", "rho_v3", "rho_e", "B1", "B2", "B3", "psi")
varnames(::typeof(cons2prim), ::IdealGlmMhdEquations2D) = ("rho", "v1", "v2", "v3", "p", "B1", "B2", "B3", "psi")
default_analysis_integrals(::IdealGlmMhdEquations2D)  = (entropy_timederivative, Val(:l2_divb), Val(:linf_divb))

# Set initial conditions at physical location `x` for time `t`
"""
    initial_condition_constant(x, t, equations::IdealGlmMhdEquations2D)

A constant initial condition to test free-stream preservation.
"""
function initial_condition_constant(x, t, equations::IdealGlmMhdEquations2D)
  rho = 1.0
  rho_v1 = 0.1
  rho_v2 = -0.2
  rho_v3 = -0.5
  rho_e = 50.0
  B1 = 3.0
  B2 = -1.2
  B3 = 0.5
  psi = 0.0
  return SVector(rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi)
end

"""
    initial_condition_convergence_test(x, t, equations::IdealGlmMhdEquations2D)

An Alfvén wave as smooth initial condition used for convergence tests.
"""
function initial_condition_convergence_test(x, t, equations::IdealGlmMhdEquations2D)
  # smooth Alfvén wave test from Derigs et al. FLASH (2016)
  # domain must be set to [0, 1/cos(α)] x [0, 1/sin(α)], γ = 5/3
  alpha = 0.25*pi
  x_perp = x[1]*cos(alpha) + x[2]*sin(alpha)
  B_perp = 0.1*sin(2.0*pi*x_perp)
  rho = 1.0
  v1 = -B_perp*sin(alpha)
  v2 =  B_perp*cos(alpha)
  v3 = 0.1*cos(2.0*pi*x_perp)
  p = 0.1
  B1 = cos(alpha) + v1
  B2 = sin(alpha) + v2
  B3 = v3
  psi = 0.0
  return prim2cons(SVector(rho, v1, v2, v3, p, B1, B2, B3, psi), equations)
end

"""
    initial_condition_weak_blast_wave(x, t, equations::IdealGlmMhdEquations2D)

A weak blast wave adapted from
- Sebastian Hennemann, Gregor J. Gassner (2020)
  A provably entropy stable subcell shock capturing approach for high order split form DG
  [arXiv: 2008.12044](https://arxiv.org/abs/2008.12044)
"""
function initial_condition_weak_blast_wave(x, t, equations::IdealGlmMhdEquations2D)
  # Adapted MHD version of the weak blast wave from Hennemann & Gassner JCP paper 2020 (Sec. 6.3)
  # Same discontinuity in the velocities but with magnetic fields
  # Set up polar coordinates
  inicenter = (0, 0)
  x_norm = x[1] - inicenter[1]
  y_norm = x[2] - inicenter[2]
  r = sqrt(x_norm^2 + y_norm^2)
  phi = atan(y_norm, x_norm)

  # Calculate primitive variables
  rho = r > 0.5 ? 1.0 : 1.1691
  v1 = r > 0.5 ? 0.0 : 0.1882 * cos(phi)
  v2 = r > 0.5 ? 0.0 : 0.1882 * sin(phi)
  p = r > 0.5 ? 1.0 : 1.245

  return prim2cons(SVector(rho, v1, v2, 0.0, p, 1.0, 1.0, 1.0, 0.0), equations)
end

"""
    initial_condition_orszag_tang(x, t, equations::IdealGlmMhdEquations2D)

The classical Orszag-Tang vortex test case. Here, the setup is taken from
- Dominik Derigs, Gregor J. Gassner, Stefanie Walch & Andrew R. Winters (2018)
  Entropy Stable Finite Volume Approximations for Ideal Magnetohydrodynamics
  [doi: 10.1365/s13291-018-0178-9](https://doi.org/10.1365/s13291-018-0178-9)
"""
function initial_condition_orszag_tang(x, t, equations::IdealGlmMhdEquations2D)
  # setup taken from Derigs et al. DMV article (2018)
  # domain must be [0, 1] x [0, 1], γ = 5/3
  rho = 1.0
  v1 = -sin(2.0*pi*x[2])
  v2 =  sin(2.0*pi*x[1])
  v3 = 0.0
  p = 1.0 / equations.gamma
  B1 = -sin(2.0*pi*x[2]) / equations.gamma
  B2 =  sin(4.0*pi*x[1]) / equations.gamma
  B3 = 0.0
  psi = 0.0
  return prim2cons(SVector(rho, v1, v2, v3, p, B1, B2, B3, psi), equations)
end

"""
    initial_condition_rotor(x, t, equations::IdealGlmMhdEquations2D)

The classical MHD rotor test case. Here, the setup is taken from
- Dominik Derigs, Gregor J. Gassner, Stefanie Walch & Andrew R. Winters (2018)
  Entropy Stable Finite Volume Approximations for Ideal Magnetohydrodynamics
  [doi: 10.1365/s13291-018-0178-9](https://doi.org/10.1365/s13291-018-0178-9)
"""
function initial_condition_rotor(x, t, equations::IdealGlmMhdEquations2D)
  # setup taken from Derigs et al. DMV article (2018)
  # domain must be [0, 1] x [0, 1], γ = 1.4
  dx = x[1] - 0.5
  dy = x[2] - 0.5
  r = sqrt(dx^2 + dy^2)
  f = (0.115 - r)/0.015
  if r <= 0.1
    rho = 10.0
    v1 = -20.0*dy
    v2 = 20.0*dx
  elseif r >= 0.115
    rho = 1.0
    v1 = 0.0
    v2 = 0.0
  else
    rho = 1.0 + 9.0*f
    v1 = -20.0*f*dy
    v2 = 20.0*f*dx
  end
  v3 = 0.0
  p = 1.0
  B1 = 5.0/sqrt(4.0*pi)
  B2 = 0.0
  B3 = 0.0
  psi = 0.0
  return prim2cons(SVector(rho, v1, v2, v3, p, B1, B2, B3, psi), equations)
end

"""
    initial_condition_blast_wave(x, t, equations::IdealGlmMhdEquations2D)

An MHD blast wave taken from
- Dominik Derigs, Gregor J. Gassner, Stefanie Walch & Andrew R. Winters (2018)
  Entropy Stable Finite Volume Approximations for Ideal Magnetohydrodynamics
  [doi: 10.1365/s13291-018-0178-9](https://doi.org/10.1365/s13291-018-0178-9)
"""
function initial_condition_blast_wave(x, t, equations::IdealGlmMhdEquations2D)
  # setup taken from Derigs et al. DMV article (2018)
  # domain must be [-0.5, 0.5] x [-0.5, 0.5], γ = 1.4
  r = sqrt(x[1]^2 + x[2]^2)
  f = (0.1 - r)/0.01
  if r <= 0.09
    p = 1000.0
  elseif r >= 0.1
    p = 0.1
  else
    p = 0.1 + 999.9*f
  end
  rho = 1.0
  v1 = 0.0
  v2 = 0.0
  v3 = 0.0
  B1 = 100.0/sqrt(4.0*pi)
  B2 = 0.0
  B3 = 0.0
  psi = 0.0
  return prim2cons(SVector(rho, v1, v2, v3, p, B1, B2, B3, psi), equations)
end

# Pre-defined source terms should be implemented as
# function source_terms_WHATEVER(u, x, t, equations::IdealGlmMhdEquations2D)

# Calculate 1D flux in for a single point
@inline function flux(u, orientation::Integer, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = u
  v1 = rho_v1/rho
  v2 = rho_v2/rho
  v3 = rho_v3/rho
  kin_en = 0.5 * rho * (v1^2 + v2^2 + v3^2)
  mag_en = 0.5*(B1^2 + B2^2 + B3^2)
  p = (equations.gamma - 1) * (rho_e - kin_en - mag_en - 0.5*psi^2)
  if orientation == 1
    f1 = rho_v1
    f2 = rho_v1*v1 + p + mag_en - B1^2
    f3 = rho_v1*v2 - B1*B2
    f4 = rho_v1*v3 - B1*B3
    f5 = (kin_en + equations.gamma*p/(equations.gamma - 1) + 2*mag_en)*v1 - B1*(v1*B1 + v2*B2 + v3*B3) + equations.c_h*psi*B1
    f6 = equations.c_h*psi
    f7 = v1*B2 - v2*B1
    f8 = v1*B3 - v3*B1
    f9 = equations.c_h*B1
  else # orientation == 2
    f1 = rho_v2
    f2 = rho_v2*v1 - B1*B2
    f3 = rho_v2*v2 + p + mag_en - B2^2
    f4 = rho_v2*v3 - B2*B3
    f5 = (kin_en + equations.gamma*p/(equations.gamma - 1) + 2*mag_en)*v2 - B2*(v1*B1 + v2*B2 + v3*B3) + equations.c_h*psi*B2
    f6 = v2*B1 - v1*B2
    f7 = equations.c_h*psi
    f8 = v2*B3 - v3*B2
    f9 = equations.c_h*B2
  end

  return SVector(f1, f2, f3, f4, f5, f6, f7, f8, f9)
end

# Calculate 1D flux for a single point in the normal direction
# Note, this directional vector is not normalized
@inline function flux(u, normal_direction::AbstractVector, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = u
  v1 = rho_v1/rho
  v2 = rho_v2/rho
  v3 = rho_v3/rho
  kin_en = 0.5 * rho * (v1^2 + v2^2 + v3^2)
  mag_en = 0.5 * (B1^2 + B2^2 + B3^2)
  p = (equations.gamma - 1) * (rho_e - kin_en - mag_en - 0.5*psi^2)

  v_normal = v1 * normal_direction[1] + v2 * normal_direction[2]
  B_normal = B1 * normal_direction[1] + B2 * normal_direction[2]
  rho_v_normal = rho * v_normal

  f1 = rho_v_normal
  f2 = rho_v_normal * v1 - B1 * B_normal + (p + mag_en) * normal_direction[1]
  f3 = rho_v_normal * v2 - B2 * B_normal + (p + mag_en) * normal_direction[2]
  f4 = rho_v_normal * v3 - B3 * B_normal
  f5 = ( (kin_en + equations.gamma*p/(equations.gamma - 1) + 2*mag_en) * v_normal
        - B_normal * (v1*B1 + v2*B2 + v3*B3) + equations.c_h * psi * B_normal )
  f6 = equations.c_h * psi * normal_direction[1] + (v2 * B1 - v1 * B2) * normal_direction[2]
  f7 = (v1 * B2 - v2 * B1) * normal_direction[1] + equations.c_h * psi * normal_direction[2]
  f8 = v_normal * B3 - v3 * B_normal
  f9 = equations.c_h * B_normal

  return SVector(f1, f2, f3, f4, f5, f6, f7, f8, f9)
end

# Calculate the nonconservative terms from Powell and Galilean invariance for the TreeMesh{2}
# OBS! This is scaled by 1/2 becuase it will cancel later with the factor of 2 in dsplit_transposed
@inline function calcflux_twopoint_nonconservative!(f1, f2, u, element,
                                                    equations::IdealGlmMhdEquations2D, dg, cache)
  for j in eachnode(dg), i in eachnode(dg)
    rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = get_node_vars(u, equations, dg, i, j, element)
    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    v3 = rho_v3 / rho

    # Powell nonconservative term: Φ^Pow = (0, B_1, B_2, B_3, v⋅B, v_1, v_2, v_3, 0)
    phi_pow = 0.5 * SVector(0, B1, B2, B3, v1*B1 + v2*B2 + v3*B3, v1, v2, v3, 0)

    # Galilean nonconservative term: Φ^Gal_{1,2} = (0, 0, 0, 0, ψ v_{1,2}, 0, 0, 0, v_{1,2})
    # x-direction
    phi_gal_x = 0.5 * SVector(0, 0, 0, 0, v1*psi, 0, 0, 0, v1)
    # y-direction
    phi_gal_y = 0.5 * SVector(0, 0, 0, 0, v2*psi, 0, 0, 0, v2)

    # add both nonconservative terms into the volume
    for l in eachnode(dg)
      _, _, _, _, _, B1, _, _, psi = get_node_vars(u, equations, dg, l, j, element)
      for v in eachvariable(equations)
        f1[v, l, i, j] += phi_pow[v] * B1 + phi_gal_x[v] * psi
      end
      _, _, _, _, _, _, B2, _, psi = get_node_vars(u, equations, dg, i, l, element)
      for v in eachvariable(equations)
        f2[v, l, i, j] += phi_pow[v] * B2 + phi_gal_y[v] * psi
      end
    end
  end

  return nothing
end

# Calculate the nonconservative terms from Powell and Galilean invariance for UnstructuredQuadMesh
# OBS! This is scaled by 1/2 becuase it will cancel later with the factor of 2 in dsplit_transposed
@inline function calcflux_twopoint_nonconservative!(f1, f2, u, element, contravariant_vectors,
                                                    equations::IdealGlmMhdEquations2D, dg, cache)
   for j in eachnode(dg), i in eachnode(dg)
    rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = get_node_vars(u, equations, dg, i, j, element)
    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    v3 = rho_v3 / rho

    # Powell nonconservative term: Φ^Pow = (0, B_1, B_2, B_3, v⋅B, v_1, v_2, v_3, 0)
    phi_pow = 0.5 * SVector(0, B1, B2, B3, v1*B1 + v2*B2 + v3*B3, v1, v2, v3, 0)

    # Galilean nonconservative term: Φ^Gal_{1,2} = (0, 0, 0, 0, ψ v_{1,2}, 0, 0, 0, v_{1,2})
    # first direction
    Ja11_ij, Ja12_ij = get_contravariant_vector(1, contravariant_vectors, i, j, element)
    phi_gal_x = 0.5 * (Ja11_ij * v1 + Ja12_ij * v2) .* SVector(0, 0, 0, 0, psi, 0, 0, 0, 1)
    # second direction
    Ja21_ij, Ja22_ij = get_contravariant_vector(2, contravariant_vectors, i, j, element)
    phi_gal_y = 0.5 * (Ja21_ij * v1 + Ja22_ij * v2) .* SVector(0, 0, 0, 0, psi, 0, 0, 0, 1)

    # add both nonconservative terms into the volume
    for l in eachnode(dg)
      _, _, _, _, _, B1, B2, _, psi = get_node_vars(u, equations, dg, l, j, element)
      Ja11_lj, Ja12_lj = get_contravariant_vector(1, contravariant_vectors, l, j, element)
      Ja11_avg = 0.5 * (Ja11_ij + Ja11_lj)
      Ja12_avg = 0.5 * (Ja12_ij + Ja12_lj)
      for v in eachvariable(equations)
        f1[v, l, i, j] += phi_pow[v] * ( Ja11_avg * B1 + Ja12_avg * B2) + phi_gal_x[v] * psi
      end
      _, _, _, _, _, B1, B2, _, psi = get_node_vars(u, equations, dg, i, l, element)
      Ja21_il, Ja22_il = get_contravariant_vector(2, contravariant_vectors, i, l, element)
      Ja21_avg = 0.5 * (Ja21_ij + Ja21_il)
      Ja22_avg = 0.5 * (Ja22_ij + Ja22_il)
      for v in eachvariable(equations)
        f2[v, l, i, j] += phi_pow[v] * ( Ja21_avg * B1 + Ja22_avg * B2) + phi_gal_y[v] * psi
      end
    end
  end

  return nothing
end

"""
    flux_derigs_etal(u_ll, u_rr, orientation, equations::IdealGlmMhdEquations2D)

Entropy conserving two-point flux by
- Derigs et al. (2018)
  Ideal GLM-MHD: About the entropy consistent nine-wave magnetic field
  divergence diminishing ideal magnetohydrodynamics equations
  [DOI: 10.1016/j.jcp.2018.03.002](https://doi.org/10.1016/j.jcp.2018.03.002)
"""
function flux_derigs_etal(u_ll, u_rr, orientation::Integer, equations::IdealGlmMhdEquations2D)
  # Unpack left and right states to get velocities, pressure, and inverse temperature (called beta)
  rho_ll, rho_v1_ll, rho_v2_ll, rho_v3_ll, rho_e_ll, B1_ll, B2_ll, B3_ll, psi_ll = u_ll
  rho_rr, rho_v1_rr, rho_v2_rr, rho_v3_rr, rho_e_rr, B1_rr, B2_rr, B3_rr, psi_rr = u_rr

  v1_ll = rho_v1_ll/rho_ll
  v2_ll = rho_v2_ll/rho_ll
  v3_ll = rho_v3_ll/rho_ll
  v1_rr = rho_v1_rr/rho_rr
  v2_rr = rho_v2_rr/rho_rr
  v3_rr = rho_v3_rr/rho_rr
  vel_norm_ll = v1_ll^2 + v2_ll^2 + v3_ll^2
  vel_norm_rr = v1_rr^2 + v2_rr^2 + v3_rr^2
  mag_norm_ll = B1_ll^2 + B2_ll^2 + B3_ll^2
  mag_norm_rr = B1_rr^2 + B2_rr^2 + B3_rr^2
  p_ll = (equations.gamma - 1)*(rho_e_ll - 0.5*rho_ll*vel_norm_ll - 0.5*mag_norm_ll - 0.5*psi_ll^2)
  p_rr = (equations.gamma - 1)*(rho_e_rr - 0.5*rho_rr*vel_norm_rr - 0.5*mag_norm_rr - 0.5*psi_rr^2)
  beta_ll = 0.5*rho_ll/p_ll
  beta_rr = 0.5*rho_rr/p_rr
  # for convenience store v⋅B
  vel_dot_mag_ll = v1_ll*B1_ll + v2_ll*B2_ll + v3_ll*B3_ll
  vel_dot_mag_rr = v1_rr*B1_rr + v2_rr*B2_rr + v3_rr*B3_rr

  # Compute the necessary mean values needed for either direction
  rho_avg  = 0.5*(rho_ll+rho_rr)
  rho_mean = ln_mean(rho_ll,rho_rr)
  beta_mean = ln_mean(beta_ll,beta_rr)
  beta_avg = 0.5*(beta_ll+beta_rr)
  v1_avg = 0.5*(v1_ll+v1_rr)
  v2_avg = 0.5*(v2_ll+v2_rr)
  v3_avg = 0.5*(v3_ll+v3_rr)
  p_mean = 0.5*rho_avg/beta_avg
  B1_avg = 0.5*(B1_ll+B1_rr)
  B2_avg = 0.5*(B2_ll+B2_rr)
  B3_avg = 0.5*(B3_ll+B3_rr)
  psi_avg = 0.5*(psi_ll+psi_rr)
  vel_norm_avg = 0.5*(vel_norm_ll+vel_norm_rr)
  mag_norm_avg = 0.5*(mag_norm_ll+mag_norm_rr)
  vel_dot_mag_avg = 0.5*(vel_dot_mag_ll+vel_dot_mag_rr)

  # Calculate fluxes depending on orientation with specific direction averages
  if orientation == 1
    f1 = rho_mean*v1_avg
    f2 = f1*v1_avg + p_mean + 0.5*mag_norm_avg - B1_avg*B1_avg
    f3 = f1*v2_avg - B1_avg*B2_avg
    f4 = f1*v3_avg - B1_avg*B3_avg
    f6 = equations.c_h*psi_avg
    f7 = v1_avg*B2_avg - v2_avg*B1_avg
    f8 = v1_avg*B3_avg - v3_avg*B1_avg
    f9 = equations.c_h*B1_avg
    # total energy flux is complicated and involves the previous eight components
    psi_B1_avg = 0.5*(B1_ll*psi_ll + B1_rr*psi_rr)
    v1_mag_avg = 0.5*(v1_ll*mag_norm_ll + v1_rr*mag_norm_rr)
    f5 = (f1*0.5*(1/(equations.gamma-1)/beta_mean - vel_norm_avg) + f2*v1_avg + f3*v2_avg +
          f4*v3_avg + f6*B1_avg + f7*B2_avg + f8*B3_avg + f9*psi_avg - 0.5*v1_mag_avg +
          B1_avg*vel_dot_mag_avg - equations.c_h*psi_B1_avg)
  else
    f1 = rho_mean*v2_avg
    f2 = f1*v1_avg - B1_avg*B2_avg
    f3 = f1*v2_avg + p_mean + 0.5*mag_norm_avg - B2_avg*B2_avg
    f4 = f1*v3_avg - B2_avg*B3_avg
    f6 = v2_avg*B1_avg - v1_avg*B2_avg
    f7 = equations.c_h*psi_avg
    f8 = v2_avg*B3_avg - v3_avg*B2_avg
    f9 = equations.c_h*B2_avg
    # total energy flux is complicated and involves the previous eight components
    psi_B2_avg = 0.5*(B2_ll*psi_ll + B2_rr*psi_rr)
    v2_mag_avg = 0.5*(v2_ll*mag_norm_ll + v2_rr*mag_norm_rr)
    f5 = (f1*0.5*(1/(equations.gamma-1)/beta_mean - vel_norm_avg) + f2*v1_avg + f3*v2_avg +
          f4*v3_avg + f6*B1_avg + f7*B2_avg + f8*B3_avg + f9*psi_avg - 0.5*v2_mag_avg +
          B2_avg*vel_dot_mag_avg - equations.c_h*psi_B2_avg)
  end

  return SVector(f1, f2, f3, f4, f5, f6, f7, f8, f9)
end

"""
    flux_hindenlang(u_ll, u_rr, orientation, equations::IdealGlmMhdEquations2D)

Entropy conserving and kinetic energy preserving two-point flux of
Hindenlang (2019), extending [`flux_ranocha`](@ref) to the MHD equations.

## References
- Hindenlang (2019)
  A new entropy conservative two-point flux for ideal MHD equations derived from
  first principles.
  Presented at HONOM 2019: European workshop on high order numerical methods
  for evolutionary PDEs, theory and applications
- Ranocha (2018)
  Generalised Summation-by-Parts Operators and Entropy Stability of Numerical Methods
  for Hyperbolic Balance Laws
  [PhD thesis, TU Braunschweig](https://cuvillier.de/en/shop/publications/7743)
- Ranocha (2020)
  Entropy Conserving and Kinetic Energy Preserving Numerical Methods for
  the Euler Equations Using Summation-by-Parts Operators
  [Proceedings of ICOSAHOM 2018](https://doi.org/10.1007/978-3-030-39647-3_42)
"""
@inline function flux_hindenlang(u_ll, u_rr, orientation::Integer, equations::IdealGlmMhdEquations2D)
  # Unpack left and right states to get velocities and pressure
  rho_ll, rho_v1_ll, rho_v2_ll, rho_v3_ll, rho_e_ll, B1_ll, B2_ll, B3_ll, psi_ll = u_ll
  rho_rr, rho_v1_rr, rho_v2_rr, rho_v3_rr, rho_e_rr, B1_rr, B2_rr, B3_rr, psi_rr = u_rr

  v1_ll = rho_v1_ll / rho_ll
  v2_ll = rho_v2_ll / rho_ll
  v3_ll = rho_v3_ll / rho_ll
  v1_rr = rho_v1_rr / rho_rr
  v2_rr = rho_v2_rr / rho_rr
  v3_rr = rho_v3_rr / rho_rr
  vel_norm_ll = v1_ll^2 + v2_ll^2 + v3_ll^2
  vel_norm_rr = v1_rr^2 + v2_rr^2 + v3_rr^2
  mag_norm_ll = B1_ll^2 + B2_ll^2 + B3_ll^2
  mag_norm_rr = B1_rr^2 + B2_rr^2 + B3_rr^2
  p_ll = (equations.gamma - 1) * (rho_e_ll - 0.5*rho_ll*vel_norm_ll - 0.5*mag_norm_ll - 0.5*psi_ll^2)
  p_rr = (equations.gamma - 1) * (rho_e_rr - 0.5*rho_rr*vel_norm_rr - 0.5*mag_norm_rr - 0.5*psi_rr^2)

  # Compute the necessary mean values needed for either direction
  rho_mean   = ln_mean(rho_ll, rho_rr)
  rho_p_mean = ln_mean(rho_ll / p_ll, rho_rr / p_rr)
  v1_avg  = 0.5 * ( v1_ll +  v1_rr)
  v2_avg  = 0.5 * ( v2_ll +  v2_rr)
  v3_avg  = 0.5 * ( v3_ll +  v3_rr)
  p_avg   = 0.5 * (  p_ll +   p_rr)
  psi_avg = 0.5 * (psi_ll + psi_rr)
  velocity_square_avg = 0.5 * (v1_ll * v1_rr + v2_ll * v2_rr + v3_ll * v3_rr)
  magnetic_square_avg = 0.5 * (B1_ll * B1_rr + B2_ll * B2_rr + B3_ll * B3_rr)

  # Calculate fluxes depending on orientation with specific direction averages
  if orientation == 1
    f1 = rho_mean * v1_avg
    f2 = f1 * v1_avg + p_avg + magnetic_square_avg - 0.5 * (B1_ll * B1_rr + B1_rr * B1_ll)
    f3 = f1 * v2_avg                               - 0.5 * (B1_ll * B2_rr + B1_rr * B2_ll)
    f4 = f1 * v3_avg                               - 0.5 * (B1_ll * B3_rr + B1_rr * B3_ll)
    #f5 below
    f6 = equations.c_h * psi_avg
    f7 = 0.5 * (v1_ll * B2_ll - v2_ll * B1_ll + v1_rr * B2_rr - v2_rr * B1_rr)
    f8 = 0.5 * (v1_ll * B3_ll - v3_ll * B1_ll + v1_rr * B3_rr - v3_rr * B1_rr)
    f9 = equations.c_h * 0.5 * (B1_ll + B1_rr)
    # total energy flux is complicated and involves the previous components
    f5 = ( f1 * ( velocity_square_avg + 1 / ((equations.gamma-1) * rho_p_mean) )
          + 0.5 * ( p_ll * v1_rr +  p_rr * v1_ll)
          + 0.5 * (v1_ll * B2_ll * B2_rr + v1_rr * B2_rr * B2_ll)
          + 0.5 * (v1_ll * B3_ll * B3_rr + v1_rr * B3_rr * B3_ll)
          - 0.5 * (v2_ll * B1_ll * B2_rr + v2_rr * B1_rr * B2_ll)
          - 0.5 * (v3_ll * B1_ll * B3_rr + v3_rr * B1_rr * B3_ll)
          + 0.5 * equations.c_h * (B1_ll * psi_rr + B1_rr * psi_ll) )
  else # orientation == 2
    f1 = rho_mean * v2_avg
    f2 = f1 * v1_avg                               - 0.5 * (B2_ll * B1_rr + B2_rr * B1_ll)
    f3 = f1 * v2_avg + p_avg + magnetic_square_avg - 0.5 * (B2_ll * B2_rr + B2_rr * B2_ll)
    f4 = f1 * v3_avg                               - 0.5 * (B2_ll * B3_rr + B2_rr * B3_ll)
    #f5 below
    f6 = 0.5 * (v2_ll * B1_ll - v1_ll * B2_ll + v2_rr * B1_rr - v1_rr * B2_rr)
    f7 = equations.c_h * psi_avg
    f8 = 0.5 * (v2_ll * B3_ll - v3_ll * B2_ll + v2_rr * B3_rr - v3_rr * B2_rr)
    f9 = equations.c_h * 0.5 * (B2_ll + B2_rr)
    # total energy flux is complicated and involves the previous components
    f5 = ( f1 * ( velocity_square_avg + 1 / ((equations.gamma-1) * rho_p_mean) )
          + 0.5 * ( p_ll * v2_rr +  p_rr * v2_ll)
          + 0.5 * (v2_ll * B1_ll * B1_rr + v2_rr * B1_rr * B1_ll)
          + 0.5 * (v2_ll * B3_ll * B3_rr + v2_rr * B3_rr * B3_ll)
          - 0.5 * (v1_ll * B2_ll * B1_rr + v1_rr * B2_rr * B1_ll)
          - 0.5 * (v3_ll * B2_ll * B3_rr + v3_rr * B2_rr * B3_ll)
          + 0.5 * equations.c_h * (B2_ll * psi_rr + B2_rr * psi_ll) )
  end

  return SVector(f1, f2, f3, f4, f5, f6, f7, f8, f9)
end

# Calculate maximum wave speed for local Lax-Friedrichs-type dissipation
@inline function max_abs_speed_naive(u_ll, u_rr, orientation::Integer, equations::IdealGlmMhdEquations2D)
  rho_ll, rho_v1_ll, rho_v2_ll, rho_v3_ll, rho_e_ll, B1_ll, B2_ll, B3_ll, psi_ll = u_ll
  rho_rr, rho_v1_rr, rho_v2_rr, rho_v3_rr, rho_e_rr, B1_rr, B2_rr, B3_rr, psi_rr = u_rr

  # Calculate velocities and fast magnetoacoustic wave speeds
  # left
  v1_ll = rho_v1_ll / rho_ll
  v2_ll = rho_v2_ll / rho_ll
  v3_ll = rho_v3_ll / rho_ll
  v_mag_ll = sqrt(v1_ll^2 + v2_ll^2 + v3_ll^2)
  cf_ll = calc_fast_wavespeed(u_ll, orientation, equations)
  # right
  v1_rr = rho_v1_rr / rho_rr
  v2_rr = rho_v2_rr / rho_rr
  v3_rr = rho_v3_rr / rho_rr
  v_mag_rr = sqrt(v1_rr^2 + v2_rr^2 + v3_rr^2)
  cf_rr = calc_fast_wavespeed(u_rr, orientation, equations)

  λ_max = max(v_mag_ll, v_mag_rr) + max(cf_ll, cf_rr)
end

@inline function max_abs_speed_naive(u_ll, u_rr, normal_direction::AbstractVector, equations::IdealGlmMhdEquations2D)

  norm_ = norm(normal_direction)
  # Normalize the vector without using `normalize` since we need to multiply by the `norm_` later
  normal_vector = normal_direction / norm_
  u_ll_rotated = rotate_to_x(u_ll, normal_vector, equations)
  u_rr_rotated = rotate_to_x(u_rr, normal_vector, equations)

  return max_abs_speed_naive(u_ll_rotated, u_rr_rotated, 1, equations) * norm_
end

"""
    min_max_speed_naive(u_ll, u_rr, orientation, equations::IdealGlmMhdEquations2D)

Calculate minimum and maximum wave speeds for HLL-type fluxes as in
- Li (2005)
  An HLLC Riemann solver for magneto-hydrodynamics
  [DOI: 10.1016/j.jcp.2004.08.020](https://doi.org/10.1016/j.jcp.2004.08.020)
"""
@inline function min_max_speed_naive(u_ll, u_rr, orientation::Integer, equations::IdealGlmMhdEquations2D)
  rho_ll, rho_v1_ll, rho_v2_ll, _, _, _, _, _, _ = u_ll
  rho_rr, rho_v1_rr, rho_v2_rr, _, _, _, _, _, _ = u_rr

  # Calculate primitive velocity variables
  v1_ll = rho_v1_ll/rho_ll
  v2_ll = rho_v2_ll/rho_ll

  v1_rr = rho_v1_rr/rho_rr
  v2_rr = rho_v2_rr/rho_rr

  # Approximate the left-most and right-most eigenvalues in the Riemann fan
  if orientation == 1 # x-direction
    c_f_ll = calc_fast_wavespeed(u_ll, orientation, equations)
    c_f_rr = calc_fast_wavespeed(u_rr, orientation, equations)
    vel_roe, c_f_roe = calc_fast_wavespeed_roe(u_ll, u_rr, orientation, equations)
    λ_min = min(v1_ll - c_f_ll, vel_roe - c_f_roe)
    λ_max = max(v1_rr + c_f_rr, vel_roe + c_f_roe)
  else # y-direction
    c_f_ll = calc_fast_wavespeed(u_ll, orientation, equations)
    c_f_rr = calc_fast_wavespeed(u_rr, orientation, equations)
    vel_roe, c_f_roe = calc_fast_wavespeed_roe(u_ll, u_rr, orientation, equations)
    λ_min = min(v2_ll - c_f_ll, vel_roe - c_f_roe)
    λ_max = max(v2_rr + c_f_rr, vel_roe + c_f_roe)
  end

  return λ_min, λ_max
end

# Very naive way to approximate the edges of the Riemann fan in the normal direction
@inline function min_max_speed_naive(u_ll, u_rr, normal_direction::AbstractVector,
                                     equations::IdealGlmMhdEquations2D)
  rho_ll, rho_v1_ll, rho_v2_ll, _, _, _, _, _, _ = u_ll
  rho_rr, rho_v1_rr, rho_v2_rr, _, _, _, _, _, _ = u_rr

  # Calculate primitive velocity variables
  v1_ll = rho_v1_ll/rho_ll
  v2_ll = rho_v2_ll/rho_ll

  v1_rr = rho_v1_rr/rho_rr
  v2_rr = rho_v2_rr/rho_rr

  # Compute wave speed estimates in each direction. Requires rotation because
  # the fast magnetoacoustic wave speed has a nonlinear dependence on the direction
  norm_ = norm(normal_direction)
  # Normalize the vector without using `normalize` since we need to multiply by the `norm_` later
  normal_vector = normal_direction / norm_

  # Rotate the velocities
  v_normal_ll = v1_ll * normal_vector[1] + v2_ll * normal_vector[2]
  v_normal_rr = v1_rr * normal_vector[1] + v2_rr * normal_vector[2]

  u_ll_rotated = rotate_to_x(u_ll, normal_vector, equations)
  u_rr_rotated = rotate_to_x(u_rr, normal_vector, equations)

  c_f_ll_rotated = calc_fast_wavespeed(u_ll_rotated, 1, equations)
  c_f_rr_rotated = calc_fast_wavespeed(u_rr_rotated, 1, equations)
  v_roe_rotated, c_f_roe_rotated = calc_fast_wavespeed_roe(u_ll_rotated, u_rr_rotated, 1, equations)

  # Estimate the min/max eigenvalues in the normal direction
  λ_min = min(v_normal_ll - c_f_ll_rotated, v_roe_rotated - c_f_roe_rotated) * norm_
  λ_max = max(v_normal_rr + c_f_rr_rotated, v_roe_rotated + c_f_roe_rotated) * norm_

  return λ_min, λ_max
end

# Called inside `FluxRotated` in `numerical_fluxes.jl` so the direction
# has been normalized prior to this rotation of the state vector
@inline function rotate_to_x(u, normal_vector, equations::IdealGlmMhdEquations2D)
  # cos and sin of the angle between the x-axis and the normalized normal_vector are
  # the normalized vector's x and y coordinates respectively (see unit circle).
  c = normal_vector[1]
  s = normal_vector[2]

  # Apply the 2D rotation matrix with normal and tangent directions of the form
  # [ 1   0    0   0   0   0    0   0   0;
  #   0  n_1  n_2  0   0   0    0   0   0;
  #   0  t_1  t_2  0   0   0    0   0   0;
  #   0   0    0   1   0   0    0   0   0;
  #   0   0    0   0   1   0    0   0   0;
  #   0   0    0   0   0  n_1  n_2  0   0;
  #   0   0    0   0   0  t_1  t_2  0   0;
  #   0   0    0   0   0   0    0   1   0;
  #   0   0    0   0   0   0    0   0   1 ]
  # where t_1 = -n_2 and t_2 = n_1.
  # Note for IdealGlmMhdEquations2D only the velocities and magnetic field variables rotate

  return SVector(u[1],
                  c * u[2] + s * u[3],
                 -s * u[2] + c * u[3],
                 u[4],
                 u[5],
                  c * u[6] + s * u[7],
                 -s * u[6] + c * u[7],
                 u[8],
                 u[9])
end

# Called inside `FluxRotated` in `numerical_fluxes.jl` so the direction
# has been normalized prior to this back-rotation of the state vector
@inline function rotate_from_x(u, normal_vector, equations::IdealGlmMhdEquations2D)
  # cos and sin of the angle between the x-axis and the normalized normal_vector are
  # the normalized vector's x and y coordinates respectively (see unit circle).
  c = normal_vector[1]
  s = normal_vector[2]

  # Apply the 2D back-rotation matrix with normal and tangent directions of the form
  # [ 1   0    0   0   0   0    0   0   0;
  #   0  n_1  t_1  0   0   0    0   0   0;
  #   0  n_2  t_2  0   0   0    0   0   0;
  #   0   0    0   1   0   0    0   0   0;
  #   0   0    0   0   1   0    0   0   0;
  #   0   0    0   0   0  n_1  t_1  0   0;
  #   0   0    0   0   0  n_2  t_2  0   0;
  #   0   0    0   0   0   0    0   1   0;
  #   0   0    0   0   0   0    0   0   1 ]
  # where t_1 = -n_2 and t_2 = n_1.
  # Note for IdealGlmMhdEquations2D the velocities and magnetic field variables back-rotate

  return SVector(u[1],
                 c * u[2] - s * u[3],
                 s * u[2] + c * u[3],
                 u[4],
                 u[5],
                 c * u[6] - s * u[7],
                 s * u[6] + c * u[7],
                 u[8],
                 u[9])
end

"""
    noncons_interface_flux(u_left, u_right, orientation, mode, equations::IdealGlmMhdEquations2D)

Strong form of non-conservative flux on a surface (Powell and GLM terms)
```math
phi^L 1/2 (B^L + B^R)_{normal} - phi^L B^L+{normal} = phi^L 1/2 (B^R - B^L)_{normal}
```
!!! note
    The non-conservative interface flux depends on the discretization. Following "modes" are available:
    * `:weak`: 'weak' formulation of split DG already includes the contribution
      ``-1/2 (phi^L B^L_{normal})`` so this mode only adds ``1/2 (phi^L B^R_{normal})``,
      analogously for the Galilean nonconservative term
    * `:whole`: This mode adds the whole non-conservative term: phi^L 1/2 (B^R-B^L)
    * `:inner`: This mode adds the split-form DG volume integral contribution: This is equivalent to
      ``(2)-(1) - 1/2 (phi^L B^L)``
!!! warning
    This is non-unique along an interface! The normal direction is super important.

For details see Section 4 of the paper
- Bohm et al. (2018)
  An entropy stable nodal discontinuous Galerkin method for the resistive MHD equations.
  Part I: Theory and numerical verification
  [DOI: 10.1016/j.jcp.2018.06.027](https://doi.org/10.1016/j.jcp.2018.06.027)
"""
@inline function noncons_interface_flux(u_left, u_right, orientation, mode,
                                        equations::IdealGlmMhdEquations2D)
  rho_ll, rho_v1_ll, rho_v2_ll, rho_v3_ll, _, B1_ll, B2_ll, B3_ll, psi_ll = u_left
  _, _, _, _, _, B1_rr, B2_rr, _, psi_rr = u_right

  # extract velocites from the left
  v1_ll  = rho_v1_ll / rho_ll
  v2_ll  = rho_v2_ll / rho_ll
  v3_ll  = rho_v3_ll / rho_ll
  v_dot_B_ll = v1_ll*B1_ll + v2_ll*B2_ll + v3_ll*B3_ll
  # extract magnetic field variable from the right and set the normal velocity
  # Note, both depend upon the orientation and need psi_rr
  if mode==:weak
    if orientation == 1 # x-direction
      v_normal = v1_ll
      B_normal = B1_rr
    else # y-direction
      v_normal = v2_ll
      B_normal = B2_rr
    end
    psi_norm = psi_rr
  elseif mode==:whole
    if orientation == 1 # x-direction
      v_normal = v1_ll
      B_normal = B1_rr - B1_ll
    else # y-direction
      v_normal = v2_ll
      B_normal = B2_rr - B2_ll
    end
    psi_norm = psi_rr - psi_ll
  else #mode==:inner
    if orientation == 1 # x-direction
      v_normal = v1_ll
      B_normal =-B1_ll
    else # y-direction
      v_normal = v2_ll
      B_normal =-B2_ll
    end
    psi_norm =-psi_ll
  end

  # compute the nonconservative flux: Powell (with B_normal) and Galilean (with v_normal)
  noncons2 = 0.5 * B_normal * B1_ll
  noncons3 = 0.5 * B_normal * B2_ll
  noncons4 = 0.5 * B_normal * B3_ll
  noncons5 = 0.5 * B_normal * v_dot_B_ll + 0.5 * v_normal * psi_ll * psi_norm
  noncons6 = 0.5 * B_normal * v1_ll
  noncons7 = 0.5 * B_normal * v2_ll
  noncons8 = 0.5 * B_normal * v3_ll
  noncons9 = 0.5 * v_normal * psi_norm

  return SVector(0, noncons2, noncons3, noncons4, noncons5, noncons6, noncons7, noncons8, noncons9)
end

# Compute surface nonconservative "flux" computation in the normal direction (2D version)
# Note, due to the non-uniqueness of this term we cannot use any fancy rotation tricks.
@inline function noncons_interface_flux(u_left, u_right, normal_direction::AbstractVector, mode,
                                        equations::IdealGlmMhdEquations2D)
  @assert mode === :weak "only :weak version of nonconservative coupling is available for curved MHD"

  rho_ll, rho_v1_ll, rho_v2_ll, rho_v3_ll, _, B1_ll, B2_ll, B3_ll, psi_ll = u_left
  _, _, _, _, _, B1_rr, B2_rr, _, psi_rr = u_right

  norm_ = norm(normal_direction)
  # Normalize the vector without using `normalize` since we need to multiply by the `norm_` later
  normal_vector = normal_direction / norm_

  # extract velocites from the left
  v1_ll  = rho_v1_ll / rho_ll
  v2_ll  = rho_v2_ll / rho_ll
  v3_ll  = rho_v3_ll / rho_ll
  v_dot_B_ll = v1_ll*B1_ll + v2_ll*B2_ll + v3_ll*B3_ll
  # extract magnetic field variable from the right and set the normal velocity
  v_normal = normal_vector[1] * v1_ll + normal_vector[2] * v2_ll
  B_normal = normal_vector[1] * B1_rr + normal_vector[2] * B2_rr
  # compute the nonconservative flux: Powell (with B_normal) and Galilean (with v_normal)
  noncons2 = 0.5 * B_normal * B1_ll
  noncons3 = 0.5 * B_normal * B2_ll
  noncons4 = 0.5 * B_normal * B3_ll
  noncons5 = 0.5 * B_normal * v_dot_B_ll + 0.5 * v_normal * psi_ll * psi_rr
  noncons6 = 0.5 * B_normal * v1_ll
  noncons7 = 0.5 * B_normal * v2_ll
  noncons8 = 0.5 * B_normal * v3_ll
  noncons9 = 0.5 * v_normal * psi_rr

  return SVector(0, noncons2, noncons3, noncons4, noncons5, noncons6, noncons7, noncons8, noncons9) * norm_
end

@inline function max_abs_speeds(u, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, _ = u
  v1 = rho_v1 / rho
  v2 = rho_v2 / rho
  v3 = rho_v3 / rho
  cf_x_direction = calc_fast_wavespeed(u, 1, equations)
  cf_y_direction = calc_fast_wavespeed(u, 2, equations)

  return abs(v1) + cf_x_direction, abs(v2) + cf_y_direction
end

# Convert conservative variables to primitive
function cons2prim(u, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = u

  v1 = rho_v1 / rho
  v2 = rho_v2 / rho
  v3 = rho_v3 / rho
  p = (equations.gamma - 1) * (rho_e - 0.5*rho*(v1^2 + v2^2 + v3^2) - 0.5*(B1^2 + B2^2 + B3^2) - 0.5*psi^2)

  return SVector(rho, v1, v2, v3, p, B1, B2, B3, psi)
end

# Convert conservative variables to entropy
@inline function cons2entropy(u, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = u

  v1 = rho_v1 / rho
  v2 = rho_v2 / rho
  v3 = rho_v3 / rho
  v_square = v1^2 + v2^2 + v3^2
  p = (equations.gamma - 1) * (rho_e - 0.5*rho*v_square - 0.5*(B1^2 + B2^2 + B3^2) - 0.5*psi^2)
  s = log(p) - equations.gamma*log(rho)
  rho_p = rho / p

  w1 = (equations.gamma - s) / (equations.gamma-1) - 0.5 * rho_p * v_square
  w2 = rho_p * v1
  w3 = rho_p * v2
  w4 = rho_p * v3
  w5 = -rho_p
  w6 = rho_p * B1
  w7 = rho_p * B2
  w8 = rho_p * B3
  w9 = rho_p * psi

  return SVector(w1, w2, w3, w4, w5, w6, w7, w8, w9)
end

# Convert primitive to conservative variables
@inline function prim2cons(prim, equations::IdealGlmMhdEquations2D)
  rho, v1, v2, v3, p, B1, B2, B3, psi = prim

  rho_v1 = rho * v1
  rho_v2 = rho * v2
  rho_v3 = rho * v3
  rho_e = p/(equations.gamma-1) + 0.5 * (rho_v1*v1 + rho_v2*v2 + rho_v3*v3) +
                                 0.5 * (B1^2 + B2^2 + B3^2) + 0.5 * psi^2

  return SVector(rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi)
end

@inline function density(u, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = u
  return rho
end

@inline function pressure(u, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = u
  p = (equations.gamma - 1)*(rho_e - 0.5 * (rho_v1^2 + rho_v2^2 + rho_v3^2) / rho
                                   - 0.5 * (B1^2 + B2^2 + B3^2)
                                   - 0.5 * psi^2)
  return p
end

@inline function density_pressure(u, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = u
  p = (equations.gamma - 1)*(rho_e - 0.5 * (rho_v1^2 + rho_v2^2 + rho_v3^2) / rho
                                   - 0.5 * (B1^2 + B2^2 + B3^2)
                                   - 0.5 * psi^2)
  return rho * p
end

# Compute the fastest wave speed for ideal MHD equations: c_f, the fast magnetoacoustic eigenvalue
@inline function calc_fast_wavespeed(cons, direction, equations::IdealGlmMhdEquations2D)
  rho, rho_v1, rho_v2, rho_v3, rho_e, B1, B2, B3, psi = cons
  v1 = rho_v1 / rho
  v2 = rho_v2 / rho
  v3 = rho_v3 / rho
  v_mag = sqrt(v1^2 + v2^2 + v3^2)
  p = (equations.gamma - 1)*(rho_e - 0.5*rho*v_mag^2 - 0.5*(B1^2 + B2^2 + B3^2) - 0.5*psi^2)
  a_square = equations.gamma * p / rho
  sqrt_rho = sqrt(rho)
  b1 = B1 / sqrt_rho
  b2 = B2 / sqrt_rho
  b3 = B3 / sqrt_rho
  b_square = b1^2 + b2^2 + b3^2
  if direction == 1 # x-direction
    c_f = sqrt(0.5*(a_square + b_square) + 0.5*sqrt((a_square + b_square)^2 - 4.0*a_square*b1^2))
  else
    c_f = sqrt(0.5*(a_square + b_square) + 0.5*sqrt((a_square + b_square)^2 - 4.0*a_square*b2^2))
  end
  return c_f
end

"""
    calc_fast_wavespeed_roe(u_ll, u_rr, direction, equations::IdealGlmMhdEquations2D)

Compute the fast magnetoacoustic wave speed using Roe averages
as given by
- Cargo and Gallice (1997)
  Roe Matrices for Ideal MHD and Systematic Construction
  of Roe Matrices for Systems of Conservation Laws
  [DOI: 10.1006/jcph.1997.5773](https://doi.org/10.1006/jcph.1997.5773)
"""
@inline function calc_fast_wavespeed_roe(u_ll, u_rr, direction, equations::IdealGlmMhdEquations2D)
  rho_ll, rho_v1_ll, rho_v2_ll, rho_v3_ll, rho_e_ll, B1_ll, B2_ll, B3_ll, psi_ll = u_ll
  rho_rr, rho_v1_rr, rho_v2_rr, rho_v3_rr, rho_e_rr, B1_rr, B2_rr, B3_rr, psi_rr = u_rr

  # Calculate primitive variables
  v1_ll = rho_v1_ll/rho_ll
  v2_ll = rho_v2_ll/rho_ll
  v3_ll = rho_v3_ll/rho_ll
  vel_norm_ll = v1_ll^2 + v2_ll^2 + v3_ll^2
  mag_norm_ll = B1_ll^2 + B2_ll^2 + B3_ll^2
  p_ll = (equations.gamma - 1)*(rho_e_ll - 0.5*rho_ll*vel_norm_ll - 0.5*mag_norm_ll - 0.5*psi_ll^2)

  v1_rr = rho_v1_rr/rho_rr
  v2_rr = rho_v2_rr/rho_rr
  v3_rr = rho_v3_rr/rho_rr
  vel_norm_rr = v1_rr^2 + v2_rr^2 + v3_rr^2
  mag_norm_rr = B1_rr^2 + B2_rr^2 + B3_rr^2
  p_rr = (equations.gamma - 1)*(rho_e_rr - 0.5*rho_rr*vel_norm_rr - 0.5*mag_norm_rr - 0.5*psi_rr^2)

  # compute total pressure which is thermal + magnetic pressures
  p_total_ll = p_ll + 0.5*mag_norm_ll
  p_total_rr = p_rr + 0.5*mag_norm_rr

  # compute the Roe density averages
  sqrt_rho_ll = sqrt(rho_ll)
  sqrt_rho_rr = sqrt(rho_rr)
  inv_sqrt_rho_add  = 1.0 / (sqrt_rho_ll + sqrt_rho_rr)
  inv_sqrt_rho_prod = 1.0 / (sqrt_rho_ll * sqrt_rho_rr)
  rho_ll_roe =  sqrt_rho_ll * inv_sqrt_rho_add
  rho_rr_roe =  sqrt_rho_rr * inv_sqrt_rho_add
  # Roe averages
  # velocities and magnetic fields
  v1_roe = v1_ll * rho_ll_roe + v1_rr * rho_rr_roe
  v2_roe = v2_ll * rho_ll_roe + v2_rr * rho_rr_roe
  v3_roe = v3_ll * rho_ll_roe + v3_rr * rho_rr_roe
  B1_roe = B1_ll * rho_ll_roe + B1_rr * rho_rr_roe
  B2_roe = B2_ll * rho_ll_roe + B2_rr * rho_rr_roe
  B3_roe = B3_ll * rho_ll_roe + B3_rr * rho_rr_roe
  # enthalpy
  H_ll  = (rho_e_ll + p_total_ll) / rho_ll
  H_rr  = (rho_e_rr + p_total_rr) / rho_rr
  H_roe = H_ll * rho_ll_roe + H_rr * rho_rr_roe
  # temporary vairable see equations (4.12) in Cargo and Gallice
  X = 0.5 * ( (B1_ll - B1_rr)^2 + (B2_ll - B2_rr)^2 + (B3_ll - B3_rr)^2 ) * inv_sqrt_rho_add^2
  # averaged components needed to compute c_f, the fast magnetoacoustic wave speed
  b_square_roe = (B1_roe^2 + B2_roe^2 + B3_roe^2) * inv_sqrt_rho_prod # scaled magnectic sum
  a_square_roe = ((2.0 - equations.gamma) * X +
                 (equations.gamma -1.0) * (H_roe - 0.5*(v1_roe^2 + v2_roe^2 + v3_roe^2) -
                                          b_square_roe)) # acoustic speed
  # finally compute the average wave speed and set the output velocity (depends on orientation)
  if direction == 1 # x-direction
    c_a_roe = B1_roe^2 * inv_sqrt_rho_prod # (squared) Alfvén wave speed
    a_star_roe = sqrt( (a_square_roe + b_square_roe)^2 - 4.0 * a_square_roe * c_a_roe )
    c_f_roe = sqrt( 0.5 * (a_square_roe + b_square_roe + a_star_roe) )
    vel_out_roe = v1_roe
  else # y-direction
    c_a_roe = B2_roe^2 * inv_sqrt_rho_prod # (squared) Alfvén wave speed
    a_star_roe = sqrt( (a_square_roe + b_square_roe)^2 - 4.0 * a_square_roe * c_a_roe )
    c_f_roe = sqrt( 0.5 * (a_square_roe + b_square_roe + a_star_roe) )
    vel_out_roe = v2_roe
  end

  return vel_out_roe, c_f_roe
end

# Calculate thermodynamic entropy for a conservative state `cons`
@inline function entropy_thermodynamic(cons, equations::IdealGlmMhdEquations2D)
  # Pressure
  p = (equations.gamma - 1) * (cons[5] - 1/2 * (cons[2]^2 + cons[3]^2 + cons[4]^2) / cons[1]
                                       - 1/2 * (cons[6]^2 + cons[7]^2 + cons[8]^2)
                                       - 1/2 * cons[9]^2)

  # Thermodynamic entropy
  s = log(p) - equations.gamma*log(cons[1])

  return s
end

# Calculate mathematical entropy for a conservative state `cons`
@inline function entropy_math(cons, equations::IdealGlmMhdEquations2D)
  S = -entropy_thermodynamic(cons, equations) * cons[1] / (equations.gamma - 1)

  return S
end

# Default entropy is the mathematical entropy
@inline entropy(cons, equations::IdealGlmMhdEquations2D) = entropy_math(cons, equations)

# Calculate total energy for a conservative state `cons`
@inline energy_total(cons, ::IdealGlmMhdEquations2D) = cons[5]

# Calculate kinetic energy for a conservative state `cons`
@inline function energy_kinetic(cons, equations::IdealGlmMhdEquations2D)
  return 0.5 * (cons[2]^2 + cons[3]^2 + cons[4]^2)/cons[1]
end

# Calculate the magnetic energy for a conservative state `cons'.
#  OBS! For non-dinmensional form of the ideal MHD magnetic pressure ≡ magnetic energy
@inline function energy_magnetic(cons, ::IdealGlmMhdEquations2D)
  return 0.5 * (cons[6]^2 + cons[7]^2 + cons[8]^2)
end

# Calculate internal energy for a conservative state `cons`
@inline function energy_internal(cons, equations::IdealGlmMhdEquations2D)
  return (energy_total(cons, equations)
          - energy_kinetic(cons, equations)
          - energy_magnetic(cons, equations)
          - cons[9]^2 / 2)
end

# Calcluate the cross helicity (\vec{v}⋅\vec{B}) for a conservative state `cons'
@inline function cross_helicity(cons, ::IdealGlmMhdEquations2D)
  return (cons[2]*cons[6] + cons[3]*cons[7] + cons[4]*cons[8]) / cons[1]
end


text:
module Example

using Twinkle, Twinkle.FlexUI
using Colors
using StaticArrays
using GeometryBasics
import Makie, GLMakie

println("Start [$(splitext(basename(@__FILE__))[1])]")

app = App()
prop!(app, :title, "Twinkle Example - Makie")

#---------------------------------------------------------------
# Define Variables
#---------------------------------------------------------------

xscale = addVariable!(app, Variable(name="xscale", type="flota64",value=1))
yscale = addVariable!(app, Variable(name="yscale", type="flota64",value=1))

image = addVariable!(app, Variable(
    name="image",
    type="image",
    value="",
))

#---------------------------------------------------------------
# Define Controls
#---------------------------------------------------------------

ui = VContainer(
    H1Label("Value Controls"),
    Slider(
        text="X Scale [1-10]",
        trailing_text="[\$()]",
        min=1,
        max=10,
        value=1,
        variable="xscale"
    ),
    Slider(
        text="Y Scale [1-10]",
        trailing_text="[\$()]",
        min=1,
        max=10,
        value=1,
        variable="yscale"
    ),
    H1Label("Static Image Example"),
    Image(
        source="\$(image)",
    ),
    H1Label("Image Viewer allowing Pan (left-drag) and Zoom (wheel)"),
    PanZoom(
        style="width: 100%; height=400px;",
        content = Container(
            direction = "row warp",
            children = [
                Image(
                    source="\$(image)",
                ),
            ]
        ),
    ),

)
FlexUI.controls!(app, ui)

#---------------------------------------------------------------
# the render function
#---------------------------------------------------------------
function render()
    Makie.inline!(true);

    f = Makie.Figure()
    Makie.Axis(f[1, 1])

    xs = LinRange(0, 10, 100)
    ys = LinRange(0, 15, 100)
    zs = [cos(xscale[]*x) * sin(yscale[]*y) for x in xs, y in ys]

    Makie.contour!(xs, ys, zs)

    # Makie.display(f)
    image["png"] = f;

end
renderFunction!(app, render)

#---------------------------------------------------------------
# Run the application
#---------------------------------------------------------------
# render()
run(app)

println("End [$(splitext(basename(@__FILE__))[1])]")

end # module

ext: jl


text:
@testset "HMC Sampling" begin
    seed!(42)
    L = 3
    D = 10
    N = 20
    nSamples = 10
    b = 5
    i = HMCSampling(; nBurnin=0, thinning=1)

    @test repr(i) == "Hamilton Monte Carlo Sampler"

    @test AGP.ρ(i) == 1
    @test AGP.is_stochastic(i) == false

    i = AGP.init_sampler!(i, L, N, nSamples, false)

    @test i.sample_store == zeros(Float64, nSamples, N, L)

    i.n_iter = nSamples
    i = AGP.init_sampler!(i, L, N, nSamples, true)

    @test i.sample_store == zeros(Float64, 2 * nSamples, N, L)

    i.n_iter = 2 * nSamples
    i = AGP.init_sampler!(i, L, N, 2, false)

    @test i.sample_store == zeros(Float64, 2, N, L)

    i.n_iter = 1
    fs = [rand(N) for _ in 1:L]
    AGP.store_variables!(i, fs)
    @test i.sample_store[1, :, 1] == fs[1]
end


text:
const JULIADB_DIR = ".juliadb"
const JULIADB_FILECACHE = "csv_metadata"
const JULIADB_INDEXFILE = "juliadb_index"

function files_from_dir(dir)
    dir = abspath(dir)
    filter(isfile, [ joinpath(dir, f) for f in readdir(dir) if !startswith(f, ".") ])
end

function format_bytes(nb)
    bytes, mb = Base.prettyprint_getunits(nb, length(Base._mem_units), Int64(1024))
    if mb == 1
        @sprintf("%d %s%s", bytes, Base._mem_units[mb], bytes==1 ? "" : "s")
    else
        @sprintf("%.3f %s", bytes, Base._mem_units[mb])
    end
end

function offset_index!(x, o)
    l = length(x)
    copyto!(columns(x.index)[1], o:o+l-1)
    x
end

function offset_index!(x::DNDSparse, o=1)
    lengths = map(a->get(a.nrows), x.domains)
    offs = [0; cumsum(lengths[1:end-1])] .+ 1
    fromchunks(delayedmap(offset_index!, x.chunks, offs))
end

Base.@deprecate loadfiles(files, delim=','; opts...) loadndsparse(files; delim=delim, opts...)

"""
`loadtable(files::Union{AbstractVector,String}; <options>)`

Load a [table](@ref Table) from CSV files.

`files` is either a vector of file paths, or a directory name.

# Options:

- `output::AbstractString` -- directory name to write the table to. By default data is loaded directly to memory. Specifying this option will allow you to load data larger than the available memory.
- `indexcols::Vector` -- columns to use as primary key columns. (defaults to [])
- `datacols::Vector` -- non-indexed columns. (defaults to all columns but indexed columns). Specify this to only load a subset of columns. In place of the name of a column, you can specify a tuple of names -- this will treat any column with one of those names as the same column, but use the first name in the tuple. This is useful when the same column changes name between CSV files. (e.g. `vendor_id` and `VendorId`)
- `distributed::Bool` -- should the output dataset be loaded as a distributed table? If true, this will use all available worker processes to load the data. (defaults to true if workers are available, false if not)
- `chunks::Int` -- number of chunks to create when loading distributed. (defaults to number of workers)
- `delim::Char` -- the delimiter character. (defaults to `,`). Use `spacedelim=true` to split by spaces.
- `spacedelim::Bool`: parse space-delimited files. `delim` has no effect if true.
- `quotechar::Char` -- quote character. (defaults to `"`)
- `escapechar::Char` -- escape character. (defaults to `"`)
- `filenamecol::Union{Symbol, Pair}` -- create a column containing the file names from where each row came from. This argument gives a name to the column. By default, `basename(name)` of the name is kept, and ".csv" suffix will be stripped. To provide a custom function to apply on the names, use a `name => Function` pair. By default, no file name column will be created.
- `header_exists::Bool` -- does header exist in the files? (defaults to true)
- `colnames::Vector{String}` -- specify column names for the files, use this with (`header_exists=false`, otherwise first row is discarded). By default column names are assumed to be present in the file.
- `samecols` -- a vector of tuples of strings where each tuple contains alternative names for the same column. For example, if some files have the name "vendor_id" and others have the name "VendorID", pass `samecols=[("VendorID", "vendor_id")]`.
- `colparsers` -- either a vector or dictionary of data types or an [`AbstractToken` object](https://juliacomputing.com/TextParse.jl/stable/#Available-AbstractToken-types-1) from [TextParse](https://juliacomputing.com/TextParse.jl/stable) package. By default, these are inferred automatically. See `type_detect_rows` option below.
- `type_detect_rows`: number of rows to use to infer the initial `colparsers` defaults to 20.
- `nastrings::Vector{String}` -- strings that are to be considered missing values. (defaults to `TextParse.NA_STRINGS`)
- `skiplines_begin::Char` -- skip some lines in the beginning of each file. (doesn't skip by default)

- `usecache::Bool`: (vestigial)
"""
function loadtable(files::Union{AbstractVector,String}; opts...)
    _loadtable(IndexedTable, files; opts...)
end

"""
`loadndsparse(files::Union{AbstractVector,String}; <options>)`

Load an [NDSparse](@ref) from CSV files.

`files` is either a vector of file paths, or a directory name.

# Options:

- `indexcols::Vector` -- columns to use as indexed columns. (by default a `1:n` implicit index is used.)
- `datacols::Vector` -- non-indexed columns. (defaults to all columns but indexed columns). Specify this to only load a subset of columns. In place of the name of a column, you can specify a tuple of names -- this will treat any column with one of those names as the same column, but use the first name in the tuple. This is useful when the same column changes name between CSV files. (e.g. `vendor_id` and `VendorId`)

All other options are identical to those in [`loadtable`](@ref)

"""
function loadndsparse(files::Union{AbstractVector,String}; opts...)
    _loadtable(NDSparse, files; opts...)
end

# Can load both NDSparse and table
function _loadtable(T, files::Union{AbstractVector,String};
                    chunks=nothing,
                    output=nothing,
                    append=false,
                    indexcols=[],
                    distributed=chunks != nothing || length(procs()) > 1,
                    usecache=false,
                    opts...)

    if isa(files, String)
        if isdir(files)
            files = files_from_dir(files)
        elseif isfile(files)
            files = [files]
        else
            throw(ArgumentError("Specified path is neither a file, " *
                                "nor a directory."))
        end
    else
        for file in files
            if !isfile(file)
                throw(ArgumentError("No file named $file."))
            end
        end
    end

    if isempty(files)
        throw(ArgumentError("Specify at least one file to load."))
    end

    if chunks === nothing && distributed
        chunks = nworkers()
    end

    if !distributed
        filegroups = [files]
    else
        if isa(chunks, Integer)
            chunks = Dagger.split_range(1:length(files), chunks)
        end

        filegroups = filter(!isempty, map(x->files[x], chunks))
    end

    loadgroup = delayed() do group
        _loadtable_serial(T, group; indexcols=indexcols, opts...)[1]
    end

    if output !== nothing && append
        prevchunks = load(output).chunks
    else
        prevchunks = []
    end

    y = fromchunks(map(loadgroup, filegroups),
                   output=output, fnoffset=length(prevchunks))
    x = fromchunks(vcat(prevchunks, y.chunks))

    if output !== nothing
        open(joinpath(output, JULIADB_INDEXFILE), "w") do io
            serialize(io, x)
        end
        _makerelative!(x, output)
    end

    if x isa DNDSparse && isempty(indexcols)
        # implicit index
        x = offset_index!(x, 1)
    end

    if !distributed
        return collect(x)
    else
        return x
    end
end

Base.@deprecate ingest(files, output; kwargs...) loadndsparse(files; output=output, kwargs...)

Base.@deprecate ingest!(files, output; kwargs...) loadndsparse(files; output=output, append=true, kwargs...)

"""
`load(dir::AbstractString)`

Load a saved `DNDSparse` from `dir` directory. Data can be saved
using the `save` function.
"""
function load(f::AbstractString; procs=workers())
    if isdir(f)
        x = open(joinpath(f, JULIADB_INDEXFILE)) do io
            deserialize(io)
        end
        _makerelative!(x, f)
        _evenlydistribute!(x, procs)
        x
    elseif isfile(f)
        MemPool.unwrap_payload(open(deserialize, f))
    else
        error("$f is not a file or directory")
    end
end

"""
    save(t::Union{DNDSparse, DIndexedTable}, destdir::AbstractString)

Saves a distributed dataset to disk in directory `destdir`. Saved data can be loaded with [`load`](@ref).
"""
function save(x::DDataset, output::AbstractString)
    if !isempty(x.chunks)
        y = fromchunks(x.chunks, output=output)
    else
        y = x
    end
    open(joinpath(output, JULIADB_INDEXFILE), "w") do io
        serialize(io, y)
    end
    _makerelative!(y, output)
    y
end

"""
    save(t::Union{NDSparse, IndexedTable}, dest::AbstractString)

Save a dataset to disk as `dest`.  Saved data can be loaded with [`load`](@ref).
"""
function save(data::Dataset, f::AbstractString)
    sz = open(f, "w") do io
        serialize(io, MemPool.MMWrap(data))
    end
    load(f)
end

function _makerelative!(t, dir::AbstractString)
    foreach(t.chunks) do c
        h = c.handle
        if isa(h, FileRef)
            c.handle = FileRef(joinpath(dir, h.file), h.size)
        end
    end
end

using MemPool

# Fix to load n/p number of files per process
function _evenlydistribute!(t, wrkrs)
    for (r, w) in zip(t.chunks, Iterators.cycle(wrkrs))
        if r.handle isa FileRef
            r.handle.force_pid[] = w
        end
    end
end

deserialize(io::AbstractSerializer, DT::Type{DNDSparse{K,V}}) where {K,V} = _deser(io, DT)
deserialize(io::AbstractSerializer, DT::Type{DIndexedTable{T,K}}) where {T,K} = _deser(io, DT)
function _deser(io::AbstractSerializer, t)
    nf = fieldcount(t)
    x = ccall(:jl_new_struct_uninit, Any, (Any,), t)
    t.mutable && Serialization.deserialize_cycle(io, x)
    for i in 1:nf
        tag = Int32(read(io.io, UInt8)::UInt8)
        if tag != Serialization.UNDEFREF_TAG
            ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), x, i-1, Serialization.handle_deserialize(io, tag))
        end
    end
    return x
end

function MemPool.mmwrite(io::AbstractSerializer, arr::StringArray)
    Serialization.serialize_type(io, MemPool.MMSer{StringArray})
    serialize(io, eltype(arr))
    MemPool.mmwrite(io, arr.buffer)
    MemPool.mmwrite(io, arr.offsets)
    MemPool.mmwrite(io, arr.lengths)
    return
end

function MemPool.mmread(::Type{StringArray}, io, mmap)
    T = deserialize(io)
    buffer  = deserialize(io)
    offsets = deserialize(io)
    lengths  = deserialize(io)
    return StringArray{T, ndims(offsets)}(buffer, offsets, lengths)
end


text:
# AUTO GENERATED FILE - DO NOT EDIT

export chart

"""
    chart(;kwargs...)

A Chart component.

Keyword arguments:
- `id` (Bool | Real | String | Dict | Array; optional)
- `option` (Bool | Real | String | Dict | Array; optional)
- `config` (Bool | Real | String | Dict | Array; optional)
- `style` (Bool | Real | String | Dict | Array; optional)
- `resizeObserver` (Dict; optional)
"""
function chart(; kwargs...)
        available_props = Symbol[:id, :option, :config, :style, :resizeObserver]
        wild_props = Symbol[]
        return Component("chart", "Chart", "dash_echarts", available_props, wild_props; kwargs...)
end

max_stars_repo_name: Covarians/dash-echarts

max_issues_repo_name: Covarians/dash-echarts

max_forks_repo_name: Covarians/dash-echarts


text:
using Pkg
Pkg.activate(".")

using Roots
using Dierckx
using PyPlot

function find_c(β)
    -log(1-β) / β
end
β = 0.01:0.005:0.99
plot(find_c.(β), β, linewidth=6)

function find_β(c)
    function f(β)
        β + exp(-c*β) - 1.0
    end
    find_zero(f, (0.00001, 1))
end

c = 1.005:0.005:4.65
plot(c, find_β.(c))

spl_c = Spline1D(c, find_β.(c))
spl_β = Spline1D(find_c.(β), β)
plot(c, spl_c.(c) - spl_β.(c))


text:
module PanelCountData

export welcome

# Write your package code here.
function welcome()
    println("welcome to Panel Count.")
end

end


text:
module FluxWiG

using Flux
using LinearAlgebra

export WiG

struct WiG{T<:AbstractFloat, L<:Union{Dense,Conv}}
    fn::L
end
(wig::WiG)(x::AbstractArray) = x .* σ.(wig.fn(x))

# Dense Constructor
function (::Type{WiG{T,Dense}})(n::Int, g::T=one(T), s::T=zero(T)) where {T<:AbstractFloat}
    Wg = Matrix{T}(I, n, n) .* g
    if s != 0
        Wg .+= randn(T, n, n) .* s
    end
    fn = Dense(param(Wg), param(zeros(T, n)))
    WiG{T,Dense}(fn)
end
(::Type{WiG})(fn::Dense{F,<:AbstractArray{T}}) where {F, T<:AbstractFloat} = WiG{T,Dense}(fn)

# Conv Constructor
function (::Type{WiG{T,Conv}})((fh, fw, n)::NTuple{3,Int}, g::T=one(T), s::T=zero(T)) where {T<:AbstractFloat}
    @assert isodd(fh) && isodd(fw) "Filter-size MUST be odd."
    padsize = (fh ÷ 2, fw ÷ 2)
    Wg = s == 0 ? zeros(T, fh, fw, n, n) : randn(T, fh, fw, n, n) .* s
    c0 = fh ÷ 2 + 1
    c1 = fw ÷ 2 + 1
    for ch in 1:n
        Wg[c0, c1, ch, ch] += g
    end
    fn = Conv(param(Wg), param(zeros(T, n)); pad=padsize)
    WiG{T,Conv}(fn)
end
(::Type{WiG})(fn::Conv{N,F,<:AbstractArray{T}}) where {N, F, T<:AbstractFloat} = WiG{T,Conv}(fn)

# params
# Flux.params(wig::WiG) = params(wig.fn)
Flux.@treelike WiG

end # module


text:
using Kinetic, ProgressMeter, Plots, LinearAlgebra, JLD2, Flux
using Flux: onecold, @epochs

###
# define neural model
###

cd(@__DIR__)
@load "data.jld2" X1 Y1 X2 Y2

device = cpu

X1 = Float32.(X1) |> device
Y1 = Float32.(Y1) |> device
X2 = Float32.(X2) |> device
Y2 = Float32.(Y2) |> device

isNewStart = true
#isNewStart = false
if isNewStart
    nn = Chain(
        Dense(7, 28, sigmoid),
        Dense(28, 56, sigmoid),
        Dense(56, 28, sigmoid),
        Dense(28, 2, sigmoid),
    )
else
    @load "nn.jld2" nn
end
nn = nn |> device

data = Flux.Data.DataLoader((X1, Y1), shuffle = true) |> device
ps = params(nn)
sqnorm(x) = sum(abs2, x)
#loss(x, y) = sum(abs2, nn(x) - y) / size(x, 2) #+ 1e-6 * sum(sqnorm, ps)
loss(x, y) = Flux.binarycrossentropy(nn(x), y)
cb = () -> println("loss: $(loss(X1, Y1))")
opt = ADAM()

@epochs 2 Flux.train!(loss, ps, data, opt, cb = Flux.throttle(cb, 1))

#sci_train!(nn, (X1, Y1), ADAM(); device = gpu, epoch = 10)
#sci_train!(nn, (X1, Y1), ADAM(); device = cpu, epoch = 1)

# accuracy
function accuracy(X, Y)
    Y1 = nn(X)

    YA1 = [onecold(Y1[:, i]) for i in axes(Y1, 2)]
    YA = [onecold(Y[:, i]) for i in axes(Y, 2)]

    accuracy = 0.0
    for i in eachindex(YA)
        if YA[i] == YA1[i]
            accuracy += 1.0
        end
    end
    accuracy /= length(YA)

    return accuracy
end

accuracy(X1, Y1)
accuracy(X2, Y2)

nn = nn |> cpu

@save "nn.jld2" nn

max_stars_repo_name: CSMMLab/Flowmachine

max_issues_repo_name: CSMMLab/Flowmachine

max_forks_repo_name: CSMMLab/Flowmachine


text:
# syntax: proto3
using ProtoBuf
import ProtoBuf.meta

mutable struct TextPluginData <: ProtoType
    version::Int32
    TextPluginData(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct TextPluginData

export TextPluginData


text:
################################################################################
# MIT License
#
# Copyright (c) 2021 Hajime Nakagami<nakagami@gmail.com>
#
# 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.
################################################################################

function fetch_records(stmt::Statement)::Vector{Vector{Any}}
    more_data = true
    results::Vector{Vector{Any}} = []

    while more_data
        _op_fetch(stmt.conn.wp, stmt.handle, calc_blr(stmt.xsqlda))
        rows_segments, more_data = _op_fetch_response(stmt.conn.wp, stmt.handle, stmt.xsqlda)
        results = vcat(results, rows_segments)
    end

    # Convert BLOB handle to data
    if any(x->x.sqltype == SQL_TYPE_BLOB, stmt.xsqlda)
        transaction = Transaction(stmt.conn.wp)
        for x in 1:length(stmt.xsqlda)
            if stmt.xsqlda[x].sqltype == SQL_TYPE_BLOB
                for i in 1:length(results)
                    if !(results[i][x] isa Missing)
                        results[i][x] = get_blob_segments(stmt.conn.wp, results[i][x], transaction.handle)
                        if stmt.xsqlda[x].sqlsubtype == 1   # TEXT
                            results[i][x] = String(results[i][x])
                        end
                    end
                end
            end
        end
        close!(transaction)
    end

    results
end

"""
    DBInterface.execute(stmt::Firebird.Statement; params) => Firebird.Cursor

Execute the prepared statement `stmt`.
"""
function DBInterface.execute(stmt::Statement, params=())::Cursor
    _op_execute(stmt.conn.wp, stmt.handle, stmt.conn.transaction.handle, params)
    _op_response(stmt.conn.wp)
    if stmt.stmt_type == isc_info_sql_stmt_select
        rows = fetch_records(stmt)
        wp = stmt.conn.wp
        _op_free_statement(wp, stmt.handle, DSQL_close)
        if wp.accept_type == ptype_lazy_send
            wp.lazy_response_count += 1
        else
            _op_response(wp)
        end
    else
        rows = []
    end
    names = [Symbol(x.aliasname) for x in stmt.xsqlda]
    types = [juliatype(x) for x in stmt.xsqlda]
    lookup = Dict(x => i for (i, x) in enumerate(names))

    Cursor(stmt, names, types, lookup, rows, 1)
end

"""
    DBInterface.execute(conn::Firebird.Connection, sql::String; params) => Firebird.Cursor

Execute the SQL `sql` statement with the database connection `conn`.
"""
function DBInterface.execute(conn::Connection, sql::AbstractString, params=())::Cursor
    stmt = DBInterface.prepare(conn, sql)
    try
        DBInterface.execute(stmt, params)
    finally
        DBInterface.close!(stmt)
    end
end


text:
# Model object
abstract type AbstractRegression end

##################################################
# PRIMAL
##################################################
const NULL_HASH = UInt64(0)

# define logistic problem
# TODO: add inheritance from MMI.Probabilistic
mutable struct LogisticRegressor{T <: Real} <: AbstractRegression
    data::AbstractDataset{T}
    # Penalty to apply (L0, L1 or L2)
    penalty::AbstractPenalty
    # Hash for current x to avoid computing twice some operations
    hash_x::UInt64
    # Shall we fit intercept as well?
    fit_intercept::Bool
    # To log evolution during optimization process
    logger::AbstractLogger
end
function LogisticRegressor(data::AbstractDataset{T},
                           penalty::AbstractPenalty) where T
    return LogisticRegressor(data, penalty, NULL_HASH, false , NoLogger())
end
function LogisticRegressor(X::AbstractArray{T, 2}, y::AbstractVector{T};
                           penalty=L2Penalty(0.0),
                           logger=NoLogger(),
                           fit_intercept=false) where T
    Dataset = issparse(X) ? SparseLogitData : LogitData
    data = Dataset(X, y)
    return LogisticRegressor(data, penalty, NULL_HASH, fit_intercept ,logger)
end

nfeatures(model::LogisticRegressor) = nfeatures(model.data)

function update!(model::LogisticRegressor, x)
    new_hash = hash(x)
    # Update prediction only if we test a new x
    if new_hash != model.hash_x
        nfeat = nfeatures(model)
        if model.fit_intercept
            predict!(model.data, view(x, 1:nfeat), x[end])
        else
            predict!(model.data, x)
        end
        model.hash_x = new_hash
    end
end

function loss(x::AbstractVector{T}, model::LogisticRegressor{T}) where T
    # Sanity check
    update!(model, x)
    if model.fit_intercept
        w = @view x[1:nfeatures(model)]
        obj = loss(w, model.data) + model.penalty(w)
    else
        obj = loss(x, model.data) + model.penalty(x)
    end
    return obj
end

function gradient!(grad::AbstractVector{T}, x::AbstractVector{T}, model::LogisticRegressor{T}) where T
    # Sanity check
    @assert length(x) == length(grad)
    # Reset gradient
    fill!(grad, 0.0)
    update!(model, x)
    if model.fit_intercept
        gradient!(grad, x, model.data, true)
        p = nfeatures(model)
        g = @view grad[1:p]
        w = @view x[1:p]
        gradient!(g, w, model.penalty)
    else
        gradient!(grad, x, model.data)
        gradient!(grad, x, model.penalty)
    end
end

function hessian!(hess::AbstractVector{T}, x::AbstractVector{T}, model::LogisticRegressor{T}) where T
    # Sanity check
    p = length(x)
    @assert length(hess) == p * (p + 1) / 2
    # Reset hessian
    fill!(hess, 0.0)
    update!(model, x)
    if model.fit_intercept
        hessian!(hess, x, model.data, true)
        hindex = Int[]
        p = nfeatures(model)
        # TODO: remove computation of hindex here.
        count = 1
        for j in 1:p
            for i in j:p
                push!(hindex, count)
                count += 1
            end
            count += 1
        end
        h = @view hess[hindex]
        w = @view x[1:p]
        hessian!(h, w, model.penalty)
    else
        hessian!(hess, x, model.data)
        hessian!(hess, x, model.penalty)
    end
end

function hessvec!(hessvec::AbstractVector{T}, x::AbstractVector{T}, vec::AbstractVector{T}, model::LogisticRegressor{T}) where T
    # Sanity check
    @assert length(vec) == length(hessvec)
    @assert length(vec) == length(x)
    # Reset hessian
    fill!(hessvec, 0.0)
    # Update yₚ
    update!(model, x)
    if model.fit_intercept
        hessvec!(hessvec, x, vec, model.data, true)
        p = nfeatures(model)
        hv = @view hessvec[1:p]
        v = @view vec[1:p]
        w = @view x[1:p]
        hessvec!(hv, w, v, model.penalty)
    else
        hessvec!(hessvec, x, vec, model.data)
        hessvec!(hessvec, x, vec, model.penalty)
    end
end

function diaghess!(diagh::AbstractVector{T}, x::AbstractVector{T},
                   model::LogisticRegressor{T}) where T
    fill!(diagh, 0.0)
    # Update yₚ
    update!(model, x)
    if model.fit_intercept
        diaghess!(diagh, x, model.data, true)
        p = nfeatures(model)
        dh = @view diagh[1:p]
        w = @view x[1:p]
        diaghess!(dh, w, model.penalty)
    else
        diaghess!(diagh, x, model.data)
        diaghess!(diagh, x, model.penalty)
    end
end

function generate_callbacks(model::LogisticRegressor)
    f = x -> loss(x, model)
    grad! = (g, x) -> gradient!(g, x, model)
    hess! = (h, x) -> hessian!(h, x, model)
    hessvec! = (h, x, v) -> hessian!(h, x, v, model)
    return f, grad!, hess!, hessvec!
end

##################################################
# DUAL
##################################################
# Define dual model
# Currently, it suports only L2 penalty
mutable struct DualLogisticRegressor{T <: Real} <: AbstractRegression
    data::AbstractDataset{T}
    penalty::L2Penalty
    hash_λ::UInt64
    logger::AbstractLogger
end
DualLogisticRegressor(data::AbstractDataset{T}, penalty::L2Penalty) where T = DualLogisticRegressor(data, penalty, NULL_HASH, NoLogger())

function DualLogisticRegressor(X::AbstractArray{T, 2}, y::AbstractVector{T};
                               penalty=L2Penalty(0.0),
                               logger=NoLogger()) where T
    if issparse(X)
        error("Sparse arrays not yet supported for dual logistic regression")
    end
    data = DualLogitData(X, y)
    return DualLogisticRegressor(data, penalty, NULL_HASH, logger)
end

function update!(model::DualLogisticRegressor, x::AbstractVector)
    new_hash = hash(x)
    if new_hash != model.hash_λ
        predict!(model.data, x)
        model.hash_λ = new_hash
    end
end

function loss(x::AbstractVector{T},
              model::DualLogisticRegressor{T}) where T
    update!(model, x)
    obj = loss(x, model.data) + model.penalty(model.data.x_pred)
    return obj
end

function gradient!(grad::AbstractVector{T}, x::AbstractVector{T},
                   model:: DualLogisticRegressor{T}) where T
    fill!(grad, 0.0)
    update!(model, x)
    gradient!(grad, x, model.data)
    # compute gradient of penalty ||X^\top θ||_2^2 in two steps
    gpenal = zeros(length(model.data.x_pred))
    gradient!(gpenal, model.data.x_pred, model.penalty)
    grad .+= model.data.X * gpenal
end

function hessian!(hess::AbstractVector{T},
                  x::AbstractVector{T},
                  model::DualLogisticRegressor{T}) where T
    fill!(hess, 0.0)
    update!(model, x)
    hessian!(hess, x, model.data)
    hessian!(hess, x, model.penalty)
end

function hessvec!(hv::AbstractVector{T},
                  x::AbstractVector{T},
                  vec::AbstractVector{T},
                  model::DualLogisticRegressor{T}) where T
    fill!(hv, 0.0)
    update!(model, x)
    hessvec!(hv, x, vec, model.data)
    hessvec!(hv, x, vec, model.penalty)
end

function generate_callbacks(model::DualLogisticRegressor)
    f = x -> loss(x, model)
    grad! = (g, x) -> gradient!(g, x, model)
    hess! = (h, x) -> hessian!(h, x, model)
    hessvec! = (h, x, v) -> hessian!(h, x, v, model)
    return f, grad!, hess!, hessvec!
end


text:
"""
Parameter state in Gibbs sampler for feature allocation model
with feature selection. This augments `State` with parameters:
`r`, `W_star`, and `p`-or-`omega`.

Usage example:
==============

```julia
julia> s = StateFS{Float64}()
julia> s.r = rand(Bool, 3, 10)
```

NOTE:
=====
The update order should be:
Z -> v -> alpha -> p -> r -> lam -> W* -> gamma -> eta -> delta -> sig^2 -> y*

Particularly, `lam` must be updated before W* and preferrably after p & r.
And `gamma` should be updated after Z and before delta.
"""
mutable struct StateFS{F <: AbstractFloat}
  r::Matrix{Bool}  # dim: IxK
  W_star::Matrix{F}  # dim: IxK
  omega::Vector{F}  # dim: number of covariates (`length(x_i)`) + 1
  theta::State{F}  # State

  # Primary constructor returns an uninitialized stateFS.
  StateFS{F}() where {F <: AbstractFloat} = new()
end

function StateFS{F}(theta::State{F}, dfs::DataFS;
                    omega_prior=nothing,
                    W_star_prior=nothing,
                    eps_r::Float64=0.0, verbose=1) where {F <: AbstractFloat}
  sfs = StateFS{F}()
  I, K = size(theta.W)

  # Set theta
  sfs.theta = theta

  # Set W*
  if W_star_prior == nothing
    sfs.W_star = theta.W * K * 5
  else
    # Want empirical mean of W_star  to start at prior mean of W_star.
    # (wstar_1 + ... + wstar_K) / K = m , where m is some mean.
    # => (wstar_1 + ... + wstar_K) = K * m
    #
    # So, setting wstar_k = w_k * K * m will lead to a
    # wstar that when normalized is w, but has a mean of m.
    sfs.W_star = theta.W * K * mean(W_star_prior)
  end

  # Set R
  if eps_r > 0.0
    sfs.r = sfs.theta.W .> eps_r
  else
    sfs.r = ones(Bool, I, K)
  end

  # Set omega
  if omega_prior == nothing
    sfs.omega = zeros(dfs.P)
  else
    sfs.omega = rand(omega_prior, dfs.P)
  end

  if verbose > 0
    println("Initial sfs.r: $(sfs.r)")
    println("Initial sfs:W_star $(sfs.W_star)")
    println("Initial sfs.theta.W: $(sfs.theta.W)")
    println("Initial sfs.theta.Z: $(sfs.theta.Z)")
    println("Inital sfs.omega: $(sfs.omega)")
  end

  return sfs
end


text:
@naiveVariationalRule(:node_type     => NLatentAutoregressiveX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (Nothing, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalNLARXOutNPPPPP)

@naiveVariationalRule(:node_type     => NLatentAutoregressiveX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, Nothing, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalNLARXIn1PNPPPP)

@naiveVariationalRule(:node_type     => NLatentAutoregressiveX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, Nothing, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalNLARXIn2PPNPPP)

@naiveVariationalRule(:node_type     => NLatentAutoregressiveX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, Nothing, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalNLARXIn3PPPNPP)

@naiveVariationalRule(:node_type     => NLatentAutoregressiveX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, Nothing, ProbabilityDistribution),
                      :name          => VariationalNLARXIn4PPPPNP)

@naiveVariationalRule(:node_type     => NLatentAutoregressiveX,
                      :outbound_type => Message{Gamma},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, Nothing),
                      :name          => VariationalNLARXIn5PPPPPN)

max_stars_repo_name: wmkouw/NLARX


text:
using Base.Test
using KrylovMethods
using LinearOperators
using MatrixDepot

println("=== Testing MINRES ===")

A = sprandn(10,10,0.1)
A = A'*A + 4*speye(10)
Af(v) = A*v
rhs = randn(10)
rhst = copy(rhs)
# test stop by maxIter
x0,f0, = minres(A,rhs,maxIter=1,rtol=1e-12,gtol=1e-12,btol=1e-20,out=1,storeInterm=true)
@test f0 == -1

x1,f1, = minres(A,rhs,maxIter=10,rtol=1e-12,gtol=1e-12,btol=1e-20,out=2)
@test all(rhs.==rhst)
x2,f2, = minres(Af,rhs,x=x0[:,1],maxIter=10,rtol=1e-12,gtol=1e-12,btol=1e-20,out=2)
@test all(rhs.==rhst)

@test f1==0
@test (norm(A*x1-rhs)/norm(rhs)) < 1e-11
@test (norm(A*x2-rhs)/norm(rhs)) < 1e-11
@test norm(x1-x2)/norm(x1) < 1e-12

println("test stop by condlim")
A = matrixdepot("fiedler",50)
rhs = randn(50)
x,f, = minres(A,rhs,maxIter=50,rtol=1e-12,gtol=1e-12,btol=1e-20,condlim=200)
@test f==-3

println("test stop by beta")
A = randn(20,4)
A = A*A'
rhs = randn(20)
x,f, = minres(A,rhs,maxIter=20,rtol=1e-20,gtol=1e-20,btol=1e-10,condlim=1e20)
@test f==-2

println("test if anything is done for zero right hand side")
A = randn(20,4)
mv = 0
function Af(x)
	mv+=1
	return Af*x
end
rhs = zeros(20)
x, = minres(Af,rhs,maxIter=20,rtol=1e-20,gtol=1e-20,btol=1e-10,condlim=1e20)
@test all(x.==0)
@test mv==0
println("=== MINRES : All tests passed. ====")


text:
using Distributions:
    UnivariateDistribution, MultivariateDistribution, MatrixDistribution, Distribution

const AMBIGUITY_MSG =
    "Ambiguous `LHS .~ RHS` or `@. LHS ~ RHS` syntax. The broadcasting " *
    "can either be column-wise following the convention of Distributions.jl or " *
    "element-wise following Julia's general broadcasting semantics. Please make sure " *
    "that the element type of `LHS` is not a supertype of the support type of " *
    "`AbstractVector` to eliminate ambiguity."

alg_str(spl::Sampler) = string(nameof(typeof(spl.alg)))

# utility funcs for querying sampler information
require_gradient(spl::Sampler) = false
require_particles(spl::Sampler) = false

_getindex(x, inds::Tuple) = _getindex(view(x, first(inds)...), Base.tail(inds))
_getindex(x, inds::Tuple{}) = x

# assume
"""
    tilde_assume(context::SamplingContext, right, vn, inds, vi)

Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
accumulate the log probability, and return the sampled value with a context associated
with a sampler.

Falls back to
```julia
tilde_assume(context.rng, context.context, context.sampler, right, vn, inds, vi)
```
"""
function tilde_assume(context::SamplingContext, right, vn, inds, vi)
    return tilde_assume(context.rng, context.context, context.sampler, right, vn, inds, vi)
end

# Leaf contexts
tilde_assume(::DefaultContext, right, vn, inds, vi) = assume(right, vn, vi)
function tilde_assume(
    rng::Random.AbstractRNG, ::DefaultContext, sampler, right, vn, inds, vi
)
    return assume(rng, sampler, right, vn, vi)
end

function tilde_assume(context::PriorContext{<:NamedTuple}, right, vn, inds, vi)
    if haskey(context.vars, getsym(vn))
        vi[vn] = vectorize(right, _getindex(getfield(context.vars, getsym(vn)), inds))
        settrans!(vi, false, vn)
    end
    return tilde_assume(PriorContext(), right, vn, inds, vi)
end
function tilde_assume(
    rng::Random.AbstractRNG,
    context::PriorContext{<:NamedTuple},
    sampler,
    right,
    vn,
    inds,
    vi,
)
    if haskey(context.vars, getsym(vn))
        vi[vn] = vectorize(right, _getindex(getfield(context.vars, getsym(vn)), inds))
        settrans!(vi, false, vn)
    end
    return tilde_assume(rng, PriorContext(), sampler, right, vn, inds, vi)
end
function tilde_assume(::PriorContext, right, vn, inds, vi)
    return assume(right, vn, vi)
end
function tilde_assume(rng::Random.AbstractRNG, ::PriorContext, sampler, right, vn, inds, vi)
    return assume(rng, sampler, right, vn, vi)
end

function tilde_assume(context::LikelihoodContext{<:NamedTuple}, right, vn, inds, vi)
    if haskey(context.vars, getsym(vn))
        vi[vn] = vectorize(right, _getindex(getfield(context.vars, getsym(vn)), inds))
        settrans!(vi, false, vn)
    end
    return tilde_assume(LikelihoodContext(), right, vn, inds, vi)
end
function tilde_assume(
    rng::Random.AbstractRNG,
    context::LikelihoodContext{<:NamedTuple},
    sampler,
    right,
    vn,
    inds,
    vi,
)
    if haskey(context.vars, getsym(vn))
        vi[vn] = vectorize(right, _getindex(getfield(context.vars, getsym(vn)), inds))
        settrans!(vi, false, vn)
    end
    return tilde_assume(rng, LikelihoodContext(), sampler, right, vn, inds, vi)
end
function tilde_assume(::LikelihoodContext, right, vn, inds, vi)
    return assume(NoDist(right), vn, vi)
end
function tilde_assume(
    rng::Random.AbstractRNG, ::LikelihoodContext, sampler, right, vn, inds, vi
)
    return assume(rng, sampler, NoDist(right), vn, vi)
end

function tilde_assume(context::MiniBatchContext, right, vn, inds, vi)
    return tilde_assume(context.context, right, vn, inds, vi)
end

function tilde_assume(rng, context::MiniBatchContext, sampler, right, vn, inds, vi)
    return tilde_assume(rng, context.context, sampler, right, vn, inds, vi)
end

function tilde_assume(context::PrefixContext, right, vn, inds, vi)
    return tilde_assume(context.context, right, prefix(context, vn), inds, vi)
end

function tilde_assume(rng, context::PrefixContext, sampler, right, vn, inds, vi)
    return tilde_assume(rng, context.context, sampler, right, prefix(context, vn), inds, vi)
end

"""
    tilde_assume!(context, right, vn, inds, vi)

Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
accumulate the log probability, and return the sampled value.

By default, calls `tilde_assume(context, right, vn, inds, vi)` and accumulates the log
probability of `vi` with the returned value.
"""
function tilde_assume!(context, right, vn, inds, vi)
    value, logp = tilde_assume(context, right, vn, inds, vi)
    acclogp!(vi, logp)
    return value
end

# observe
"""
    tilde_observe(context::SamplingContext, right, left, vname, vinds, vi)

Handle observed variables with a `context` associated with a sampler.

Falls back to
```julia
tilde_observe(context.rng, context.context, context.sampler, right, left, vname, vinds, vi)
```
"""
function tilde_observe(context::SamplingContext, right, left, vname, vinds, vi)
    return tilde_observe(
        context.rng, context.context, context.sampler, right, left, vname, vinds, vi
    )
end

"""
    tilde_observe(context::SamplingContext, right, left, vi)

Handle observed constants with a `context` associated with a sampler.

Falls back to `tilde_observe(context.context, context.sampler, right, left, vi)`.
"""
function tilde_observe(context::SamplingContext, right, left, vi)
    return tilde_observe(context.context, context.sampler, right, left, vi)
end

# Leaf contexts
tilde_observe(::DefaultContext, right, left, vi) = observe(right, left, vi)
function tilde_observe(::DefaultContext, sampler, right, left, vi)
    return observe(sampler, right, left, vi)
end
tilde_observe(::PriorContext, right, left, vi) = 0
tilde_observe(::PriorContext, sampler, right, left, vi) = 0
tilde_observe(::LikelihoodContext, right, left, vi) = observe(right, left, vi)
function tilde_observe(::LikelihoodContext, sampler, right, left, vi)
    return observe(sampler, right, left, vi)
end

# `MiniBatchContext`
function tilde_observe(context::MiniBatchContext, right, left, vi)
    return context.loglike_scalar * tilde_observe(context.context, right, left, vi)
end
function tilde_observe(context::MiniBatchContext, right, left, vname, vi)
    return context.loglike_scalar * tilde_observe(context.context, right, left, vname, vi)
end

# `PrefixContext`
function tilde_observe(context::PrefixContext, right, left, vname, vi)
    return tilde_observe(context.context, right, left, prefix(context, vname), vi)
end
function tilde_observe(context::PrefixContext, right, left, vi)
    return tilde_observe(context.context, right, left, vi)
end

"""
    tilde_observe!(context, right, left, vname, vinds, vi)

Handle observed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
accumulate the log probability, and return the observed value.

Falls back to `tilde_observe!(context, right, left, vi)` ignoring the information about variable name
and indices; if needed, these can be accessed through this function, though.
"""
function tilde_observe!(context, right, left, vname, vinds, vi)
    return tilde_observe!(context, right, left, vi)
end

"""
    tilde_observe(context, right, left, vi)

Handle observed constants, e.g., `1.0 ~ Normal()`, accumulate the log probability, and
return the observed value.

By default, calls `tilde_observe(context, right, left, vi)` and accumulates the log
probability of `vi` with the returned value.
"""
function tilde_observe!(context, right, left, vi)
    logp = tilde_observe(context, right, left, vi)
    acclogp!(vi, logp)
    return left
end

function assume(rng, spl::Sampler, dist)
    return error("DynamicPPL.assume: unmanaged inference algorithm: $(typeof(spl))")
end

function observe(spl::Sampler, weight)
    return error("DynamicPPL.observe: unmanaged inference algorithm: $(typeof(spl))")
end

# fallback without sampler
function assume(dist::Distribution, vn::VarName, vi)
    r = vi[vn]
    return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn))
end

# SampleFromPrior and SampleFromUniform
function assume(
    rng::Random.AbstractRNG,
    sampler::Union{SampleFromPrior,SampleFromUniform},
    dist::Distribution,
    vn::VarName,
    vi,
)
    if haskey(vi, vn)
        # Always overwrite the parameters with new ones for `SampleFromUniform`.
        if sampler isa SampleFromUniform || is_flagged(vi, vn, "del")
            unset_flag!(vi, vn, "del")
            r = init(rng, dist, sampler)
            vi[vn] = vectorize(dist, r)
            settrans!(vi, false, vn)
            setorder!(vi, vn, get_num_produce(vi))
        else
            r = vi[vn]
        end
    else
        r = init(rng, dist, sampler)
        push!(vi, vn, r, dist, sampler)
        settrans!(vi, false, vn)
    end

    return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn))
end

# default fallback (used e.g. by `SampleFromPrior` and `SampleUniform`)
observe(sampler::AbstractSampler, right, left, vi) = observe(right, left, vi)
function observe(right::Distribution, left, vi)
    increment_num_produce!(vi)
    return Distributions.loglikelihood(right, left)
end

# .~ functions

# assume
"""
    dot_tilde_assume(context::SamplingContext, right, left, vn, inds, vi)

Handle broadcasted assumed variables, e.g., `x .~ MvNormal()` (where `x` does not occur in the
model inputs), accumulate the log probability, and return the sampled value for a context
associated with a sampler.

Falls back to
```julia
dot_tilde_assume(context.rng, context.context, context.sampler, right, left, vn, inds, vi)
```
"""
function dot_tilde_assume(context::SamplingContext, right, left, vn, inds, vi)
    return dot_tilde_assume(
        context.rng, context.context, context.sampler, right, left, vn, inds, vi
    )
end

# `DefaultContext`
function dot_tilde_assume(::DefaultContext, right, left, vns, inds, vi)
    return dot_assume(right, left, vns, vi)
end

function dot_tilde_assume(rng, ::DefaultContext, sampler, right, left, vns, inds, vi)
    return dot_assume(rng, sampler, right, vns, left, vi)
end

# `LikelihoodContext`
function dot_tilde_assume(
    context::LikelihoodContext{<:NamedTuple}, right, left, vn, inds, vi
)
    return if haskey(context.vars, getsym(vn))
        var = _getindex(getfield(context.vars, getsym(vn)), inds)
        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
        set_val!(vi, _vns, _right, _left)
        settrans!.(Ref(vi), false, _vns)
        dot_tilde_assume(LikelihoodContext(), _right, _left, _vns, inds, vi)
    else
        dot_tilde_assume(LikelihoodContext(), right, left, vn, inds, vi)
    end
end
function dot_tilde_assume(
    rng::Random.AbstractRNG,
    context::LikelihoodContext{<:NamedTuple},
    sampler,
    right,
    left,
    vn,
    inds,
    vi,
)
    return if haskey(context.vars, getsym(vn))
        var = _getindex(getfield(context.vars, getsym(vn)), inds)
        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
        set_val!(vi, _vns, _right, _left)
        settrans!.(Ref(vi), false, _vns)
        dot_tilde_assume(rng, LikelihoodContext(), sampler, _right, _left, _vns, inds, vi)
    else
        dot_tilde_assume(rng, LikelihoodContext(), sampler, right, left, vn, inds, vi)
    end
end
function dot_tilde_assume(context::LikelihoodContext, right, left, vn, inds, vi)
    return dot_assume(NoDist.(right), left, vn, vi)
end
function dot_tilde_assume(
    rng::Random.AbstractRNG, context::LikelihoodContext, sampler, right, left, vn, inds, vi
)
    return dot_assume(rng, sampler, NoDist.(right), vn, left, vi)
end

# `PriorContext`
function dot_tilde_assume(context::PriorContext{<:NamedTuple}, right, left, vn, inds, vi)
    return if haskey(context.vars, getsym(vn))
        var = _getindex(getfield(context.vars, getsym(vn)), inds)
        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
        set_val!(vi, _vns, _right, _left)
        settrans!.(Ref(vi), false, _vns)
        dot_tilde_assume(PriorContext(), _right, _left, _vns, inds, vi)
    else
        dot_tilde_assume(PriorContext(), right, left, vn, inds, vi)
    end
end
function dot_tilde_assume(
    rng::Random.AbstractRNG,
    context::PriorContext{<:NamedTuple},
    sampler,
    right,
    left,
    vn,
    inds,
    vi,
)
    return if haskey(context.vars, getsym(vn))
        var = _getindex(getfield(context.vars, getsym(vn)), inds)
        _right, _left, _vns = unwrap_right_left_vns(right, var, vn)
        set_val!(vi, _vns, _right, _left)
        settrans!.(Ref(vi), false, _vns)
        dot_tilde_assume(rng, PriorContext(), sampler, _right, _left, _vns, inds, vi)
    else
        dot_tilde_assume(rng, PriorContext(), sampler, right, left, vn, inds, vi)
    end
end
function dot_tilde_assume(context::PriorContext, right, left, vn, inds, vi)
    return dot_assume(right, left, vn, vi)
end
function dot_tilde_assume(
    rng::Random.AbstractRNG, context::PriorContext, sampler, right, left, vn, inds, vi
)
    return dot_assume(rng, sampler, right, vn, left, vi)
end

# `MiniBatchContext`
function dot_tilde_assume(context::MiniBatchContext, right, left, vn, inds, vi)
    return dot_tilde_assume(context.context, right, left, vn, inds, vi)
end

function dot_tilde_assume(
    rng, context::MiniBatchContext, sampler, right, left, vn, inds, vi
)
    return dot_tilde_assume(rng, context.context, sampler, right, left, vn, inds, vi)
end

# `PrefixContext`
function dot_tilde_assume(context::PrefixContext, right, left, vn, inds, vi)
    return dot_tilde_assume(context.context, right, prefix.(Ref(context), vn), inds, vi)
end

function dot_tilde_assume(rng, context::PrefixContext, sampler, right, left, vn, inds, vi)
    return dot_tilde_assume(
        rng, context.context, sampler, right, prefix.(Ref(context), vn), inds, vi
    )
end

"""
    dot_tilde_assume!(context, right, left, vn, inds, vi)

Handle broadcasted assumed variables, e.g., `x .~ MvNormal()` (where `x` does not occur in the
model inputs), accumulate the log probability, and return the sampled value.

Falls back to `dot_tilde_assume(context, right, left, vn, inds, vi)`.
"""
function dot_tilde_assume!(context, right, left, vn, inds, vi)
    value, logp = dot_tilde_assume(context, right, left, vn, inds, vi)
    acclogp!(vi, logp)
    return value
end

# `dot_assume`
function dot_assume(
    dist::MultivariateDistribution, var::AbstractMatrix, vns::AbstractVector{<:VarName}, vi
)
    @assert length(dist) == size(var, 1)
    # NOTE: We cannot work with `var` here because we might have a model of the form
    #
    #     m = Vector{Float64}(undef, n)
    #     m .~ Normal()
    #
    # in which case `var` will have `undef` elements, even if `m` is present in `vi`.
    r = vi[vns]
    lp = sum(zip(vns, eachcol(r))) do vn, ri
        return Bijectors.logpdf_with_trans(dist, ri, istrans(vi, vn))
    end
    return r, lp
end

function dot_assume(
    rng,
    spl::Union{SampleFromPrior,SampleFromUniform},
    dist::MultivariateDistribution,
    vns::AbstractVector{<:VarName},
    var::AbstractMatrix,
    vi,
)
    @assert length(dist) == size(var, 1)
    r = get_and_set_val!(rng, vi, vns, dist, spl)
    lp = sum(Bijectors.logpdf_with_trans(dist, r, istrans(vi, vns[1])))
    return r, lp
end

function dot_assume(
    dists::Union{Distribution,AbstractArray{<:Distribution}},
    var::AbstractArray,
    vns::AbstractArray{<:VarName},
    vi,
)
    # NOTE: We cannot work with `var` here because we might have a model of the form
    #
    #     m = Vector{Float64}(undef, n)
    #     m .~ Normal()
    #
    # in which case `var` will have `undef` elements, even if `m` is present in `vi`.
    r = reshape(vi[vec(vns)], size(vns))
    lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans(vi, vns[1])))
    return r, lp
end

function dot_assume(
    rng,
    spl::Union{SampleFromPrior,SampleFromUniform},
    dists::Union{Distribution,AbstractArray{<:Distribution}},
    vns::AbstractArray{<:VarName},
    var::AbstractArray,
    vi,
)
    r = get_and_set_val!(rng, vi, vns, dists, spl)
    # Make sure `r` is not a matrix for multivariate distributions
    lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans(vi, vns[1])))
    return r, lp
end
function dot_assume(rng, spl::Sampler, ::Any, ::AbstractArray{<:VarName}, ::Any, ::Any)
    return error(
        "[DynamicPPL] $(alg_str(spl)) doesn't support vectorizing assume statement"
    )
end

function get_and_set_val!(
    rng,
    vi,
    vns::AbstractVector{<:VarName},
    dist::MultivariateDistribution,
    spl::Union{SampleFromPrior,SampleFromUniform},
)
    n = length(vns)
    if haskey(vi, vns[1])
        # Always overwrite the parameters with new ones for `SampleFromUniform`.
        if spl isa SampleFromUniform || is_flagged(vi, vns[1], "del")
            unset_flag!(vi, vns[1], "del")
            r = init(rng, dist, spl, n)
            for i in 1:n
                vn = vns[i]
                vi[vn] = vectorize(dist, r[:, i])
                settrans!(vi, false, vn)
                setorder!(vi, vn, get_num_produce(vi))
            end
        else
            r = vi[vns]
        end
    else
        r = init(rng, dist, spl, n)
        for i in 1:n
            vn = vns[i]
            push!(vi, vn, r[:, i], dist, spl)
            settrans!(vi, false, vn)
        end
    end
    return r
end

function get_and_set_val!(
    rng,
    vi,
    vns::AbstractArray{<:VarName},
    dists::Union{Distribution,AbstractArray{<:Distribution}},
    spl::Union{SampleFromPrior,SampleFromUniform},
)
    if haskey(vi, vns[1])
        # Always overwrite the parameters with new ones for `SampleFromUniform`.
        if spl isa SampleFromUniform || is_flagged(vi, vns[1], "del")
            unset_flag!(vi, vns[1], "del")
            f = (vn, dist) -> init(rng, dist, spl)
            r = f.(vns, dists)
            for i in eachindex(vns)
                vn = vns[i]
                dist = dists isa AbstractArray ? dists[i] : dists
                vi[vn] = vectorize(dist, r[i])
                settrans!(vi, false, vn)
                setorder!(vi, vn, get_num_produce(vi))
            end
        else
            r = reshape(vi[vec(vns)], size(vns))
        end
    else
        f = (vn, dist) -> init(rng, dist, spl)
        r = f.(vns, dists)
        push!.(Ref(vi), vns, r, dists, Ref(spl))
        settrans!.(Ref(vi), false, vns)
    end
    return r
end

function set_val!(
    vi, vns::AbstractVector{<:VarName}, dist::MultivariateDistribution, val::AbstractMatrix
)
    @assert size(val, 2) == length(vns)
    foreach(enumerate(vns)) do (i, vn)
        vi[vn] = val[:, i]
    end
    return val
end
function set_val!(
    vi,
    vns::AbstractArray{<:VarName},
    dists::Union{Distribution,AbstractArray{<:Distribution}},
    val::AbstractArray,
)
    @assert size(val) == size(vns)
    foreach(CartesianIndices(val)) do ind
        dist = dists isa AbstractArray ? dists[ind] : dists
        vi[vns[ind]] = vectorize(dist, val[ind])
    end
    return val
end

# observe
"""
    dot_tilde_observe(context::SamplingContext, right, left, vi)

Handle broadcasted observed constants, e.g., `[1.0] .~ MvNormal()`, accumulate the log
probability, and return the observed value for a context associated with a sampler.

Falls back to `dot_tilde_observe(context.context, context.sampler, right, left, vi)`.
"""
function dot_tilde_observe(context::SamplingContext, right, left, vi)
    return dot_tilde_observe(context.context, context.sampler, right, left, vi)
end

# Leaf contexts
dot_tilde_observe(::DefaultContext, right, left, vi) = dot_observe(right, left, vi)
function dot_tilde_observe(::DefaultContext, sampler, right, left, vi)
    return dot_observe(sampler, right, left, vi)
end
dot_tilde_observe(::PriorContext, right, left, vi) = 0
dot_tilde_observe(::PriorContext, sampler, right, left, vi) = 0
function dot_tilde_observe(context::LikelihoodContext, right, left, vi)
    return dot_observe(right, left, vi)
end
function dot_tilde_observe(context::LikelihoodContext, sampler, right, left, vi)
    return dot_observe(sampler, right, left, vi)
end

# `MiniBatchContext`
function dot_tilde_observe(context::MiniBatchContext, right, left, vi)
    return context.loglike_scalar * dot_tilde_observe(context.context, right, left, vi)
end

# `PrefixContext`
function dot_tilde_observe(context::PrefixContext, right, left, vi)
    return dot_tilde_observe(context.context, right, left, vi)
end

"""
    dot_tilde_observe!(context, right, left, vname, vinds, vi)

Handle broadcasted observed values, e.g., `x .~ MvNormal()` (where `x` does occur in the model inputs),
accumulate the log probability, and return the observed value.

Falls back to `dot_tilde_observe!(context, right, left, vi)` ignoring the information about variable
name and indices; if needed, these can be accessed through this function, though.
"""
function dot_tilde_observe!(context, right, left, vn, inds, vi)
    return dot_tilde_observe!(context, right, left, vi)
end

"""
    dot_tilde_observe!(context, right, left, vi)

Handle broadcasted observed constants, e.g., `[1.0] .~ MvNormal()`, accumulate the log
probability, and return the observed value.

Falls back to `dot_tilde_observe(context, right, left, vi)`.
"""
function dot_tilde_observe!(context, right, left, vi)
    logp = dot_tilde_observe(context, right, left, vi)
    acclogp!(vi, logp)
    return left
end

# Falls back to non-sampler definition.
function dot_observe(::AbstractSampler, dist, value, vi)
    return dot_observe(dist, value, vi)
end
function dot_observe(dist::MultivariateDistribution, value::AbstractMatrix, vi)
    increment_num_produce!(vi)
    return Distributions.loglikelihood(dist, value)
end
function dot_observe(dists::Distribution, value::AbstractArray, vi)
    increment_num_produce!(vi)
    return Distributions.loglikelihood(dists, value)
end
function dot_observe(dists::AbstractArray{<:Distribution}, value::AbstractArray, vi)
    increment_num_produce!(vi)
    return sum(Distributions.loglikelihood.(dists, value))
end


text:
using Documenter, LiveServer

makedocs(
    modules = [LiveServer],
    format = Documenter.HTML(
        # Use clean URLs, unless built as a "local" build
        prettyurls = !("local" in ARGS),
        # custom CSS
        assets = ["assets/custom.css"]
        ),
    sitename = "LiveServer.jl",
    authors  = "Jonas Asprion, Thibaut Lienart",
    pages    = [
        "Home" => "index.md",
        "Manual" => [
            "Functionalities" => "man/functionalities.md",
            "Extending LiveServer" => "man/extending_ls.md",
            "LiveServer + Literate" => "man/ls+lit.md"
            ],
        "Library" => [
            "Public"    => "lib/public.md",
            "Internals" => "lib/internals.md",
            ],
        ], # end page
)

deploydocs(
    repo = "github.com/asprionj/LiveServer.jl.git"
)


text:
using Associates
import Associates: unwrapsinglet
using Test

x=[1,2]
unwrapsinglet(x)

v = collect(1:1:26)

collect(zip('A':'Z',v))

D1 = Dict(zip('A':'Z',v))

D2 = Dict(zip(v,rand(26)))

m1= GMap(D1)

dom(m1)

codom(m1)

m2= GMap(D2)
dom(m2)
codom(m2)

m1
m2

inv(m1).forward

using AbstractTrees
Tree(m2.forward)

# function (+)(D1::Dict, D2::Dict)
#     for k in keys(D1)
#         if k in keys(D2)
#             push!(D2[k], k)
#         else
#             D2[k] = D1[k]
#         end
#     end
# end

Dict(zip(["1","b","z"], [1,2,3])) ∪ Dict(zip(["1","b","z"], [4,5,6]))

@testset "Associates.jl" begin
    # Write your tests here.

end