/* Local includes. */

#include "dust.h"

#include "chemistry.h"
#include "cooling.h"
#include "dust_constants.h"
#include "dust_properties.h"
#include "dust_yield_tables.h"

#include <string.h>

/**
 * @brief redistribute any dust mass back to element abundances
 * on star particle formation according to dust composition, to
 * represent astration
 *
 * @param p The gas particles.
 * @param dp Global dust parameters for initialisation.
 */
void dust_redistribute_masses_astration(struct spart *sp, const struct part *p,
                                        const struct dustevo_props *dp) {

  /**
   * iterate through grain species and element types and
   * redistribute dust abundances to element abundances,
   * according to composition array
   */

  float grain_mass;
  int eldx;

  for (int grain = 0; grain < dust_grain_species_count; grain++) {
    grain_mass = p->dust_data.grain_mass_fraction[grain];
    for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
      eldx = dp->grain_element_indices[grain][elem];
      /* put fraction of grain mass back into constituent element (astration) in
       * paired mode */
      if (dp->pair_to_cooling) {
        sp->chemistry_data.metal_mass_fraction[eldx] +=
            grain_mass * dp->grain_element_mfrac[grain][elem];
      }
    }
  }
}

/**
 * @brief Prints the dust evolution model to stdout.
 *
 * @param dust #dustevo_props struct.
 */
void dustevo_print_backend(const struct dustevo_props *dp) {
  message(
      "Running with an approximated McKinnon et al. (2016) dust evolution "
      "model.");
}

/**
 * @brief initialise structure housing global dust parametrisation.
 * In particular, flags and values set in the parameter file,
 * and any hard-coded properties
 *
 * @param dp Global dust parameters for initialisation
 * @param params The parsed parameter file.
 * @param phys_const The physical constants in internal units.
 * @param us The current internal system of units.
 */
