#ifndef SWIFT_DUST_T20_H
#define SWIFT_DUST_T20_H

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

/* Local includes */
#include "dust_properties.h"

void dust_redistribute_masses_astration(struct spart* sp, const struct part* p,
                                        const struct dustevo_props* dp);

void dust_redistribute_masses_cooling_step(struct part* p,
                                           const struct dustevo_props* dp);

void dustevo_print_backend(const struct dustevo_props* dp);

float dustevo_new_depfac(float depletion_value, float reduction_factor);

void dust_recompute_composition(double* comp_array, const double mg_specfrac,
                                const double fe_specfrac,
                                const struct dustevo_props* dp);

void dust_evolve_part(const struct phys_const* phys_const,
                      const struct unit_system* us,
                      const struct cosmology* cosmo,
                      const struct hydro_props* hydro_properties,
                      const struct entropy_floor_properties* floor_props,
                      const struct cooling_function_data* cooling,
                      const struct dustevo_props* dp, struct part* restrict p,
                      struct xpart* restrict xp, const float dt,
                      const float dt_therm, const double time);

void dustevo_props_dump(const struct dustevo_props* dp, FILE* stream);

void dustevo_props_restore(struct dustevo_props* dp, struct feedback_props* fp,
                           FILE* stream);

void dustevo_set_jagged_grain_arrays(struct dustevo_props* dp);

void dustevo_set_jagged_cgrain_arrays(struct dustevo_props* dp);

void dustevo_check_init_abundance(struct dustevo_props* dp,
                                  struct swift_params* params);

/**
 * @brief Return a string containing the name of a given #grain_species.
 */
__attribute__((always_inline)) INLINE static const char* dust_get_grain_name(
    enum grain_species grain) {

  static const char* dust_grain_species_names[dust_grain_species_count] = {
      "GraphiteLarge", "MgSilicatesLarge", "FeSilicatesLarge",
      "GraphiteSmall", "MgSilicatesSmall", "FeSilicatesSmall"};

  return dust_grain_species_names[grain];
}

#ifdef HAVE_HDF5
void dust_read_depletion(hid_t id, float** log_depletion_fractions,
                         const int table_cooling_N_redshifts,
                         const int table_cooling_N_temperature,
                         const int table_cooling_N_metallicity,
                         const int table_cooling_N_density,
                         const int table_cooling_N_elementtypes);
#endif

void depletion_correct_rates_split(
    float* cooling_array_heating_rate, float* cooling_array_cooling_rate,
    float* log_depletion_fractions, const int table_cooling_N_redshifts,
    const int table_cooling_N_temperature,
    const int table_cooling_N_metallicity, const int table_cooling_N_density,
    const int table_cooling_N_elementtypes, const int table_cooling_N_cooltypes,
    const int table_cooling_N_heattypes);

void depletion_correct_rates(
    float* cooling_array_heating_rate, float* cooling_array_cooling_rate,
    float* log_depletion_fractions, const int table_cooling_N_redshifts,
    const int table_cooling_N_temperature,
    const int table_cooling_N_metallicity, const int table_cooling_N_density,
    const int table_cooling_N_elementtypes, const int table_cooling_N_cooltypes,
    const int table_cooling_N_heattypes);

/**
 * @brief Prepares a particle for the smooth dust calculation.
 *
 * Zeroes all the relevant arrays in preparation for the sums taking place in
 * the various smooth dust tasks
 *
 * @param p The particle to act upon
 * @param dp #chemistry_global_data containing chemistry informations.
 */
__attribute__((always_inline)) INLINE static void dust_init_part(
    struct part* restrict p, const struct dustevo_props* dp) {

  struct dust_part_data* cpd = &p->dust_data;

  for (int grain = 0; grain < dust_grain_species_count; ++grain) {
    cpd->dgrain_mass_fraction[grain] = 0.0f;
  }
  cpd->dgrain_iron_mass_fraction_from_SNIa = 0.0f;
}

__attribute__((always_inline)) INLINE static void dust_first_init_part(
    const struct phys_const* restrict phys_const,
    const struct unit_system* restrict us,
    const struct cosmology* restrict cosmo, const struct dustevo_props* dp,
    const struct chemistry_global_data* cd, struct part* restrict p,
    struct xpart* restrict xp) {

  // Add initialization of fields in dust_part_data struct.
  for (int grain = 0; grain < dust_grain_species_count; ++grain) {
    p->dust_data.grain_mass_fraction[grain] =
        dp->initial_grain_mass_fraction[grain];
  }

  // if running in paired mode, subtract away depleted part from element arrays
  if (dp->pair_to_cooling) {
    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        const int eldx = dp->grain_element_indices[grain][elem];
        p->chemistry_data.metal_mass_fraction[eldx] -=
            dp->initial_grain_mass_fraction[grain] *
            dp->grain_element_mfrac[grain][elem];
      }
    }
  }
  p->dust_data.grain_iron_mass_fraction_from_SNIa = 0.0f;
  dust_init_part(p, dp);
}

