// Copyright (c) 2017-2026, Lawrence Livermore National Security, LLC and other CEED contributors.
// All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
//
// SPDX-License-Identifier: BSD-2-Clause
//
// This file is part of CEED:  http://github.com/ceed

/// @file
/// Shock tube initial condition and Euler equation operator for Navier-Stokes example using PETSc - modified from eulervortex.h

// Model from:
//   On the Order of Accuracy and Numerical Performance of Two Classes of Finite Volume WENO Schemes, Zhang, Zhang, and Shu (2011).
#include <ceed/types.h>
#ifndef CEED_RUNNING_JIT_PASS
#include <math.h>
#include <stdbool.h>
#endif

#include "utils.h"

typedef struct SetupContextShock_ *SetupContextShock;
struct SetupContextShock_ {
  CeedScalar theta0;
  CeedScalar thetaC;
  CeedScalar P0;
  CeedScalar N;
  CeedScalar cv;
  CeedScalar cp;
  CeedScalar time;
  CeedScalar mid_point;
  CeedScalar P_high;
  CeedScalar rho_high;
  CeedScalar P_low;
  CeedScalar rho_low;
};

typedef struct ShockTubeContext_ *ShockTubeContext;
struct ShockTubeContext_ {
  CeedScalar Cyzb;
  CeedScalar Byzb;
  CeedScalar c_tau;
  bool       implicit;
  bool       yzb;
  int        stabilization;
};

// *****************************************************************************
// This function sets the initial conditions
//
//   Temperature:
//     T   = P / (rho * R)
//   Density:
//     rho = 1.0        if x <= mid_point
//         = 0.125      if x >  mid_point
//   Pressure:
//     P   = 1.0        if x <= mid_point
//         = 0.1        if x >  mid_point
//   Velocity:
//     u   = 0
//   Velocity/Momentum Density:
//     Ui  = rho ui
//   Total Energy:
//     E   = P / (gamma - 1) + rho (u u)/2
//
// Constants:
//   cv              ,  Specific heat, constant volume
//   cp              ,  Specific heat, constant pressure
//   mid_point       ,  Location of initial domain mid_point
//   gamma  = cp / cv,  Specific heat ratio
//
// *****************************************************************************

// *****************************************************************************
// This helper function provides support for the exact, time-dependent solution (currently not implemented) and IC formulation for Euler traveling
// vortex
// *****************************************************************************
CEED_QFUNCTION_HELPER CeedInt Exact_ShockTube(CeedInt dim, CeedScalar time, const CeedScalar X[], CeedInt Nf, CeedScalar q[], void *ctx) {
  // Context
  const SetupContextShock context   = (SetupContextShock)ctx;
  const CeedScalar        mid_point = context->mid_point;  // Midpoint of the domain
  const CeedScalar        P_high    = context->P_high;     // Driver section pressure
  const CeedScalar        rho_high  = context->rho_high;   // Driver section density
  const CeedScalar        P_low     = context->P_low;      // Driven section pressure
  const CeedScalar        rho_low   = context->rho_low;    // Driven section density

  // Setup
  const CeedScalar gamma = 1.4;   // ratio of specific heats
  const CeedScalar x     = X[0];  // Coordinates

  CeedScalar rho, P, u[3] = {0.};

  // Initial Conditions
  if (x <= mid_point + 200 * CEED_EPSILON) {
    rho = rho_high;
    P   = P_high;
  } else {
    rho = rho_low;
    P   = P_low;
  }

  // Assign exact solution
  q[0] = rho;
  q[1] = rho * u[0];
  q[2] = rho * u[1];
  q[3] = rho * u[2];
  q[4] = P / (gamma - 1.0) + rho * (u[0] * u[0]) / 2.;

  return 0;
}