void dustevo_props_init_backend(struct dustevo_props *dp,
                                struct swift_params *params,
                                struct feedback_props *fp,
                                struct cooling_function_data *cooling,
                                const struct phys_const *phys_const,
                                const struct unit_system *us) {

  /**
   * initialise structure housing global dust parametrisation.
   * First, set global parameters:
   *
   * initialise dustevo_props struct object and store:
   * - Parameter file flags (e.g. with_accretion, with_sputtering)
   * - Parameter file values (e.g diffusion boost, clumping factor)
   * - Hard coded values (e.g. grain radii, density, composition)
   *
   *
   * Then, using the yield tables from feedback_props, compute
   * and set dust yields for each of the  dust species
   *
   * Finally apply corections to the cooling and heating rates
   * to remove implicit dust, via the scale_out_table_depletion
   * function.
   **/

  /* read some parameters */
  /* operation modes */
  dp->pair_to_cooling =
      parser_get_opt_param_int(params, "DustEvolution:pair_to_cooling", 0);
  dp->with_subgrid_props =
      parser_get_opt_param_int(params, "DustEvolution:use_subgrid_props", 1);
  dp->with_sputtering =
      parser_get_opt_param_int(params, "DustEvolution:use_sputtering", 1);
  dp->with_SNII_destruction =
      parser_get_opt_param_int(params, "DustEvolution:use_SNII_destruction", 1);
  dp->with_accretion =
      parser_get_opt_param_int(params, "DustEvolution:use_accretion", 1);

  /* initial abundances */
  dp->initial_grain_mass_fraction[dust_grain_species_C] =
      parser_get_opt_param_float(
          params, "DustEvolution:inital_abundance_depletedC", 0.);
  dp->initial_grain_mass_fraction[dust_grain_species_O] =
      parser_get_opt_param_float(
          params, "DustEvolution:inital_abundance_depletedO", 0.);
  dp->initial_grain_mass_fraction[dust_grain_species_Mg] =
      parser_get_opt_param_float(
          params, "DustEvolution:inital_abundance_depletedMg", 0.);
  dp->initial_grain_mass_fraction[dust_grain_species_Si] =
      parser_get_opt_param_float(
          params, "DustEvolution:inital_abundance_depletedSi", 0.);
  dp->initial_grain_mass_fraction[dust_grain_species_Fe] =
      parser_get_opt_param_float(
          params, "DustEvolution:inital_abundance_depletedFe", 0.);

  /* global parameters */
  dp->clumping_factor =
      parser_get_opt_param_float(params, "DustEvolution:clumping_factor", 10.);
  dp->diffusion_rate_boost = parser_get_opt_param_float(
      params, "DustEvolution:diffusion_boost_factor", 1.0);

  /* initialise 1D arrays */
  memset(dp->abundance_pattern, 0, sizeof(dp->abundance_pattern));
  memset(dp->atomic_weight, 0, sizeof(dp->atomic_weight));
  memset(dp->condensation_frac, 0, sizeof(dp->condensation_frac));

  /* set assumed abundance patterns to Wiersma et al (2009a) */
  dp->abundance_pattern[chemistry_element_H] = 7.0649785e-01;
  dp->abundance_pattern[chemistry_element_He] = 2.8055534e-01;
  dp->abundance_pattern[chemistry_element_C] = 2.0665436e-03;
  dp->abundance_pattern[chemistry_element_N] = 8.3562563e-04;
  dp->abundance_pattern[chemistry_element_O] = 5.4926244e-03;
  dp->abundance_pattern[chemistry_element_Ne] = 1.4144605e-03;
  dp->abundance_pattern[chemistry_element_Mg] = 5.9070642e-04;
  dp->abundance_pattern[chemistry_element_Si] = 6.8258739e-04;
  dp->abundance_pattern[chemistry_element_Fe] = 1.1032152e-03;

  dp->solar_metallicity = 0.0127;

  /* set element atomic weights */
  dp->atomic_weight[chemistry_element_H] = 1.0079;
  dp->atomic_weight[chemistry_element_He] = 4.0026;
  dp->atomic_weight[chemistry_element_C] = 12.0107;
  dp->atomic_weight[chemistry_element_N] = 14.0067;
  dp->atomic_weight[chemistry_element_O] = 15.9994;
  dp->atomic_weight[chemistry_element_Ne] = 20.1797;
  dp->atomic_weight[chemistry_element_Mg] = 24.305;
  dp->atomic_weight[chemistry_element_Si] = 28.0855;
  dp->atomic_weight[chemistry_element_Fe] = 55.845;

  /* set condensation fractions */
  dp->condensation_frac[dust_grain_species_C] = direct_condensation_fraction_C;
  dp->condensation_frac[dust_grain_species_O] = direct_condensation_fraction_O;
  dp->condensation_frac[dust_grain_species_Mg] =
      direct_condensation_fraction_Mg;
  dp->condensation_frac[dust_grain_species_Si] =
      direct_condensation_fraction_Si;
  dp->condensation_frac[dust_grain_species_Fe] =
      direct_condensation_fraction_Fe;

  /* set element count contributing to each grain */
  dp->grain_element_count[dust_grain_species_C] = 1;
  dp->grain_element_count[dust_grain_species_O] = 1;
  dp->grain_element_count[dust_grain_species_Mg] = 1;
  dp->grain_element_count[dust_grain_species_Si] = 1;
  dp->grain_element_count[dust_grain_species_Fe] = 1;

  /* number of SNII per unit stellar mass */
  dp->specific_numSNII = snii_destruction_coeff / phys_const->const_solar_mass;

  /* set 2d grain arrays */
  dustevo_set_jagged_arrays(dp);

  /** NOTE 1: total metallicity yields untouched here, so Z represents the
   * conserved dust + gas phase metals **/

  /** NOTE 2: only the IMF resampled tables are modified in fp, while plain
   * .yield arrays are unchanged (the original yield tables are only used to
   * compute these and are already modified via the SNII yield factors) **/

  initialise_dyield_tables(fp, dp);
  compute_SNII_dyield(fp, dp);
  compute_AGB_dyield(fp, dp);
  // print_dyield_tables(fp, dp);
}

/**
 * @brief Write dustevo props to the given FILE as a stream of bytes.
 *
 * @param dp the struct
 * @param stream the file stream
 */
void dustevo_props_dump(const struct dustevo_props *dp, FILE *stream) {

  /* Zero all pointers in the dustevo_props struct. */
  struct dustevo_props dp_copy = *dp;

  dp_copy.logfD = NULL;

  dp_copy.dyield_AGB.yield_IMF_resampled = NULL;
  dp_copy.dyield_SNII.yield_IMF_resampled = NULL;

  /* iterate through pointer arrays and zero */
  for (int grain = 0; grain < dust_grain_species_count; grain++) {
    dp_copy.grain_element_indices[grain] = NULL;
    dp_copy.grain_element_mfrac[grain] = NULL;
  }

  restart_write_blocks((void *)&dp_copy, sizeof(struct dustevo_props), 1,
                       stream, "dp", "dustevo properties");
}

