/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Authors: Aaron Dotter, Emily Boudreaux
//   Last Modified: March 06, 2025
//
//   4DSSE is free software; you can use it and/or modify
//   it under the terms and restrictions the GNU General Library Public
//   License version 3 (GPLv3) as published by the Free Software Foundation.
//
//   4DSSE is distributed in the hope that it will be useful,
//   but WITHOUT ANY WARRANTY; without even the implied warranty of
//   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//   See the GNU Library General Public License for more details.
//
//   You should have received a copy of the GNU Library General Public License
//   along with this software; if not, write to the Free Software
//   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
//
// *********************************************************************** */
// this file contains a C++ conversion of Frank Timmes' fortran code
// helmholtz.f90, which implements the Helmholtz Free Energy EOS described
// by Timmes & Swesty (2000) doi:10.1086/313304  -- Aaron Dotter 2025

#ifndef HELM_H
#define HELM_H

#define IMAX 541
#define JMAX 201

#include <array>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <string>
#include <format>

/**
 * @brief 2D array template alias.
 * @tparam T Type of the array elements.
 * @tparam J Number of columns.
 * @tparam I Number of rows.
 */
template <typename T, std::size_t J, std::size_t I>
using array2D = std::array<std::array<T, J>, I>;

double** heap_allocate_contiguous_2D_memory(int rows, int cols);

void heap_deallocate_contiguous_2D_memory(double **array);

namespace helmholtz
{
    const double tlo = 3.0, thi = 13.0;
    const double dlo = -12.0, dhi = 15.0;
    const double tstp = (thi - tlo) / (JMAX - 1), tstpi = 1 / tstp;
    const double dstp = (dhi - dlo) / (IMAX - 1), dstpi = 1 / dstp;

    /**
     * @struct HELMTable
     * @brief Structure to hold the Helmholtz EOS table data.
     */
    struct HELMTable
    {
        bool loaded = false;
        const int imax = IMAX, jmax = JMAX;
        double t[JMAX], d[IMAX];
        double** f;
        double** fd;
        double** ft;
        double** fdd;
        double** ftt;
        double** fdt;
        double** fddt;
        double** fdtt;
        double** fddtt;
        double** dpdf;
        double** dpdfd;
        double** dpdft;
        double** dpdfdt;
        double** ef;
        double** efd;
        double** eft;
        double** efdt;
        double** xf;
        double** xfd;
        double** xft;
        double** xfdt;

        double dt_sav[JMAX], dt2_sav[JMAX], dti_sav[JMAX], dt2i_sav[JMAX];
        double dd_sav[IMAX], dd2_sav[IMAX], ddi_sav[IMAX];

        // Constructor
        HELMTable() {
            f = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            fd = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            ft = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            fdd = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            ftt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            fdt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            fddt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            fdtt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            fddtt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            dpdf = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            dpdfd = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            dpdft = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            dpdfdt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            ef = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            efd = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            eft = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            efdt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            xf = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            xfd = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            xft = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
            xfdt = heap_allocate_contiguous_2D_memory(IMAX, JMAX);
        }

        // Destructor
        ~HELMTable() {
            heap_deallocate_contiguous_2D_memory(f);
            heap_deallocate_contiguous_2D_memory(fd);
            heap_deallocate_contiguous_2D_memory(ft);
            heap_deallocate_contiguous_2D_memory(fdd);
            heap_deallocate_contiguous_2D_memory(ftt);
            heap_deallocate_contiguous_2D_memory(fdt);
            heap_deallocate_contiguous_2D_memory(fddt);
            heap_deallocate_contiguous_2D_memory(fdtt);
            heap_deallocate_contiguous_2D_memory(fddtt);
            heap_deallocate_contiguous_2D_memory(dpdf);
            heap_deallocate_contiguous_2D_memory(dpdfd);
            heap_deallocate_contiguous_2D_memory(dpdft);
            heap_deallocate_contiguous_2D_memory(dpdfdt);
            heap_deallocate_contiguous_2D_memory(ef);
            heap_deallocate_contiguous_2D_memory(efd);
            heap_deallocate_contiguous_2D_memory(eft);
            heap_deallocate_contiguous_2D_memory(efdt);
            heap_deallocate_contiguous_2D_memory(xf);
            heap_deallocate_contiguous_2D_memory(xfd);
            heap_deallocate_contiguous_2D_memory(xft);
            heap_deallocate_contiguous_2D_memory(xfdt);
        }

