// 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
/// Density current initial condition and operator for Navier-Stokes example using PETSc

// Model from:
//   Semi-Implicit Formulations of the Navier-Stokes Equations: Application to
//   Nonhydrostatic Atmospheric Modeling, Giraldo, Restelli, and Lauter (2010).
#include <ceed/types.h>
#ifndef CEED_RUNNING_JIT_PASS
#include <math.h>
#endif

#include "newtonian_state.h"
#include "newtonian_types.h"
#include "utils.h"

typedef struct DensityCurrentContext_ *DensityCurrentContext;
struct DensityCurrentContext_ {
  CeedScalar                       theta0;
  CeedScalar                       thetaC;
  CeedScalar                       P0;
  CeedScalar                       N;
  CeedScalar                       rc;
  CeedScalar                       center[3];
  CeedScalar                       dc_axis[3];
  struct NewtonianIdealGasContext_ newtonian_ctx;
};

// *****************************************************************************
// This function sets the initial conditions and the boundary conditions
//
// These initial conditions are given in terms of potential temperature and Exner pressure and then converted to density and total energy.
//   Initial momentum density is zero.
//
// Initial Conditions:
//   Potential Temperature:
//     theta = thetabar + delta_theta
//       thetabar   = theta0 exp( N**2 z / g )
//       delta_theta = r <= rc : thetaC(1 + cos(pi r/rc)) / 2
//                     r > rc : 0
//         r        = sqrt( (x - xc)**2 + (y - yc)**2 + (z - zc)**2 )
//         with (xc,yc,zc) center of domain, rc characteristic radius of thermal bubble
//   Exner Pressure:
//     Pi = Pibar + deltaPi
//       Pibar      = 1. + g**2 (exp( - N**2 z / g ) - 1) / (cp theta0 N**2)
//       deltaPi    = 0 (hydrostatic balance)
//   Velocity/Momentum Density:
//     Ui = ui = 0
//
// Conversion to Conserved Variables:
//   rho = P0 Pi**(cv/Rd) / (Rd theta)
//   E   = rho (cv T + (u u)/2 + g z)
//
//  Boundary Conditions:
//    Mass Density:
//      0.0 flux
//    Momentum Density:
//      0.0
//    Energy Density:
//      0.0 flux
//
// Constants:
//   theta0          ,  Potential temperature constant
//   thetaC          ,  Potential temperature perturbation
//   P0              ,  Pressure at the surface
//   N               ,  Brunt-Vaisala frequency
//   cv              ,  Specific heat, constant volume
//   cp              ,  Specific heat, constant pressure
//   Rd     = cp - cv,  Specific heat difference
//   g               ,  Gravity
//   rc              ,  Characteristic radius of thermal bubble
//   center          ,  Location of bubble center
//   dc_axis         ,  Axis of density current cylindrical anomaly, or {0,0,0} for spherically symmetric
// *****************************************************************************

// *****************************************************************************
// This helper function provides support for the exact, time-dependent solution
//   (currently not implemented) and IC formulation for density current
// *****************************************************************************
CEED_QFUNCTION_HELPER State Exact_DC(CeedInt dim, CeedScalar time, const CeedScalar X[], CeedInt Nf, void *ctx) {
  // Context
  const DensityCurrentContext context = (DensityCurrentContext)ctx;
  const CeedScalar            theta0  = context->theta0;
  const CeedScalar            thetaC  = context->thetaC;
  const CeedScalar            P0      = context->P0;
  const CeedScalar            N       = context->N;
  const CeedScalar            rc      = context->rc;
  const CeedScalar           *center  = context->center;
  const CeedScalar           *dc_axis = context->dc_axis;
  NewtonianIdealGasContext    gas     = &context->newtonian_ctx;
  const CeedScalar            cp      = gas->cp;
  const CeedScalar            cv      = gas->cv;
  const CeedScalar            Rd      = cp - cv;
  const CeedScalar           *g_vec   = gas->g;
  const CeedScalar            g       = -g_vec[2];

  // Setup
  // -- Coordinates
  const CeedScalar x = X[0];
  const CeedScalar y = X[1];
  const CeedScalar z = X[2];

  // -- Potential temperature, density current
  CeedScalar rr[3] = {x - center[0], y - center[1], z - center[2]};
  // (I - q q^T) r: distance from dc_axis (or from center if dc_axis is the zero vector)
  for (CeedInt i = 0; i < 3; i++) rr[i] -= dc_axis[i] * Dot3(dc_axis, rr);
  const CeedScalar r           = sqrt(Dot3(rr, rr));
  const CeedScalar delta_theta = r <= rc ? thetaC * (1. + cos(M_PI * r / rc)) / 2. : 0.;
  const CeedScalar theta       = theta0 * exp(Square(N) * z / g) + delta_theta;

  // -- Exner pressure, hydrostatic balance
  const CeedScalar Pi = 1. + Square(g) * (exp(-Square(N) * z / g) - 1.) / (cp * theta0 * Square(N));

  // Initial Conditions
  CeedScalar Y[5] = {0.};
  Y[0]            = P0 * pow(Pi, cp / Rd);
  Y[1]            = 0.0;
  Y[2]            = 0.0;
  Y[3]            = 0.0;
  Y[4]            = Pi * theta;

  return StateFromY(gas, Y);
}

// *****************************************************************************
// This QFunction sets the initial conditions for density current
// *****************************************************************************
CEED_QFUNCTION(ICsDC)(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];

  const DensityCurrentContext    context = (DensityCurrentContext)ctx;
  const NewtonianIdealGasContext gas     = &context->newtonian_ctx;

  CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
    const CeedScalar x[]  = {X[0][i], X[1][i], X[2][i]};
    State            s    = Exact_DC(3, 0., x, 5, ctx);
    CeedScalar       q[5] = {0};
    StateToQ(gas, s, q, gas->state_var);
    for (CeedInt j = 0; j < 5; j++) q0[j][i] = q[j];
  }
  return 0;
}