/**
 * @brief Initialise the dust properties of a black hole with
 * the dust properties of the gas it is born from.
 *
 * Black holes don't store fractions so we need to use element masses.
 *
 * @param bp_data The black hole data to initialise.
 * @param p_data The gas data to use.
 * @param gas_mass The mass of the gas particle.
 */
__attribute__((always_inline)) INLINE static void dust_bpart_from_part(
    struct dust_bpart_data* bp_data, const struct dust_part_data* p_data,
    const double gas_mass) {

  for (int i = 0; i < dust_grain_species_count; ++i) {
    bp_data->grain_mass[i] = p_data->grain_mass_fraction[i] * gas_mass;
  }
  bp_data->grain_iron_mass_from_SNIa =
      p_data->grain_iron_mass_fraction_from_SNIa * gas_mass;
}

/**
 * @brief Add the dust data of a gas particle to a black hole.
 *
 * Black holes don't store fractions so we need to add element masses.
 *
 * @param bp_data The black hole data to add to.
 * @param p_data The gas data to use.
 * @param gas_mass The mass of the gas particle.
 */
__attribute__((always_inline)) INLINE static void dust_add_part_to_bpart(
    struct dust_bpart_data* bp_data, const struct dust_part_data* p_data,
    const double gas_mass) {

  for (int i = 0; i < dust_grain_species_count; ++i) {
    bp_data->grain_mass[i] += p_data->grain_mass_fraction[i] * gas_mass;
  }
  bp_data->grain_iron_mass_from_SNIa +=
      p_data->grain_iron_mass_fraction_from_SNIa * gas_mass;
}

/**
 * @brief Transfer dust data of a gas particle to a black hole.
 *
 * In contrast to `dust_add_part_to_bpart`, only a fraction of the
 * masses stored in the gas particle are transferred here. Absolute masses
 * of the gas particle are adjusted as well.
 * Black holes don't store fractions so we need to add element masses.
 *
 * @param bp_data The black hole data to add to.
 * @param p_data The gas data to use.
 * @param nibble_mass The mass to be removed from the gas particle.
 * @param nibble_fraction The fraction of the (original) mass of the gas
 *        particle that is removed.
 */
__attribute__((always_inline)) INLINE static void dust_transfer_part_to_bpart(
    struct dust_bpart_data* bp_data, struct dust_part_data* p_data,
    const double nibble_mass, const double nibble_fraction) {

  for (int i = 0; i < dust_grain_species_count; ++i)
    bp_data->grain_mass[i] += p_data->grain_mass_fraction[i] * nibble_mass;
  bp_data->grain_iron_mass_from_SNIa +=
      p_data->grain_iron_mass_fraction_from_SNIa * nibble_mass;
}

/**
 * @brief Add the dust data of a black hole to another one.
 *
 * @param bp_data The black hole data to add to.
 * @param swallowed_data The black hole data to use.
 */
__attribute__((always_inline)) INLINE static void dust_add_bpart_to_bpart(
    struct dust_bpart_data* bp_data,
    const struct dust_bpart_data* swallowed_data) {

  for (int i = 0; i < dust_grain_species_count; ++i) {
    bp_data->grain_mass[i] += swallowed_data->grain_mass[i];
  }
  bp_data->grain_iron_mass_from_SNIa +=
      swallowed_data->grain_iron_mass_from_SNIa;
}

/**
 * @brief Returns the total dust mass carried by a particle.
 */
__attribute__((always_inline)) INLINE static float dust_get_total_mass(
    const struct part* p) {
  float mass = 0.f;
  for (int i = 0; i < dust_grain_species_count; ++i) {
    mass += p->dust_data.grain_mass_fraction[i] * p->mass;
  }

  return mass;
}

float dust_get_local_dust_to_gas_ratio(void);

double dust_get_large_to_small_ratio(void);

#define dust_check_for_NaNs(p) dust_check_for_NaNs_(p, __FUNCTION__, __LINE__);

__attribute__((always_inline)) INLINE static void dust_check_for_NaNs_(
    const struct part* p, const char* function, const int line) {

  for (int i = 0; i < dust_grain_species_count; ++i) {
    if (isnan(p->dust_data.grain_mass_fraction[i]))
      error("Entry i in the dust array is NaN! in function '%s', line %d",
            function, line);
    if (p->dust_data.grain_mass_fraction[i] < 0.)
      error("Entry i in the dust array is < 0.! in function '%s', line %d",
            function, line);
  }
}

#endif /* SWIFT_DUST_T20_PROPERTIES_H */