// Copyright (c) 2017-2022, 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
/// Helper functions for computing stabilization terms of a newtonian simulation

#ifndef stabilization_h
#define stabilization_h

#include <ceed.h>

#include "newtonian_state.h"

// *****************************************************************************
// Helper function for computing the variation in primitive variables, given Tau_d
// *****************************************************************************
CEED_QFUNCTION_HELPER void dYFromTau(CeedScalar Y[5], CeedScalar Tau_d[3], CeedScalar dY[5]) {
  dY[0] = Tau_d[0] * Y[0];
  dY[1] = Tau_d[1] * Y[1];
  dY[2] = Tau_d[1] * Y[2];
  dY[3] = Tau_d[1] * Y[3];
  dY[4] = Tau_d[2] * Y[4];
}

// *****************************************************************************
// Helper functions for computing the stabilization terms
// *****************************************************************************
CEED_QFUNCTION_HELPER void StabilizationMatrix(NewtonianIdealGasContext gas, State s, CeedScalar Tau_d[3], CeedScalar R[5], const CeedScalar x[3],
                                               CeedScalar stab[5][3]) {
  CeedScalar        dY[5];
  const CeedScalar  dx_i[3] = {0};
  StateConservative dF[3];
  // Zero stab so all future terms can safely sum into it
  for (CeedInt i = 0; i < 5; i++) {
    for (CeedInt j = 0; j < 3; j++) stab[i][j] = 0;
  }
  dYFromTau(R, Tau_d, dY);
  State ds = StateFromY_fwd(gas, s, dY, x, dx_i);
  FluxInviscid_fwd(gas, s, ds, dF);
  for (CeedInt i = 0; i < 3; i++) {
    CeedScalar dF_i[5];
    UnpackState_U(dF[i], dF_i);
    for (CeedInt j = 0; j < 5; j++) stab[j][i] += dF_i[j];
  }
}

CEED_QFUNCTION_HELPER void Stabilization(NewtonianIdealGasContext gas, State s, CeedScalar Tau_d[3], State ds[3], CeedScalar U_dot[5],
                                         const CeedScalar body_force[5], const CeedScalar x[3], CeedScalar stab[5][3]) {
  // -- Stabilization method: none (Galerkin), SU, or SUPG
  CeedScalar R[5] = {0};
  switch (gas->stabilization) {
    case STAB_NONE:
      break;
    case STAB_SU:
      FluxInviscidStrong(gas, s, ds, R);
      break;
    case STAB_SUPG:
      FluxInviscidStrong(gas, s, ds, R);
      for (CeedInt j = 0; j < 5; j++) R[j] += U_dot[j] - body_force[j];
      break;
  }
  StabilizationMatrix(gas, s, Tau_d, R, x, stab);
}

// *****************************************************************************
// Helper function for computing Tau elements (stabilization constant)
//   Model from:
//     PHASTA
//
//   Tau[i] = itau=0 which is diagonal-Shakib (3 values still but not spatial)
// *****************************************************************************
CEED_QFUNCTION_HELPER void Tau_diagPrim(NewtonianIdealGasContext gas, State s, const CeedScalar dXdx[3][3], const CeedScalar dt,
                                        CeedScalar Tau_d[3]) {
  // Context
  const CeedScalar Ctau_t = gas->Ctau_t;
  const CeedScalar Ctau_v = gas->Ctau_v;
  const CeedScalar Ctau_C = gas->Ctau_C;
  const CeedScalar Ctau_M = gas->Ctau_M;
  const CeedScalar Ctau_E = gas->Ctau_E;
  const CeedScalar cv     = gas->cv;
  const CeedScalar mu     = gas->mu;
  const CeedScalar u[3]   = {s.Y.velocity[0], s.Y.velocity[1], s.Y.velocity[2]};
  const CeedScalar rho    = s.U.density;

  CeedScalar gijd[6];
  CeedScalar tau;
  CeedScalar dts;
  CeedScalar fact;

  gijd[0] = dXdx[0][0] * dXdx[0][0] + dXdx[1][0] * dXdx[1][0] + dXdx[2][0] * dXdx[2][0];

  gijd[1] = dXdx[0][0] * dXdx[0][1] + dXdx[1][0] * dXdx[1][1] + dXdx[2][0] * dXdx[2][1];

  gijd[2] = dXdx[0][1] * dXdx[0][1] + dXdx[1][1] * dXdx[1][1] + dXdx[2][1] * dXdx[2][1];

  gijd[3] = dXdx[0][0] * dXdx[0][2] + dXdx[1][0] * dXdx[1][2] + dXdx[2][0] * dXdx[2][2];

  gijd[4] = dXdx[0][1] * dXdx[0][2] + dXdx[1][1] * dXdx[1][2] + dXdx[2][1] * dXdx[2][2];

  gijd[5] = dXdx[0][2] * dXdx[0][2] + dXdx[1][2] * dXdx[1][2] + dXdx[2][2] * dXdx[2][2];

  dts = Ctau_t / dt;

  tau = rho * rho *
            ((4. * dts * dts) + u[0] * (u[0] * gijd[0] + 2. * (u[1] * gijd[1] + u[2] * gijd[3])) + u[1] * (u[1] * gijd[2] + 2. * u[2] * gijd[4]) +
             u[2] * u[2] * gijd[5]) +
        Ctau_v * mu * mu *
            (gijd[0] * gijd[0] + gijd[2] * gijd[2] + gijd[5] * gijd[5] + +2. * (gijd[1] * gijd[1] + gijd[3] * gijd[3] + gijd[4] * gijd[4]));

  fact = sqrt(tau);

  Tau_d[0] = Ctau_C * fact / (rho * (gijd[0] + gijd[2] + gijd[5])) * 0.125;

  Tau_d[1] = Ctau_M / fact;
  Tau_d[2] = Ctau_E / (fact * cv);

  // consider putting back the way I initially had it
  // Ctau_E * Tau_d[1] /cv to avoid a division if the compiler is smart enough to see that cv IS a constant that it could invert once for all elements
  // but in that case energy tau is scaled by the product of Ctau_E * Ctau_M
  // OR we could absorb cv into Ctau_E but this puts more burden on user to know how to change constants with a change of fluid or units.  Same for
  // Ctau_v * mu * mu IF AND ONLY IF we don't add viscosity law =f(T)
}

// *****************************************************************************

#endif  // stabilization_h