/**
 * @brief Restore a dustevo_props struct from the given FILE as a stream of
 * bytes.
 *
 * @param dp the struct
 * @param fp #feedback_props structure
 * @param stream the file stream
 */
void dustevo_props_restore(struct dustevo_props *dp, struct feedback_props *fp,
                           FILE *stream) {

  restart_read_blocks((void *)dp, sizeof(struct dustevo_props), 1, stream, NULL,
                      "dustevo properties");
  dustevo_set_jagged_arrays(dp);
  initialise_dyield_tables(fp, dp);
  compute_SNII_dyield(fp, dp);
  compute_AGB_dyield(fp, dp);
}

/**
 * @brief Set jagged arrays in the dustevo_props structure
 *
 * @param dp the struct
 */
void dustevo_set_jagged_arrays(struct dustevo_props *dp) {
  /* set grain composition */
  float gcomp1[1] = {1.};
  float gcomp2[1] = {1.};
  float gcomp3[1] = {1.};
  float gcomp4[1] = {1.};
  float gcomp5[1] = {1.};

  dp->grain_element_mfrac[dust_grain_species_C] = gcomp1;
  dp->grain_element_mfrac[dust_grain_species_O] = gcomp2;
  dp->grain_element_mfrac[dust_grain_species_Mg] = gcomp3;
  dp->grain_element_mfrac[dust_grain_species_Si] = gcomp4;
  dp->grain_element_mfrac[dust_grain_species_Fe] = gcomp5;

  /* allocate memory */
  dp->grain_element_mfrac[dust_grain_species_C] =
      malloc(dp->grain_element_count[dust_grain_species_C] * sizeof(int));
  dp->grain_element_mfrac[dust_grain_species_O] =
      malloc(dp->grain_element_count[dust_grain_species_O] * sizeof(int));
  dp->grain_element_mfrac[dust_grain_species_Mg] =
      malloc(dp->grain_element_count[dust_grain_species_Mg] * sizeof(int));
  dp->grain_element_mfrac[dust_grain_species_Si] =
      malloc(dp->grain_element_count[dust_grain_species_Si] * sizeof(int));
  dp->grain_element_mfrac[dust_grain_species_Fe] =
      malloc(dp->grain_element_count[dust_grain_species_Fe] * sizeof(int));

  /* deep copy grain arrays */
  memcpy(dp->grain_element_mfrac[dust_grain_species_C], &gcomp1,
         dp->grain_element_count[dust_grain_species_C] * sizeof(int));
  memcpy(dp->grain_element_mfrac[dust_grain_species_O], &gcomp2,
         dp->grain_element_count[dust_grain_species_O] * sizeof(int));
  memcpy(dp->grain_element_mfrac[dust_grain_species_Mg], &gcomp3,
         dp->grain_element_count[dust_grain_species_Mg] * sizeof(int));
  memcpy(dp->grain_element_mfrac[dust_grain_species_Si], &gcomp4,
         dp->grain_element_count[dust_grain_species_Si] * sizeof(int));
  memcpy(dp->grain_element_mfrac[dust_grain_species_Fe], &gcomp5,
         dp->grain_element_count[dust_grain_species_Fe] * sizeof(int));

  /* set chemistry indices composition */
  int gidx1[1] = {chemistry_element_C};
  int gidx2[1] = {chemistry_element_O};
  int gidx3[1] = {chemistry_element_Mg};
  int gidx4[1] = {chemistry_element_Si};
  int gidx5[1] = {chemistry_element_Fe};

  /* allocate memory */
  dp->grain_element_indices[dust_grain_species_C] =
      malloc(dp->grain_element_count[dust_grain_species_C] * sizeof(int));
  dp->grain_element_indices[dust_grain_species_O] =
      malloc(dp->grain_element_count[dust_grain_species_O] * sizeof(int));
  dp->grain_element_indices[dust_grain_species_Mg] =
      malloc(dp->grain_element_count[dust_grain_species_Mg] * sizeof(int));
  dp->grain_element_indices[dust_grain_species_Si] =
      malloc(dp->grain_element_count[dust_grain_species_Si] * sizeof(int));
  dp->grain_element_indices[dust_grain_species_Fe] =
      malloc(dp->grain_element_count[dust_grain_species_Fe] * sizeof(int));

  /* deep copy grain arrays */
  memcpy(dp->grain_element_indices[dust_grain_species_C], &gidx1,
         dp->grain_element_count[dust_grain_species_C] * sizeof(int));
  memcpy(dp->grain_element_indices[dust_grain_species_O], &gidx2,
         dp->grain_element_count[dust_grain_species_O] * sizeof(int));
  memcpy(dp->grain_element_indices[dust_grain_species_Mg], &gidx3,
         dp->grain_element_count[dust_grain_species_Mg] * sizeof(int));
  memcpy(dp->grain_element_indices[dust_grain_species_Si], &gidx4,
         dp->grain_element_count[dust_grain_species_Si] * sizeof(int));
  memcpy(dp->grain_element_indices[dust_grain_species_Fe], &gidx5,
         dp->grain_element_count[dust_grain_species_Fe] * sizeof(int));
}

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) {

  const float X = p->chemistry_data.metal_mass_fraction[chemistry_element_H];
  const float Z = p->chemistry_data.metal_mass_fraction_total;
  int eldx;

  /* metal mass fraction in dust before evolution step */
  float D_pre = 0.;
  for (int grain = 0; grain < dust_grain_species_count; grain++) {
    D_pre += p->dust_data.grain_mass_fraction[grain];
  }

  /* no dust, no accretion or destruction. return. */
  if (D_pre <= 0.) return;

  /* set temperature and density corresponding to either subgrid or hydro */
  float rho, temp;

  if (dp->with_subgrid_props) {
    rho = p->cooling_data.subgrid_dens;
    temp = p->cooling_data.subgrid_temp;
  } else {
    rho = hydro_get_physical_density(p, cosmo);
    temp = cooling_get_temperature(phys_const, hydro_properties, us, cosmo,
                                   cooling, p, xp);
  }

  /* --------------- Common Parameters --------------- */

  /* grain radius in cgs, hard coded for now, 0.1 micron. */
  const double grain_radius_cgs = representative_grain_size;

  /* convert Hydrogen mass fraction into Hydrogen number density */
  const double n_H = rho * X / phys_const->const_proton_mass;

  /* convert some properties to cgs  */
  const double n_H_cgs =
      n_H * units_cgs_conversion_factor(us, UNIT_CONV_NUMBER_DENSITY);
  const double dt_cgs = dt * units_cgs_conversion_factor(us, UNIT_CONV_TIME);

  /*------------------------- Sputtering -------------------------*/

  double deltf_sput;
  if (dp->with_sputtering) {
    /* compute fractional change in grain species mass this timestep */
    const double dfdr_cgs = 3. / grain_radius_cgs;
    const double drdt_grain_cgs =
        -sputtering_timescale_norm * n_H_cgs / (pow(2e6 / temp, 2.5) + 1);
    const double dfdt_cgs = drdt_grain_cgs * dfdr_cgs;
    deltf_sput = dfdt_cgs * dt_cgs;
  } else
    deltf_sput = 0.;

  /*------------------------- SNII shocks -------------------------*/

  double deltf_SNII;

  if (dp->with_SNII_destruction && (xp->sf_data.SFR > 0.)) {
    const float eps_eff = 0.1;
    const double unit_mass_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_MASS);
    const double unit_time_cgs =
        units_cgs_conversion_factor(us, UNIT_CONV_TIME);
    /* SNII rate assuming constant SFR, in cgs units */
    const float rate_SNII =
        xp->sf_data.SFR * dp->specific_numSNII / unit_time_cgs;

    /* Yamasawa et. al (2011): the cgs mass of material swept up in SNII shocks
     */
    const float m_swept = snii_destruction_coeff * pow(n_H_cgs, -0.202) *
                          pow((Z / 0.0127) + 0.039, -0.289);

    /* destruction timescale of dust in cgs */
    const float tau_dest =
        p->mass * unit_mass_cgs / (eps_eff * m_swept * rate_SNII);

    /* fractional change in dust due to dust destruction */
    deltf_SNII = -dt_cgs / tau_dest;
  } else {
    deltf_SNII = 0.;
  }

  /* ------------------------- Accretion ------------------------- */

  float deltf_acc[dust_grain_species_count] = {0.};
  if (dp->with_accretion) {
    /* accretion timescale multiplicative terms  (Hirashita & Voshchinnikov
       2013), implicitly assuming S_acc = 0.3 */
    const float grain_radius_term = grain_radius_cgs / 1e-5;
    /* const float Z_term  		= 0.0127 / Z; /\* Z_sun accessible
     * somewhere?
     * *\/
     */
    const float n_H_term = 1.e3 / n_H_cgs;
    const float T_term = sqrt(10. / temp);

    const float mult_terms = grain_radius_term * /*Z_term */ n_H_term * T_term;

    /* conversion factor converted to CGS from yr */
    const float tau_ref = 6.31e15 * mult_terms / dp->clumping_factor;

    /* Compute diffuse fraction of refractory elements for accretion */
    double diffuse_fraction;
    for (int grain = 0; grain < dust_grain_species_count; grain++) {
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        eldx = dp->grain_element_indices[grain][elem];
        if (p->chemistry_data.metal_mass_fraction[eldx] > 0) {
          /* NB: this works because for M16 model grains composed purely of
           * individual elements */
          if (dp->pair_to_cooling) {
            diffuse_fraction = (p->chemistry_data.metal_mass_fraction[eldx] /
                                (p->dust_data.grain_mass_fraction[grain] +
                                 p->chemistry_data.metal_mass_fraction[eldx]));
          } else {
            /* in paired mode, metal mass fractions include implicit dust */
            diffuse_fraction = ((p->chemistry_data.metal_mass_fraction[eldx] -
                                 p->dust_data.grain_mass_fraction[grain]) /
                                p->chemistry_data.metal_mass_fraction[eldx]);
          }

          deltf_acc[elem] = dt_cgs * diffuse_fraction / tau_ref;
        } else {
          deltf_acc[elem] = 0.;
        }
      }
    }
  }

  /* ------------------------- Combine ------------------------- */

  float delta_dfrac;
  float new_grain_mfrac;
  float D_post = 0.;

  for (int grain = 0; grain < dust_grain_species_count; grain++) {
    delta_dfrac = p->dust_data.grain_mass_fraction[grain] *
                  (exp(deltf_sput + deltf_SNII + deltf_acc[grain]) - 1.);
    /* limit grain growth */
    if (dp->pair_to_cooling) {
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        eldx = dp->grain_element_indices[grain][elem];
        /* grain growth limited by available mass in a consitituent element */
        delta_dfrac =
            min(delta_dfrac, p->chemistry_data.metal_mass_fraction[eldx] /
                                 dp->grain_element_mfrac[grain][elem]);
      }
    }

    else {
      float diffuse_mfrac;
      for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
        eldx = dp->grain_element_indices[grain][elem];
        /* grain growth limited by available mass in a consitituent element */
        diffuse_mfrac = p->chemistry_data.metal_mass_fraction[eldx] -
                        (dp->grain_element_mfrac[grain][elem] *
                         p->dust_data.grain_mass_fraction[grain]);

        delta_dfrac = min(delta_dfrac,
                          diffuse_mfrac / dp->grain_element_mfrac[grain][elem]);
      }
    }

    new_grain_mfrac = p->dust_data.grain_mass_fraction[grain] + delta_dfrac;

    /* update particle grain mass fractions */
    p->dust_data.grain_mass_fraction[grain] += delta_dfrac;

    if (new_grain_mfrac > FLT_MIN) {
      if (dp->pair_to_cooling) {
        for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
          eldx = dp->grain_element_indices[grain][elem];
          /* update constituent element abundances in paired mode */
          p->chemistry_data.metal_mass_fraction[eldx] -=
              delta_dfrac * dp->grain_element_mfrac[grain][elem];
        }
      }
      D_post += p->dust_data.grain_mass_fraction[grain];
    }

    else {
      p->dust_data.grain_mass_fraction[grain] = FLT_MIN;
      if (dp->pair_to_cooling) {
        for (int elem = 0; elem < dp->grain_element_count[grain]; elem++) {
          eldx = dp->grain_element_indices[grain][elem];
          /* update constituent element abundances in paired mode */
          p->chemistry_data.metal_mass_fraction[eldx] -=
              new_grain_mfrac * dp->grain_element_mfrac[grain][elem];
        }
      }
      D_post += FLT_MIN;
    }
  }
  if (0) {
    message("delta dust: %e, T %e rho %e", (D_post - D_pre) / D_pre, temp, rho);
  }
}