/*!******************************************************************
 * \file blob2d.cxx
 *
 *       2D simulations
 *
 *        NR Walkden, B Dudson  20 January 2012
 *******************************************************************/

#include <bout/derivs.hxx>         // To use DDZ()
#include <bout/invert_laplace.hxx> // Laplacian inversion
#include <bout/physicsmodel.hxx>   // Commonly used BOUT++ components

/// 2D drift-reduced model, mainly used for blob studies
///
///
class Blob2D : public PhysicsModel {
private:
  // Evolving variables
  Field3D n, omega; ///< Density and vorticity

  // Auxilliary variables
  Field3D phi; ///< Electrostatic potential

  // Parameters
  BoutReal rho_s;   ///< Bohm gyro radius
  BoutReal Omega_i; ///< Ion cyclotron frequency
  BoutReal c_s;     ///< Bohm sound speed

  // Constants to calculate the parameters
  BoutReal Te0; ///< Isothermal temperature [eV]
  BoutReal B0;  ///< Constant magnetic field [T]
  BoutReal m_i; ///< Ion mass [kg]
  BoutReal m_e; ///< Electron mass [kg]
  BoutReal e;   ///< Electron charge [C]

  BoutReal D_n, D_vort; ///< Diffusion coefficients
  BoutReal R_c;         ///< Radius of curvature
  BoutReal L_par;       ///< Parallel connection length

  // Model options
  bool boussinesq;   ///< Use the Boussinesq approximation in vorticity
  bool compressible; ///< If allow inclusion of n grad phi term in density evolution
  bool sheath;       ///< Sheath connected?

  std::unique_ptr<Laplacian> phiSolver{
      nullptr}; ///< Performs Laplacian inversions to calculate phi

protected:
  int init(bool UNUSED(restarting)) override {

    /******************Reading options *****************/

    auto& globalOptions = Options::root();
    auto& options = globalOptions["model"];

    // Load system parameters
    Te0 = options["Te0"].doc("Temperature in eV").withDefault(30.0);

    e = options["e"].withDefault(1.602e-19);
    m_i = options["m_i"].withDefault(2 * 1.667e-27);
    m_e = options["m_e"].withDefault(9.11e-31);

    D_vort = options["D_vort"].doc("Viscous diffusion coefficient").withDefault(0.0);
    D_n = options["D_n"].doc("Density diffusion coefficient").withDefault(0.0);

    R_c = options["R_c"].doc("Radius of curvature").withDefault(1.5);
    L_par = options["L_par"].doc("Parallel connection length").withDefault(10.0);

    B0 = options["B0"].doc("Value of magnetic field strength").withDefault(0.35);

    // System option switches

    compressible = options["compressible"]
                       .doc("Compressible ExB term in density equation")
                       .withDefault(false);
    boussinesq = options["boussinesq"]
                     .doc("Use Boussinesq approximation in vorticity")
                     .withDefault(true);
    sheath = options["sheath"].doc("Sheath closure").withDefault(true);

    /***************Calculate the Parameters **********/

    Omega_i = e * B0 / m_i;    // Cyclotron Frequency
    c_s = sqrt(e * Te0 / m_i); // Bohm sound speed
    rho_s = c_s / Omega_i;     // Bohm gyro-radius

    output.write(
        "\n\n\t----------Parameters: ------------ \n\tOmega_i = {:e} /s,\n\tc_s = "
        "{:e} m/s,\n\trho_s = {:e} m\n",
        Omega_i, c_s, rho_s);

    // Calculate delta_*, blob size scaling
    output.write("\tdelta_* = rho_s * (dn/n) * {:e} ",
                 pow(L_par * L_par / (R_c * rho_s), 1. / 5));

    /************ Create a solver for potential ********/

    Options& boussinesq_options = Options::root()["phiBoussinesq"];
    Options& non_boussinesq_options = Options::root()["phiSolver"];

    if (boussinesq) {
      // BOUT.inp section "phiBoussinesq"
      phiSolver = Laplacian::create(&boussinesq_options);
      // Mark other section as conditionally used so we don't get errors from unused
      // options
      non_boussinesq_options.setConditionallyUsed();
    } else {
      // BOUT.inp section "phiSolver"
      phiSolver = Laplacian::create(&non_boussinesq_options);
      // Mark other section as conditionally used so we don't get errors from unused
      // options
      boussinesq_options.setConditionallyUsed();
    }
    phi = 0.0; // Starting guess for first solve (if iterative)

    /************ Tell BOUT++ what to solve ************/
    SOLVE_FOR(n, omega);

    return 0;
  }

  /// Add variables to the output. This can be used to calculate
  /// diagnostics
  ///
  /// @param[inout] state  A nested dictionary that can be added to
  void outputVars(Options& state) override {
    // Set time-varying quantity (assignRepeat)
    state["phi"].assignRepeat(phi).setAttributes({{"units", "V"},
                                                  {"conversion", Te0},
                                                  {"standard_name", "potential"},
                                                  {"long_name", "Plasma potential"}});

    // Force updates to non-varying quantities
    state["rho_s"].force(rho_s).setAttributes(
        {{"units", "m"},
         {"conversion", 1},
         {"standard_name", "length normalisation"},
         {"long_name", "Gyro-radius length normalisation"}});

    state["c_s"].force(c_s).setAttributes({{"units", "m/s"},
                                           {"conversion", 1},
                                           {"standard_name", "velocity normalisation"},
                                           {"long_name", "Sound speed normalisation"}});

    state["Omega_i"].force(Omega_i).setAttributes(
        {{"units", "s^-1"},
         {"conversion", 1},
         {"standard_name", "frequency normalisation"},
         {"long_name", "Cyclotron frequency normalisation"}});
  }

  int rhs(BoutReal UNUSED(t)) override {

    // Run communications
    ////////////////////////////////////////////////////////////////////////////
    mesh->communicate(n, omega);

    // Invert div(n grad(phi)) = grad(n) grad(phi) + n Delp_perp^2(phi) = omega
    ////////////////////////////////////////////////////////////////////////////

    if (!boussinesq) {
      // Including full density in vorticit inversion
      phiSolver->setCoefC(n); // Update the 'C' coefficient. See invert_laplace.hxx
      phi = phiSolver->solve(omega / n, phi); // Use previous solution as guess
    } else {
      // Background density only (1 in normalised units)
      phi = phiSolver->solve(omega);
    }

    mesh->communicate(phi);

    // Density Evolution
    /////////////////////////////////////////////////////////////////////////////

    ddt(n) = -bracket(phi, n, BRACKET_ARAKAWA) // ExB term
             + 2 * DDZ(n) * (rho_s / R_c)      // Curvature term
             + D_n * Delp2(n);                 // Diffusion term
    if (compressible) {
      ddt(n) -= 2 * n * DDZ(phi) * (rho_s / R_c); // ExB Compression term
    }

    if (sheath) {
      // Sheath closure
      ddt(n) += n * phi * (rho_s / L_par);
    }

    // Vorticity evolution
    /////////////////////////////////////////////////////////////////////////////

    ddt(omega) = -bracket(phi, omega, BRACKET_ARAKAWA) // ExB term
                 + 2 * DDZ(n) * (rho_s / R_c) / n      // Curvature term
                 + D_vort * Delp2(omega) / n           // Viscous diffusion term
        ;

    if (sheath) {
      ddt(omega) += phi * (rho_s / L_par);
    }

    return 0;
  }
};

// Define a standard main() function
BOUTMAIN(Blob2D);
