/*******************************************************************************
 * This file is part of SWIFT.
 * Copyright (c) 2019 Matthieu Schaller (schaller@strw.leidenuniv.nl)
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 ******************************************************************************/
#ifndef SWIFT_COLIBRE_COOLING_SUBGRID_H
#define SWIFT_COLIBRE_COOLING_SUBGRID_H

/* Config parameters. */
#include <config.h>

/* Local includes. */
#include "colibre_tables_restrict.h"
#include "cooling.h"
#include "equation_of_state.h"

/**
 * @brief Compute the subgrid properties of the gas.
 * For particles on the entropy floor, we use pressure equilibrium to
 * infer the properties of the particle.
 * Note that we return the density in physical coordinates.
 *
 * @return depends on the input for isub:
 *   isub = subgrid_density: the physical subgrid density in internal units
 *   isub = subgrid_temperature: the subgrid temperature temperature
 *   isub = subgrid_HI_fraction: the subgrid HI fraction (n_HI / n_H)
 *   isub = subgrid_HII_fraction: the subgrid HII fraction (n_HII / n_H)
 *   isub = subgrid_H2_fraction: the subgrid HII fraction (n_H2 / n_H)
 *
 * @param cooling The properties of the cooling scheme.
 * @param phys_const The physical constants.
 * @param floor_props The properties of the entropy floor.
 * @param cosmo The cosmological model.
 * @param rho_phys Density of the gas in internal physical units.
 * @param logZZsol Logarithm base 10 of the gas' metallicity in units of solar
 * metallicity.
 * @param XH The Hydrogen abundance of the gas.
 * @param P_phys Pressure of the gas in internal physical units.
 * @param log10_T The logarithm base 10 of the temperature of the gas.
 * @param log10_u_EOS_max_cgs The logarithm base 10 of the maximal energy in cgs
 * to be considered on the EOS at the density of the gas.
 * @param HII_region Is this patch of gas in an HII region?
 * @param abundance_ratio element abundance ratio of gas particle
 * @param log_u_cgs The logarithm base 10 of the internal energy in cgs
 * @param isub which subgrid property to calculate
 */