// *****************************************************************************
// Helper function for computing flux Jacobian
// *****************************************************************************
CEED_QFUNCTION_HELPER void ConvectiveFluxJacobian_Euler(CeedScalar dF[3][5][5], const CeedScalar rho, const CeedScalar u[3], const CeedScalar E,
                                                        const CeedScalar gamma) {
  CeedScalar u_sq = u[0] * u[0] + u[1] * u[1] + u[2] * u[2];  // Velocity square
  for (CeedInt i = 0; i < 3; i++) {                           // Jacobian matrices for 3 directions
    for (CeedInt j = 0; j < 3; j++) {                         // Rows of each Jacobian matrix
      dF[i][j + 1][0] = ((i == j) ? ((gamma - 1.) * (u_sq / 2.)) : 0.) - u[i] * u[j];
      for (CeedInt k = 0; k < 3; k++) {  // Columns of each Jacobian matrix
        dF[i][0][k + 1]     = ((i == k) ? 1. : 0.);
        dF[i][j + 1][k + 1] = ((j == k) ? u[i] : 0.) + ((i == k) ? u[j] : 0.) - ((i == j) ? u[k] : 0.) * (gamma - 1.);
        dF[i][4][k + 1]     = ((i == k) ? (E * gamma / rho - (gamma - 1.) * u_sq / 2.) : 0.) - (gamma - 1.) * u[i] * u[k];
      }
      dF[i][j + 1][4] = ((i == j) ? (gamma - 1.) : 0.);
    }
    dF[i][4][0] = u[i] * ((gamma - 1.) * u_sq - E * gamma / rho);
    dF[i][4][4] = u[i] * gamma;
  }
}

// *****************************************************************************
// Helper function for calculating the covariant length scale in the direction of some 3 element input vector
//
// Where
//  vec         = vector that length is measured in the direction of
//  h           = covariant element length along vec
// *****************************************************************************
CEED_QFUNCTION_HELPER CeedScalar Covariant_length_along_vector(CeedScalar vec[3], const CeedScalar dXdx[3][3]) {
  CeedScalar vec_norm            = sqrt(vec[0] * vec[0] + vec[1] * vec[1] + vec[2] * vec[2]);
  CeedScalar vec_dot_jacobian[3] = {0.0};
  for (CeedInt i = 0; i < 3; i++) {
    for (CeedInt j = 0; j < 3; j++) {
      vec_dot_jacobian[i] += dXdx[j][i] * vec[i];
    }
  }
  CeedScalar norm_vec_dot_jacobian =
      sqrt(vec_dot_jacobian[0] * vec_dot_jacobian[0] + vec_dot_jacobian[1] * vec_dot_jacobian[1] + vec_dot_jacobian[2] * vec_dot_jacobian[2]);
  CeedScalar h = 2.0 * vec_norm / norm_vec_dot_jacobian;
  return h;
}

// *****************************************************************************
// Helper function for computing Tau elements (stabilization constant)
//   Model from:
//     Stabilized Methods for Compressible Flows, Hughes et al 2010
//
//   Spatial criterion #2 - Tau is a 3x3 diagonal matrix
//   Tau[i] = c_tau h[i] Xi(Pe) / rho(A[i]) (no sum)
//
// Where
//   c_tau     = stabilization constant (0.5 is reported as "optimal")
//   h[i]      = 2 length(dxdX[i])
//   Pe        = Peclet number ( Pe = sqrt(u u) / dot(dXdx,u) diffusivity )
//   Xi(Pe)    = coth Pe - 1. / Pe (1. at large local Peclet number )
//   rho(A[i]) = spectral radius of the convective flux Jacobian i, wave speed in direction i
// *****************************************************************************
CEED_QFUNCTION_HELPER void Tau_spatial(CeedScalar Tau_x[3], const CeedScalar dXdx[3][3], const CeedScalar u[3], const CeedScalar sound_speed,
                                       const CeedScalar c_tau) {
  for (CeedInt i = 0; i < 3; i++) {
    // length of element in direction i
    CeedScalar h = 2 / sqrt(dXdx[0][i] * dXdx[0][i] + dXdx[1][i] * dXdx[1][i] + dXdx[2][i] * dXdx[2][i]);
    // fastest wave in direction i
    CeedScalar fastest_wave = fabs(u[i]) + sound_speed;
    Tau_x[i]                = c_tau * h / fastest_wave;
  }
}

// *****************************************************************************
// This QFunction sets the initial conditions for shock tube
// *****************************************************************************
CEED_QFUNCTION(ICsShockTube)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
  const CeedScalar(*X)[CEED_Q_VLA] = (const CeedScalar(*)[CEED_Q_VLA])in[0];
  CeedScalar(*q0)[CEED_Q_VLA]      = (CeedScalar(*)[CEED_Q_VLA])out[0];

  CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
    const CeedScalar x[] = {X[0][i], X[1][i], X[2][i]};
    CeedScalar       q[5];

    Exact_ShockTube(3, 0., x, 5, q, ctx);

    for (CeedInt j = 0; j < 5; j++) q0[j][i] = q[j];
  }
  return 0;
}

