#include "macrodef.h"
#include "bssn_rhs.h"
#include "share_func.h"
#include "tool.h"
#include <vector>

namespace
{
    // Reuse the temporary workspace across block calls to avoid repeated heap churn
    // in the EScalar wrapper. MPI ranks execute this path sequentially, so a single
    // process-local buffer is sufficient here.
    std::vector<double> g_escalar_tmp_store;
}

#ifdef fortran1
#define f_frpotential frpotential
#endif
#ifdef fortran2
#define f_frpotential FRPOTENTIAL
#endif
#ifdef fortran3
#define f_frpotential frpotential_
#endif

extern "C"
{
    void f_frpotential(int *, double *, double *, double *);
}

int f_compute_rhs_bssn_escalar_c(int *ex, double &T,
                                 double *X, double *Y, double *Z,
                                 double *chi, double *trK,
                                 double *dxx, double *gxy, double *gxz, double *dyy, double *gyz, double *dzz,
                                 double *Axx, double *Axy, double *Axz, double *Ayy, double *Ayz, double *Azz,
                                 double *Gamx, double *Gamy, double *Gamz,
                                 double *Lap, double *betax, double *betay, double *betaz,
                                 double *dtSfx, double *dtSfy, double *dtSfz,
                                 double *Sphi, double *Spi,
                                 double *chi_rhs, double *trK_rhs,
                                 double *gxx_rhs, double *gxy_rhs, double *gxz_rhs, double *gyy_rhs, double *gyz_rhs, double *gzz_rhs,
                                 double *Axx_rhs, double *Axy_rhs, double *Axz_rhs, double *Ayy_rhs, double *Ayz_rhs, double *Azz_rhs,
                                 double *Gamx_rhs, double *Gamy_rhs, double *Gamz_rhs,
                                 double *Lap_rhs, double *betax_rhs, double *betay_rhs, double *betaz_rhs,
                                 double *dtSfx_rhs, double *dtSfy_rhs, double *dtSfz_rhs,
                                 double *Sphi_rhs, double *Spi_rhs,
                                 double *rho, double *Sx, double *Sy, double *Sz,
                                 double *Sxx, double *Sxy, double *Sxz, double *Syy, double *Syz, double *Szz,
                                 double *Gamxxx, double *Gamxxy, double *Gamxxz, double *Gamxyy, double *Gamxyz, double *Gamxzz,
                                 double *Gamyxx, double *Gamyxy, double *Gamyxz, double *Gamyyy, double *Gamyyz, double *Gamyzz,
                                 double *Gamzxx, double *Gamzxy, double *Gamzxz, double *Gamzyy, double *Gamzyz, double *Gamzzz,
                                 double *Rxx, double *Rxy, double *Rxz, double *Ryy, double *Ryz, double *Rzz,
                                 double *ham_Res, double *movx_Res, double *movy_Res, double *movz_Res,
                                 double *Gmx_Res, double *Gmy_Res, double *Gmz_Res,
                                 int &Symmetry, int &Lev, double &eps, int &co)
{
    const int nx = ex[0], ny = ex[1], nz = ex[2];
    const int all = nx * ny * nz;

    const size_t workspace_size = size_t(all) * 17;
    if (g_escalar_tmp_store.size() < workspace_size)
        g_escalar_tmp_store.resize(workspace_size);

    double *tmp_ptr = g_escalar_tmp_store.data();
    auto alloc_tmp = [&](int n = 1) -> double *
    {
        double *ptr = tmp_ptr;
        tmp_ptr += size_t(all) * n;
        return ptr;
    };

    double *chix = alloc_tmp(), *chiy = alloc_tmp(), *chiz = alloc_tmp();
    double *Kx = alloc_tmp(), *Ky = alloc_tmp(), *Kz = alloc_tmp();
    double *fxx = alloc_tmp(), *fxy = alloc_tmp(), *fxz = alloc_tmp();
    double *fyy = alloc_tmp(), *fyz = alloc_tmp(), *fzz = alloc_tmp();
    double *Lapx = alloc_tmp(), *Lapy = alloc_tmp(), *Lapz = alloc_tmp();
    double *V = alloc_tmp(), *dVdSphi = alloc_tmp();

    const double ZEO = 0.0, ONE = 1.0, TWO = 2.0, HALF = 0.5;
    const double SSS[3] = {1.0, 1.0, 1.0};

    fderivs(ex, chi, chix, chiy, chiz, X, Y, Z, 1.0, 1.0, 1.0, Symmetry, Lev);
    fderivs(ex, Lap, Lapx, Lapy, Lapz, X, Y, Z, 1.0, 1.0, 1.0, Symmetry, Lev);
    fderivs(ex, Sphi, Kx, Ky, Kz, X, Y, Z, 1.0, 1.0, 1.0, Symmetry, Lev);
    fdderivs(ex, Sphi, fxx, fxy, fxz, fyy, fyz, fzz, X, Y, Z, 1.0, 1.0, 1.0, Symmetry, Lev);

    f_frpotential(ex, Sphi, V, dVdSphi);

    for (int i = 0; i < all; ++i)
    {
        const double alpn1 = Lap[i] + ONE;
        const double chin1 = chi[i] + ONE;
        const double gxx = dxx[i] + ONE;
        const double gyy = dyy[i] + ONE;
        const double gzz = dzz[i] + ONE;
        const double det = gxx * gyy * gzz + gxy[i] * gyz[i] * gxz[i] + gxz[i] * gxy[i] * gyz[i]
                         - gxz[i] * gyy * gxz[i] - gxy[i] * gxy[i] * gzz - gxx * gyz[i] * gyz[i];
        const double gupxx = (gyy * gzz - gyz[i] * gyz[i]) / det;
        const double gupxy = -(gxy[i] * gzz - gyz[i] * gxz[i]) / det;
        const double gupxz = (gxy[i] * gyz[i] - gyy * gxz[i]) / det;
        const double gupyy = (gxx * gzz - gxz[i] * gxz[i]) / det;
        const double gupyz = -(gxx * gyz[i] - gxy[i] * gxz[i]) / det;
        const double gupzz = (gxx * gyy - gxy[i] * gxy[i]) / det;

        Sphi_rhs[i] = alpn1 * Spi[i];

        Spi_rhs[i] = gupxx * fxx[i] + gupyy * fyy[i] + gupzz * fzz[i]
                   + TWO * (gupxy * fxy[i] + gupxz * fxz[i] + gupyz * fyz[i])
                   - ((Gamx[i] + (gupxx * chix[i] + gupxy * chiy[i] + gupxz * chiz[i]) / TWO / chin1) * Kx[i]
                   +  (Gamy[i] + (gupxy * chix[i] + gupyy * chiy[i] + gupyz * chiz[i]) / TWO / chin1) * Ky[i]
                   +  (Gamz[i] + (gupxz * chix[i] + gupyz * chiy[i] + gupzz * chiz[i]) / TWO / chin1) * Kz[i]);

        Spi_rhs[i] = Spi_rhs[i] * alpn1
                   + gupxx * Lapx[i] * Kx[i] + gupxy * Lapx[i] * Ky[i] + gupxz * Lapx[i] * Kz[i]
                   + gupxy * Lapy[i] * Kx[i] + gupyy * Lapy[i] * Ky[i] + gupyz * Lapy[i] * Kz[i]
                   + gupxz * Lapz[i] * Kx[i] + gupyz * Lapz[i] * Ky[i] + gupzz * Lapz[i] * Kz[i];

        Spi_rhs[i] = Spi_rhs[i] * chin1 + alpn1 * (trK[i] * Spi[i] - dVdSphi[i]);

        rho[i] = chin1 * ((gupxx * Kx[i] * Kx[i] + gupyy * Ky[i] * Ky[i] + gupzz * Kz[i] * Kz[i]) * HALF
               + gupxy * Kx[i] * Ky[i] + gupxz * Kx[i] * Kz[i] + gupyz * Ky[i] * Kz[i])
               + Spi[i] * Spi[i] * HALF + V[i];
        Sx[i] = -Spi[i] * Kx[i];
        Sy[i] = -Spi[i] * Ky[i];
        Sz[i] = -Spi[i] * Kz[i];

        const double pressure = (rho[i] - Spi[i] * Spi[i]) / chin1;
        Sxx[i] = Kx[i] * Kx[i] - pressure * gxx;
        Sxy[i] = Kx[i] * Ky[i] - pressure * gxy[i];
        Sxz[i] = Kx[i] * Kz[i] - pressure * gxz[i];
        Syy[i] = Ky[i] * Ky[i] - pressure * gyy;
        Syz[i] = Ky[i] * Kz[i] - pressure * gyz[i];
        Szz[i] = Kz[i] * Kz[i] - pressure * gzz;
    }

    if (f_compute_rhs_bssn(ex, T, X, Y, Z,
                           chi, trK,
                           dxx, gxy, gxz, dyy, gyz, dzz,
                           Axx, Axy, Axz, Ayy, Ayz, Azz,
                           Gamx, Gamy, Gamz,
                           Lap, betax, betay, betaz,
                           dtSfx, dtSfy, dtSfz,
                           chi_rhs, trK_rhs,
                           gxx_rhs, gxy_rhs, gxz_rhs, gyy_rhs, gyz_rhs, gzz_rhs,
                           Axx_rhs, Axy_rhs, Axz_rhs, Ayy_rhs, Ayz_rhs, Azz_rhs,
                           Gamx_rhs, Gamy_rhs, Gamz_rhs,
                           Lap_rhs, betax_rhs, betay_rhs, betaz_rhs,
                           dtSfx_rhs, dtSfy_rhs, dtSfz_rhs,
                           rho, Sx, Sy, Sz,
                           Sxx, Sxy, Sxz, Syy, Syz, Szz,
                           Gamxxx, Gamxxy, Gamxxz, Gamxyy, Gamxyz, Gamxzz,
                           Gamyxx, Gamyxy, Gamyxz, Gamyyy, Gamyyz, Gamyzz,
                           Gamzxx, Gamzxy, Gamzxz, Gamzyy, Gamzyz, Gamzzz,
                           Rxx, Rxy, Rxz, Ryy, Ryz, Rzz,
                           ham_Res, movx_Res, movy_Res, movz_Res,
                           Gmx_Res, Gmy_Res, Gmz_Res,
                           Symmetry, Lev, eps, co))
        return 1;

    lopsided_kodis(ex, X, Y, Z, Sphi, Sphi_rhs, betax, betay, betaz, Symmetry, SSS, eps);
    lopsided_kodis(ex, X, Y, Z, Spi, Spi_rhs, betax, betay, betaz, Symmetry, SSS, eps);

    for (int i = 0; i < all; ++i)
    {
        if (Sphi_rhs[i] != Sphi_rhs[i] || Spi_rhs[i] != Spi_rhs[i] || rho[i] != rho[i])
            return 1;
    }

    return 0;
}