        friend std::ostream& operator<<(std::ostream& os, const helmholtz::HELMTable& table) {
            if (!table.loaded) {
                os << "HELMTable not loaded\n";
                return os;
            }
            double tMin = std::numeric_limits<double>::max(), tMax = std::numeric_limits<double>::lowest();
            double dMin = std::numeric_limits<double>::max(), dMax = std::numeric_limits<double>::lowest();

            for (const auto& temp : table.t) {
                if (temp < tMin) tMin = temp;
                if (temp > tMax) tMax = temp;
            }

            for (const auto& dens : table.d) {
                if (dens < dMin) dMin = dens;
                if (dens > dMax) dMax = dens;
            }

            os << "HELMTable Data:\n";
            os << "  imax: " << table.imax << ", jmax: " << table.jmax << "\n";
            os << "  Temperature Range: [" << tMin << ", " << tMax << "]\n";
            os << "  Density Range: [" << dMin << ", " << dMax << "]\n";

            return os;
        }


    };

    /**
     * @struct EOSInput
     * @brief Structure to hold the input parameters for the EOS calculation.
     */
    struct EOSInput
    {
        double T; ///< Temperature.
        double rho; ///< Density.
        double abar; ///< Mean atomic mass.
        double zbar; ///< Mean atomic number.

        friend std::ostream& operator<<(std::ostream& os, const helmholtz::EOSInput& eosInput) {
            os << "EOSInput Data:\n";
            os << "  Temperature: " << eosInput.T << "\n";
            os << "  Density: " << eosInput.rho << "\n";
            os << "  Mean Atomic Mass: " << eosInput.abar << "\n";
            os << "  Mean Atomic Number: " << eosInput.zbar << "\n";

            return os;
        }
    };

    /**
     * @struct EOS
     * @brief Structure to hold the output parameters and derivatives of the EOS calculation.
     */
    struct EOS
    {
        // output
        double ye, etaele, xnefer; //
        double ptot, pgas, prad;   // pressure
        double etot, egas, erad;   // energy
        double stot, sgas, srad;   // entropy

        // derivatives
        double dpresdd, dpresdt, dpresda, dpresdz;
        double dentrdd, dentrdt, dentrda, dentrdz;
        double denerdd, denerdt, denerda, denerdz;

        // these could become functions:
        double chiT, chiRho, csound, grad_ad;  // derived quantities
        double gamma1, gamma2, gamma3, cV, cP; // derived quantities
        double dse, dpe, dsp;                  // Maxwell relations

        friend std::ostream& operator<<(std::ostream& os, const helmholtz::EOS& eos) {
            os << "EOS Data:\n" << std::setw(20) << std::left;
            os << "  Electron Fraction: " << std::format("{0:24.16e}",eos.ye) << "\n";
            os << "  Electron Chemical Potential: " << std::format("{0:24.16e}",eos.etaele) << "\n";
            os << "  Electron Number Density: " << std::format("{0:24.16e}",eos.xnefer) << "\n";
            os << "  Total Pressure: " << std::format("{0:24.16e}",eos.ptot) << "\n";
	    os << "  dPres/dRho: " << std::format("{0:24.16e}",eos.dpresdd) << "\n";
	    os << "  dPres/dT: " << std::format("{0:24.16e}",eos.dpresdt) << "\n";
	    os << "  Gas Pressure: " << std::format("{0:24.16e}",eos.pgas) << "\n";
            os << "  Radiation Pressure: " << std::format("{0:24.16e}",eos.prad) << "\n";
            os << "  Total Energy: " << std::format("{0:24.16e}",eos.etot) << "\n";
	    os << "  dEner/dRho: " << std::format("{0:24.16e}",eos.denerdd) << "\n";
	    os << "  dEner/dT: " << std::format("{0:24.16e}",eos.denerdt) << "\n";
            os << "  Gas Energy: " << std::format("{0:24.16e}",eos.egas) << "\n";
	    os << "  Radiation Energy: " << std::format("{0:24.16e}",eos.erad) << "\n";
	    os << "  Total Entropy: " << std::format("{0:24.16e}",eos.stot) << "\n";
	    os << "  dEntr/dRho: " << std::format("{0:24.16e}",eos.dentrdd) << "\n";
	    os << "  dEntr/dT: " << std::format("{0:24.16e}",eos.dentrdt) << "\n";
	    os << "  Gas Entropy: " << std::format("{0:24.16e}",eos.sgas) << "\n";
	    os << "  Radiation Entropy: " << std::format("{0:24.16e}",eos.srad);
            return os;
        }
    };