// *****************************************************************************
// This QFunction implements the following formulation of Euler equations with explicit time stepping method
//
// This is 3D Euler for compressible gas dynamics in conservation form with state variables of density, momentum density, and total energy density.
//
// State Variables: q = ( rho, U1, U2, U3, E )
//   rho - Mass Density
//   Ui  - Momentum Density,      Ui = rho ui
//   E   - Total Energy Density,  E  = P / (gamma - 1) + rho (u u)/2
//
// Euler Equations:
//   drho/dt + div( U )                   = 0
//   dU/dt   + div( rho (u x u) + P I3 )  = 0
//   dE/dt   + div( (E + P) u )           = 0
//
// Equation of State:
//   P = (gamma - 1) (E - rho (u u) / 2)
//
// Constants:
//   cv              ,  Specific heat, constant volume
//   cp              ,  Specific heat, constant pressure
//   g               ,  Gravity
//   gamma  = cp / cv,  Specific heat ratio
// *****************************************************************************
CEED_QFUNCTION(EulerShockTube)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
  const CeedScalar(*q)[CEED_Q_VLA]     = (const CeedScalar(*)[CEED_Q_VLA])in[0];
  const CeedScalar(*dq)[5][CEED_Q_VLA] = (const CeedScalar(*)[5][CEED_Q_VLA])in[1];
  const CeedScalar(*q_data)            = in[2];
  CeedScalar(*v)[CEED_Q_VLA]           = (CeedScalar(*)[CEED_Q_VLA])out[0];
  CeedScalar(*dv)[5][CEED_Q_VLA]       = (CeedScalar(*)[5][CEED_Q_VLA])out[1];

  const CeedScalar gamma = 1.4;

  ShockTubeContext context = (ShockTubeContext)ctx;
  const CeedScalar Cyzb    = context->Cyzb;
  const CeedScalar Byzb    = context->Byzb;
  const CeedScalar c_tau   = context->c_tau;

  CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
    // Setup
    // -- Interp in
    const CeedScalar rho      = q[0][i];
    const CeedScalar u[3]     = {q[1][i] / rho, q[2][i] / rho, q[3][i] / rho};
    const CeedScalar E        = q[4][i];
    const CeedScalar drho[3]  = {dq[0][0][i], dq[1][0][i], dq[2][0][i]};
    const CeedScalar dU[3][3] = {
        {dq[0][1][i], dq[1][1][i], dq[2][1][i]},
        {dq[0][2][i], dq[1][2][i], dq[2][2][i]},
        {dq[0][3][i], dq[1][3][i], dq[2][3][i]}
    };
    const CeedScalar dE[3] = {dq[0][4][i], dq[1][4][i], dq[2][4][i]};
    CeedScalar       wdetJ, dXdx[3][3];
    QdataUnpack_3D(Q, i, q_data, &wdetJ, dXdx);
    // dU/dx
    CeedScalar du[3][3]        = {{0}};
    CeedScalar drhodx[3]       = {0};
    CeedScalar dEdx[3]         = {0};
    CeedScalar dUdx[3][3]      = {{0}};
    CeedScalar dXdxdXdxT[3][3] = {{0}};
    for (CeedInt j = 0; j < 3; j++) {
      for (CeedInt k = 0; k < 3; k++) {
        du[j][k] = (dU[j][k] - drho[k] * u[j]) / rho;
        drhodx[j] += drho[k] * dXdx[k][j];
        dEdx[j] += dE[k] * dXdx[k][j];
        for (CeedInt l = 0; l < 3; l++) {
          dUdx[j][k] += dU[j][l] * dXdx[l][k];
          dXdxdXdxT[j][k] += dXdx[j][l] * dXdx[k][l];  // dXdx_j,k * dXdx_k,j
        }
      }
    }

    const CeedScalar E_kinetic = 0.5 * rho * (u[0] * u[0] + u[1] * u[1] + u[2] * u[2]), E_internal = E - E_kinetic,
                     P = E_internal * (gamma - 1);  // P = pressure

    // The Physics
    // Zero v and dv so all future terms can safely sum into it
    for (CeedInt j = 0; j < 5; j++) {
      v[j][i] = 0;
      for (CeedInt k = 0; k < 3; k++) dv[k][j][i] = 0;
    }

    // -- Density
    // ---- u rho
    for (CeedInt j = 0; j < 3; j++) dv[j][0][i] += wdetJ * (rho * u[0] * dXdx[j][0] + rho * u[1] * dXdx[j][1] + rho * u[2] * dXdx[j][2]);
    // -- Momentum
    // ---- rho (u x u) + P I3
    for (CeedInt j = 0; j < 3; j++) {
      for (CeedInt k = 0; k < 3; k++) {
        dv[k][j + 1][i] += wdetJ * ((rho * u[j] * u[0] + (j == 0 ? P : 0)) * dXdx[k][0] + (rho * u[j] * u[1] + (j == 1 ? P : 0)) * dXdx[k][1] +
                                    (rho * u[j] * u[2] + (j == 2 ? P : 0)) * dXdx[k][2]);
      }
    }
    // -- Total Energy Density
    // ---- (E + P) u
    for (CeedInt j = 0; j < 3; j++) dv[j][4][i] += wdetJ * (E + P) * (u[0] * dXdx[j][0] + u[1] * dXdx[j][1] + u[2] * dXdx[j][2]);

    // -- YZB stabilization
    if (context->yzb) {
      CeedScalar drho_norm    = 0.0;    // magnitude of the density gradient
      CeedScalar j_vec[3]     = {0.0};  // unit vector aligned with the density gradient
      CeedScalar h_shock      = 0.0;    // element lengthscale
      CeedScalar acoustic_vel = 0.0;    // characteristic velocity, acoustic speed
      CeedScalar tau_shock    = 0.0;    // timescale
      CeedScalar nu_shock     = 0.0;    // artificial diffusion

      // Unit vector aligned with the density gradient
      drho_norm = sqrt(drhodx[0] * drhodx[0] + drhodx[1] * drhodx[1] + drhodx[2] * drhodx[2]);
      for (CeedInt j = 0; j < 3; j++) j_vec[j] = drhodx[j] / (drho_norm + 1e-20);

      if (drho_norm == 0.0) {
        nu_shock = 0.0;
      } else {
        h_shock = Covariant_length_along_vector(j_vec, dXdx);
        h_shock /= Cyzb;
        acoustic_vel = sqrt(gamma * P / rho);
        tau_shock    = h_shock / (2 * acoustic_vel) * pow(drho_norm * h_shock / rho, Byzb);
        nu_shock     = fabs(tau_shock * acoustic_vel * acoustic_vel);
      }

      for (CeedInt j = 0; j < 3; j++) dv[j][0][i] -= wdetJ * nu_shock * drhodx[j];

      for (CeedInt k = 0; k < 3; k++) {
        for (CeedInt j = 0; j < 3; j++) dv[j][k][i] -= wdetJ * nu_shock * du[k][j];
      }

      for (CeedInt j = 0; j < 3; j++) dv[j][4][i] -= wdetJ * nu_shock * dEdx[j];
    }

    // Stabilization
    // Need the Jacobian for the advective fluxes for stabilization
    //    indexed as: jacob_F_conv[direction][flux component][solution component]
    CeedScalar jacob_F_conv[3][5][5] = {{{0.}}};
    ConvectiveFluxJacobian_Euler(jacob_F_conv, rho, u, E, gamma);

    // dqdx collects drhodx, dUdx and dEdx in one vector
    CeedScalar dqdx[5][3];
    for (CeedInt j = 0; j < 3; j++) {
      dqdx[0][j] = drhodx[j];
      dqdx[4][j] = dEdx[j];
      for (CeedInt k = 0; k < 3; k++) dqdx[k + 1][j] = dUdx[k][j];
    }

    // strong_conv = dF/dq * dq/dx    (Strong convection)
    CeedScalar strong_conv[5] = {0};
    for (CeedInt j = 0; j < 3; j++) {
      for (CeedInt k = 0; k < 5; k++) {
        for (CeedInt l = 0; l < 5; l++) strong_conv[k] += jacob_F_conv[j][k][l] * dqdx[l][j];
      }
    }

    // Stabilization
    // -- Tau elements
    const CeedScalar sound_speed = sqrt(gamma * P / rho);
    CeedScalar       Tau_x[3]    = {0.};
    Tau_spatial(Tau_x, dXdx, u, sound_speed, c_tau);

    CeedScalar stab[5][3] = {0};
    switch (context->stabilization) {
      case 0:  // Galerkin
        break;
      case 1:  // SU
        for (CeedInt j = 0; j < 3; j++) {
          for (CeedInt k = 0; k < 5; k++) {
            for (CeedInt l = 0; l < 5; l++) {
              stab[k][j] += jacob_F_conv[j][k][l] * Tau_x[j] * strong_conv[l];
            }
          }
        }
        for (CeedInt j = 0; j < 5; j++) {
          for (CeedInt k = 0; k < 3; k++) dv[k][j][i] -= wdetJ * (stab[j][0] * dXdx[k][0] + stab[j][1] * dXdx[k][1] + stab[j][2] * dXdx[k][2]);
        }
        break;
    }
  }
  return 0;
}