#ifndef SWIFT_DUST_IO_T20_H
#define SWIFT_DUST_IO_T20_H

#include "dust.h"
#include "engine.h"
#include "io_properties.h"

/**
 * @brief Conversion function for adding explicit depletion back to metal arrays
 *
 * @param e The #engine.
 * @param p The particle array.
 * @param xp The extended particle array.
 * @param ret Output array.
 **/
INLINE static void convert_part_add_back_dust_for_chemistry(
    const struct engine* e, const struct part* p, const struct xpart* xp,
    float* ret) {

  /* first, set to return internal (diffuse phase) elemental abundances */
  for (int i = 0; i < chemistry_element_count; i++) {
    ret[i] = (float)p->chemistry_data.metal_mass_fraction[i];
  }

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (dp->pair_to_cooling) {
    /* put fraction of grain mass back into constituent element
     * so output arrays represent total metal mass fraction
     * (depleted + diffuse) instead of diffuse arrays used
     * internally when running in paired mode. */
    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      const float grain_mass = p->dust_data.grain_mass_fraction[grain];
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        const int eldx = dp->grain_element_indices[grain][elem];
        ret[eldx] += grain_mass * dp->grain_element_mfrac[grain][elem];
      }
    }
  }
}

INLINE static void convert_part_add_back_dust_for_reduced_chemistry(
    const struct engine* e, const struct part* p, const struct xpart* xp,
    float* ret) {

  /* first, set to return internal (diffuse phase) elemental abundances */
  ret[0] = (float)p->chemistry_data.metal_mass_fraction[chemistry_element_H];
  ret[1] = (float)p->chemistry_data.metal_mass_fraction[chemistry_element_He];

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (dp->pair_to_cooling) {
    /* put fraction of grain mass back into constituent element
     * so output arrays represent total metal mass fraction
     * (depleted + diffuse) instead of diffuse arrays used
     * internally when running in paired mode. */
    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      const float grain_mass = p->dust_data.grain_mass_fraction[grain];
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        const int eldx = dp->grain_element_indices[grain][elem];

        if (eldx == chemistry_element_H)
          ret[0] += grain_mass * dp->grain_element_mfrac[grain][elem];
        if (eldx == chemistry_element_He)
          ret[1] += grain_mass * dp->grain_element_mfrac[grain][elem];
      }
    }
  }
}

/**
 * @brief Conversion function to create diffuse chemicals field
 *
 * @param e The #engine.
 * @param p The particle array.
 * @param xp The extended particle array.
 * @param ret Output array.
 */
INLINE static void convert_part_diffuse_element_mass_fractions(
    const struct engine* e, const struct part* p, const struct xpart* xp,
    float* ret) {

  /* first, set to return internal (diffuse phase) elemental abundances */
  for (int i = 0; i < chemistry_element_count; i++) {
    ret[i] = (float)p->chemistry_data.metal_mass_fraction[i];
  }

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (!dp->pair_to_cooling) {
    /* put fraction of grain mass back into constituent element
     * so output arrays represent total metal mass fraction
     * (depleted + diffuse) instead of diffuse arrays used
     * internally when running in paired mode. */
    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      const float grain_mass = p->dust_data.grain_mass_fraction[grain];
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        const int eldx = dp->grain_element_indices[grain][elem];
        ret[eldx] -= grain_mass * dp->grain_element_mfrac[grain][elem];
      }
    }
  }
}

/**
 * @brief Conversion function for adding back depleted SNIa iron
 *
 * @param e The #engine.
 * @param p The particle array.
 * @param xp The extended particle array.
 * @param ret Output array.
 */
INLINE static void convert_part_add_back_dust_for_snia_iron(
    const struct engine* e, const struct part* p, const struct xpart* xp,
    float* ret) {

  float snia_iron = p->chemistry_data.iron_mass_fraction_from_SNIa;

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (dp->pair_to_cooling) {
    snia_iron += p->dust_data.grain_iron_mass_fraction_from_SNIa;
  }

  *ret = snia_iron;
}

/**
 * @brief Conversion function to create diffuse SNIa iron field
 *
 * @param e The #engine.
 * @param p The particle array.
 * @param xp The extended particle array.
 * @param ret Output array.
 */