    // interpolating polynomial function definitions

    /**
     * @brief Interpolating polynomial function psi0.
     * @param z Input value.
     * @return Result of the polynomial.
     */
    double psi0(double z);

    /**
     * @brief Derivative of the interpolating polynomial function psi0.
     * @param z Input value.
     * @return Derivative of the polynomial.
     */
    double dpsi0(double z);

    /**
     * @brief Second derivative of the interpolating polynomial function psi0.
     * @param z Input value.
     * @return Second derivative of the polynomial.
     */
    double ddpsi0(double z);

    /**
     * @brief Interpolating polynomial function psi1.
     * @param z Input value.
     * @return Result of the polynomial.
     */
    double psi1(double z);

    /**
     * @brief Derivative of the interpolating polynomial function psi1.
     * @param z Input value.
     * @return Derivative of the polynomial.
     */
    double dpsi1(double z);

    /**
     * @brief Second derivative of the interpolating polynomial function psi1.
     * @param z Input value.
     * @return Second derivative of the polynomial.
     */
    double ddpsi1(double z);

    /**
     * @brief Interpolating polynomial function psi2.
     * @param z Input value.
     * @return Result of the polynomial.
     */
    double psi2(double z);

    /**
     * @brief Derivative of the interpolating polynomial function psi2.
     * @param z Input value.
     * @return Derivative of the polynomial.
     */
    double dpsi2(double z);

    /**
     * @brief Second derivative of the interpolating polynomial function psi2.
     * @param z Input value.
     * @return Second derivative of the polynomial.
     */
    double ddpsi2(double z);

    /**
     * @brief Interpolating polynomial function xpsi0.
     * @param z Input value.
     * @return Result of the polynomial.
     */
    double xpsi0(double z);

    /**
     * @brief Derivative of the interpolating polynomial function xpsi0.
     * @param z Input value.
     * @return Derivative of the polynomial.
     */
    double xdpsi0(double z);

    /**
     * @brief Interpolating polynomial function xpsi1.
     * @param z Input value.
     * @return Result of the polynomial.
     */
    double xpsi1(double z);

    /**
     * @brief Derivative of the interpolating polynomial function xpsi1.
     * @param z Input value.
     * @return Derivative of the polynomial.
     */
    double xdpsi1(double z);

    /**
     * @brief Interpolating polynomial function h3.
     * @param fi Array of coefficients.
     * @param w0t Weight 0 for temperature.
     * @param w1t Weight 1 for temperature.
     * @param w0mt Weight 0 for temperature (minus).
     * @param w1mt Weight 1 for temperature (minus).
     * @param w0d Weight 0 for density.
     * @param w1d Weight 1 for density.
     * @param w0md Weight 0 for density (minus).
     * @param w1md Weight 1 for density (minus).
     * @return Result of the polynomial.
     */
    double h3(double fi[36], double w0t, double w1t, double w0mt, double w1mt,
              double w0d, double w1d, double w0md, double w1md);

    /**
     * @brief Interpolating polynomial function h5.
     * @param fi Array of coefficients.
     * @param w0t Weight 0 for temperature.
     * @param w1t Weight 1 for temperature.
     * @param w2t Weight 2 for temperature.
     * @param w0mt Weight 0 for temperature (minus).
     * @param w1mt Weight 1 for temperature (minus).
     * @param w2mt Weight 2 for temperature (minus).
     * @param w0d Weight 0 for density.
     * @param w1d Weight 1 for density.
     * @param w2d Weight 2 for density.
     * @param w0md Weight 0 for density (minus).
     * @param w1md Weight 1 for density (minus).
     * @param w2md Weight 2 for density (minus).
     * @return Result of the polynomial.
     */
    double h5(double fi[36], double w0t, double w1t, double w2t, double w0mt,
              double w1mt, double w2mt, double w0d, double w1d, double w2d,
              double w0md, double w1md, double w2md);

    /**
     * @brief Read the Helmholtz EOS table from a file.
     * @param filename Path to the file containing the table.
     * @return HELMTable structure containing the table data.
     */
    HELMTable read_helm_table(const std::string filename);

    /**
     * @brief Calculate the Helmholtz EOS components.
     * @param q EOSInput parameters for the EOS calculation.
     * @param table HELMTable structure containing the table data.
     * @return EOS structure containing the calculated quantities.
     */
    EOS get_helm_EOS(EOSInput &q, const HELMTable &table);

}

#endif // HELM_H