double compute_subgrid_property(
    const struct cooling_function_data *cooling,
    const struct phys_const *phys_const,
    const struct entropy_floor_properties *floor_props,
    const struct cosmology *cosmo, const float rho_phys, const float logZZsol,
    const float XH, const float P_phys, const float log10_T,
    const float log10_u_EOS_max_cgs, const int HII_region,
    const float abundance_ratio[colibre_cooling_N_elementtypes],
    const double log_u_cgs, const enum colibre_subgrid_properties isub) {

  /* Set the value for returning the values based on the hydro properties,
   * e.g. for particles above the EOS, or if no EOS is present */

  double standard_return;

  /* Start by computing the subgrid property as if it were not subgrid */

  if (isub == colibre_compute_subgrid_density) {
    standard_return = rho_phys;
  } else if (isub == colibre_compute_subgrid_temperature) {
    standard_return = exp10(log10_T);
  } else {
    /* This should never happen! */
    standard_return = -1.;
    error("Unknown subgrid property to calculate");
  }

  /* Return case 1: */
  /* Particle is above the equation of state, nothing to be done */
  if (log_u_cgs >= log10_u_EOS_max_cgs) {
    return standard_return;
  }

  /* Return case 2: */
  /* Particle is in an HII region */
  /* get density by mapping particle onto T_HII and mu_HII at equal pressure */
  if (HII_region) {
    if (isub == colibre_compute_subgrid_density) {
      const double mu_HII = get_mu_HII_region(cooling, XH);
      const double P_cgs =
          (double)P_phys * cooling->colibre_table.pressure_to_cgs;

      const double rho_cgs = P_cgs * cooling->colibre_table.proton_mass_cgs /
                             (exp10(cooling->colibre_table.log10_kB_cgs) *
                              mu_HII * XH * cooling->HIIregion_temp);

      return rho_cgs * cooling->colibre_table.density_from_cgs;

    } else if (isub == colibre_compute_subgrid_temperature) {
      return cooling->HIIregion_temp;
    } else {
      error("Unknown subgrid property to calculate");
      return 0.;
    }
  }

  /* Get the Hydrogen number density */
  const double n_H = XH * rho_phys / phys_const->const_proton_mass;
  const double n_H_cgs = n_H * cooling->colibre_table.number_density_to_cgs;

  /* Get index along the different table axis for this particle */
  int imet, iden, item;
  float dmet, dden, dtem;

  cooling_get_index_1d(cooling->colibre_table.Metallicity,
                       colibre_cooling_N_metallicity, logZZsol, &imet, &dmet);
  cooling_get_index_1d(cooling->colibre_table.nH, colibre_cooling_N_density,
                       log10f(n_H_cgs), &iden, &dden);
  cooling_get_index_1d(cooling->colibre_table.Temp,
                       colibre_cooling_N_temperature, log10_T, &item, &dtem);

  float log10_Teq =
      interpolation_3d(cooling->colibre_table.logTeq,         /* */
                       0 /*ired*/, imet, iden,                /* */
                       cooling->colibre_table.dz, dmet, dden, /* */
                       colibre_cooling_N_loaded_redshifts,    /* */
                       colibre_cooling_N_metallicity,         /* */
                       colibre_cooling_N_density);

  /* Return case 3:
   * The tabulated thermal equilibrium temperature for the SPH property
   * is larger than the SPH/EOS gas temperature. This is possible for low
   * densities and high metallicities, but also if non-equilibrium effects
   * increase the cooling rate compared to the equilibrium cooling rate.
   * The subgrid properties cannot be found by projecting the SPH properties
   * on the thermal equilibrium function and we return the SPH properties */
  if (log10_Teq >= log10_T) {
    return standard_return;
  }

  /* Particle is on the EOS and the SPH/EOS temperature is above the thermal
   * equilibrium temperature.
   * The subgrid property is calculated by projecting the hydro properties
   * on the thermal equilibrium temperature where the pressure equals
   * the hydro pressure */

  const float P_cgs = P_phys * cooling->colibre_table.pressure_to_cgs;
  const float log10_P_cgs = log10f(P_cgs);
  /* Maximal equilibrium pressure (n = n_max, and Teq = Teq(n_max)) from the
   * table at the current redshift and metallicity */
  const float log10_Peq_max_cgs = interpolation_3d_no_z(
      cooling->colibre_table.logPeq,                   /* */
      0 /*ired*/, imet, colibre_cooling_N_density - 1, /* */
      cooling->colibre_table.dz, dmet, 0.,             /* */
      colibre_cooling_N_loaded_redshifts,              /* */
      colibre_cooling_N_metallicity,                   /* */
      colibre_cooling_N_density);

  /* Return case 4:
   * EOS pressure (logP) is larger than maximum Peq (possible for very
   * steep EOS); use the equilibrium temperature from the highest density bin */
  if (log10_P_cgs > log10_Peq_max_cgs) {

    if (isub == colibre_compute_subgrid_density) {
      const double rho_cgs =
          exp10(cooling->colibre_table.nH[colibre_cooling_N_density - 1]) *
          cooling->colibre_table.proton_mass_cgs / XH;

      return rho_cgs * cooling->colibre_table.density_from_cgs;

    } else if (isub == colibre_compute_subgrid_temperature) {

      const float log10_T_at_Peq = interpolation_3d_no_z(
          cooling->colibre_table.logTeq, 0 /*ired*/, imet,
          colibre_cooling_N_density - 1, cooling->colibre_table.dz, dmet, 0.,
          colibre_cooling_N_loaded_redshifts, colibre_cooling_N_metallicity,
          colibre_cooling_N_density);

      return exp10(log10_T_at_Peq);

    } else {
      error("Unknown subgrid property to calculate");
      return 0.;
    }
  }

  /* Find equilibrium property by increasing the density
   * Note that the logPeq table is neither equally spaced nor
   * necessarilly monotically increasing.
   * We hence loop over densities and pick the first one where
   * log10_P < log10_Peq. We start with the next resolved density (index
   * iden+1), as we require the subgrid density to be larger than the physical
   * density */

  float log10_Peq_prev = log10_P_cgs;

  for (int i = iden + 1; i < colibre_cooling_N_density; i++) {

    float log10_n_at_Peq;
    float log10_Peq_current = interpolation_3d_no_z(
        cooling->colibre_table.logPeq, 0 /*ired*/, imet, i,
        cooling->colibre_table.dz, dmet, 0., colibre_cooling_N_loaded_redshifts,
        colibre_cooling_N_metallicity, colibre_cooling_N_density);

    if (log10_Peq_current >= log10_P_cgs) {

      if (log10_Peq_current == log10_Peq_prev) {
        /* the interpolated pressure exactly equals the input pressure
         * on the first density iteration (where log10_Peq_prev = log10_P_cgs),
         * no interpolation necessary */

        log10_n_at_Peq = cooling->colibre_table.nH[i];
      } else {
        /* first density bin where log P_eq >= log P_SPH
         * the solution is therefore between log10_Peq_prev and
         * log10_Peq_current */

        const float delta_P_eq = (log10_P_cgs - log10_Peq_prev) /
                                 (log10_Peq_current - log10_Peq_prev);

        /* Interpolate to get the density at equilibrium */
        log10_n_at_Peq = cooling->colibre_table.nH[i - 1] +
                         delta_P_eq * (cooling->colibre_table.nH[i] -
                                       cooling->colibre_table.nH[i - 1]);
      }
      /* Return case 5:
       * SPH density and temperatures are projected onto the thermal
       * equilibrium temperature function Teq(density, metallicity/abundances,
       * redshift) for equal pressure */

      if (isub == colibre_compute_subgrid_density) {

        const double rho_cgs =
            exp10(log10_n_at_Peq) * cooling->colibre_table.proton_mass_cgs / XH;

        /* only project to higher densities */
        if (rho_cgs * cooling->colibre_table.density_from_cgs < rho_phys) {
          return standard_return;
        } else {
          return rho_cgs * cooling->colibre_table.density_from_cgs;
        }

      } else if (isub == colibre_compute_subgrid_temperature) {

        /* We found a valid density: Get the index along the density
         * axis of the tables. */
        int iden_eq;
        float dden_eq;

        cooling_get_index_1d(cooling->colibre_table.nH,
                             colibre_cooling_N_density, log10_n_at_Peq,
                             &iden_eq, &dden_eq);

        /* Finish by interpolating the tables to get the temperature */
        const float log10_T_at_Peq =
            interpolation_3d(cooling->colibre_table.logTeq,            /* */
                             0 /*ired*/, imet, iden_eq,                /* */
                             cooling->colibre_table.dz, dmet, dden_eq, /* */
                             colibre_cooling_N_loaded_redshifts,       /* */
                             colibre_cooling_N_metallicity,            /* */
                             colibre_cooling_N_density);

        /* only project to lower temperatures */
        if (log10_T_at_Peq > log10_T) {
          return standard_return;
        } else {
          return exp10(log10_T_at_Peq);
        }

      } else {

        error("Unknown subgrid property to calculate");
        return 0.;
      }
    }
    log10_Peq_prev = log10_Peq_current;
  }

  /* Return case 6:
   * Nothing worked, return the property at the SPH density and temperature
   * The code should not reach this point during the simulation, but still
   * does before and at the first timestep */
  return standard_return;
}

#endif /* SWIFT_COLIBRE_COOLING_SUBGRID_H */