INLINE static void convert_part_diffuse_snia_iron(const struct engine* e,
                                                  const struct part* p,
                                                  const struct xpart* xp,
                                                  float* ret) {

  float snia_iron = p->chemistry_data.iron_mass_fraction_from_SNIa;

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (!dp->pair_to_cooling) {
    snia_iron -= p->dust_data.grain_iron_mass_fraction_from_SNIa;
  }

  *ret = snia_iron;
}

/**
 * @brief Conversion function for adding explicit depletion back to metal arrays
 *
 * @param e The #engine.
 * @param bp The particle array.
 * @param ret Output array.
 **/
INLINE static void convert_bpart_add_back_dust_for_chemistry(
    const struct engine* e, const struct bpart* bp, float* ret) {

  /* first, set to return internal (diffuse phase) elemental abundances */
  for (int i = 0; i < chemistry_element_count; i++) {
    ret[i] = (float)bp->chemistry_data.metal_mass[i];
  }

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (dp->pair_to_cooling) {
    /* put masses of grain mass back into constituent element
     * so output arrays represent total metal mass
     * (depleted + diffuse) instead of diffuse arrays used
     * internally when running in paired mode. */
    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      const float grain_mass = bp->dust_data.grain_mass[grain];
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        const int eldx = dp->grain_element_indices[grain][elem];
        ret[eldx] += grain_mass * dp->grain_element_mfrac[grain][elem];
      }
    }
  }
}

/**
 * @brief Conversion function to create diffuse chemicals field
 *
 *
 * @param e The #engine.
 * @param bp The particle array.
 * @param ret Output array.
 */
INLINE static void convert_bpart_diffuse_element_mass_fractions(
    const struct engine* e, const struct bpart* bp, float* ret) {

  /* first, set to return internal (diffuse phase) elemental abundances */
  for (int i = 0; i < chemistry_element_count; i++) {
    ret[i] = (float)bp->chemistry_data.metal_mass[i];
  }

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (!dp->pair_to_cooling) {
    /* put fraction of grain mass back into constituent element
     * so output arrays represent total metal mass fraction
     * (depleted + diffuse) instead of diffuse arrays used
     * internally when running in paired mode. */
    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      const float grain_mass = bp->dust_data.grain_mass[grain];
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        const int eldx = dp->grain_element_indices[grain][elem];
        ret[eldx] -= grain_mass * dp->grain_element_mfrac[grain][elem];
      }
    }
  }
}

/**
 * @brief Conversion function for adding back depleted SNIa iron
 *
 * @param e The #engine.
 * @param bp The particle array.
 * @param ret Output array.
 */
INLINE static void convert_bpart_add_back_dust_for_snia_iron(
    const struct engine* e, const struct bpart* bp, float* ret) {

  float snia_iron = bp->chemistry_data.iron_mass_from_SNIa;

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (dp->pair_to_cooling) {
    snia_iron += bp->dust_data.grain_iron_mass_from_SNIa;
  }

  *ret = snia_iron;
}

/**
 * @brief Conversion function to create diffuse SNIa iron field
 *
 *
 * @param e The #engine.
 * @param bp The particle array.
 * @param ret Output array.
 */
INLINE static void convert_bpart_diffuse_snia_iron(const struct engine* e,
                                                   const struct bpart* bp,
                                                   float* ret) {

  float snia_iron = bp->chemistry_data.iron_mass_from_SNIa;

  const struct dustevo_props* dp = e->dustevo;

  /* running in paired mode? */
  if (!dp->pair_to_cooling) {
    snia_iron -= bp->dust_data.grain_iron_mass_from_SNIa;
  }

  *ret = snia_iron;
}

/**
 * @brief Conversion function to compute the total dust mass
 *
 * @param e The #engine.
 * @param p The particle array.
 * @param xp The particle array (extra data).
 * @param ret Output array.
 */
INLINE static void convert_part_total_dust_mass_fractions(
    const struct engine* e, const struct part* p, const struct xpart* xp,
    float* ret) {

  *ret = 0;
  for (int i = 0; i < dust_grain_species_count; ++i)
    *ret += p->dust_data.grain_mass_fraction[i];
}

/**
 * @brief Specifies which particle fields to write to a dataset
 *
 * @param parts The particle array.

 * @param list The list of i/o properties to write.
 * @param with_cosmology flag indicating if running with cosmology.
 *
 * @return Returns the number of fields to write.
 */
