/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Authors: Aaron Dotter, Emily Boudreaux
//   Last Modified: March 20, 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

#pragma once


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

#include "debug.h"

namespace serif::eos::helmholtz {
    static constexpr int IMAX = 541;
    static constexpr int JMAX = 201;
    /**
     * @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);

    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(const std::array<double, 36> &fi, const double w0t, const double w1t, const double w0mt, const double w1mt,
      const double w0d, const double w1d, const double w0md, const 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(const std::array<double, 36> &fi, const double w0t, const double w1t, const double w2t, const double w0mt,
          const double w1mt, const double w2mt, const double w0d, const double w1d, const double w2d,
          const double w0md, const double w1md, const 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.
     */
    std::unique_ptr<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);

}