INLINE static int dust_write_particles(const struct part* parts,
                                       const struct xpart* xparts,
                                       struct io_props* list,
                                       const int with_cosmology) {

  /* List what we want to write */
  list[0] = io_make_output_field("DustMassFractions", FLOAT,
                                 dust_grain_species_count, UNIT_CONV_NO_UNITS,
                                 0.f, parts, dust_data.grain_mass_fraction,
                                 "Fractions of the particles' masses that are "
                                 "in a given species of dust grain");

  list[1] = io_make_output_field_convert_part(
      "TotalDustMassFractions", FLOAT, 1, UNIT_CONV_NO_UNITS, 0.f, parts,
      xparts, convert_part_total_dust_mass_fractions,
      "Fractions of the particles' masses that are in dust (sum of all grains "
      "species)");

  return 2;
}

/**
 * @brief Specifies which black hole particle fields to write to a dataset
 *
 * @param bparts The black hole particle array.
 * @param list The list of i/o properties to write.
 *
 * @return Returns the number of fields to write.
 */
INLINE static int dust_write_bparticles(const struct bpart* bparts,
                                        struct io_props* list) {

  /* List what we want to write */
  list[0] = io_make_output_field(
      "DustMasses", FLOAT, dust_grain_species_count, UNIT_CONV_MASS, 0.f,
      bparts, dust_data.grain_mass,
      "Mass contents of the BH particles in a given dust grain species");

  return 1;
}

#ifdef HAVE_HDF5

/**
 * @brief Writes the current model of dust evolution to the file
 * @param h_grp The HDF5 group in which to write
 * @param h_grp_columns The HDF5 group containing named columns
 * @param e The #engine.
 */
INLINE static void dust_write_flavour(hid_t h_grp, hid_t h_grp_columns,
                                      const struct engine* e) {

  io_write_attribute_s(h_grp, "Dust Model", "Trayford et. al (2024)");

  /* Create an array of grain names */
  const int grain_name_length = 32;
  char grain_names[dust_grain_species_count][grain_name_length];
  for (int grain = 0; grain < dust_grain_species_count; ++grain) {
    sprintf(grain_names[grain], "%s",
            dust_get_grain_name((enum grain_species)grain));
  }

  /* Add to the named columns */
  hsize_t dims[1] = {dust_grain_species_count};
  hid_t type = H5Tcopy(H5T_C_S1);
  H5Tset_size(type, grain_name_length);
  hid_t space = H5Screate_simple(1, dims, NULL);
  hid_t dset = H5Dcreate(h_grp_columns, "DustMassFractions", type, space,
                         H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
  H5Dwrite(dset, type, H5S_ALL, H5S_ALL, H5P_DEFAULT, grain_names[0]);
  H5Dclose(dset);
  /* dset = H5Dcreate(h_grp_columns, "MetalDiffusionRates", type, space, */
  /*                  H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT); */
  /* H5Dwrite(dset, type, H5S_ALL, H5S_ALL, H5P_DEFAULT, element_names[0]); */
  /* H5Dclose(dset); */

  H5Tclose(type);
  H5Sclose(space);
}

/**
 * @brief Create and write array for mapping elements to dust grains
 * @param h_grp The HDF5 group in which to write
 * @param e The #engine.
 */
INLINE static void dust_write_composition(hid_t h_grp, const struct engine* e) {

  const struct dustevo_props* dp = e->dustevo;

  /* Create an array mapping grains to elements */
  float grain_mapping[dust_grain_species_count][chemistry_element_count];
  bzero(grain_mapping,
        dust_grain_species_count * chemistry_element_count * sizeof(float));

  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];
      grain_mapping[grain][eldx] = dp->grain_element_mfrac[grain][elem];
    }
  }

  /* Write out array */
  hsize_t dims[2] = {dust_grain_species_count, chemistry_element_count};
  hid_t space = H5Screate_simple(2, dims, NULL);
  hid_t dset = H5Dcreate(h_grp, "GrainToElementMapping", H5T_NATIVE_FLOAT,
                         space, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
  if (dset < 0) error("Error while creating grain to element map array");

  H5Dwrite(dset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
           &grain_mapping[0][0]);
  io_write_attribute_s(dset, "Description",
                       "Mass fraction of respective grain type (row)"
                       "constituted by given element (column)");
  H5Dclose(dset);
  H5Sclose(space);
}

#endif

#endif /* SWIFT_DUST_T20_PROPERTIES_H */