/*******************************************************************************
 * This file is part of SWIFT.
 * Copyright (c) 2018 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_FEEDBACK_COLIBRE_H
#define SWIFT_FEEDBACK_COLIBRE_H

#include "SNIa_DTD.h"
#include "cooling.h"
#include "cosmology.h"
#include "dust.h"
#include "engine.h"
#include "error.h"
#include "feedback_properties.h"
#include "hydro_properties.h"
#include "part.h"
#include "random.h"
#include "rays.h"
#include "units.h"

void compute_stellar_evolution(const struct feedback_props* feedback_props,
                               const struct cosmology* cosmo,
                               const struct dustevo_props* dustevo_props,
                               struct spart* sp, const struct unit_system* us,
                               const double age, const double dt,
                               const double time_beg_of_step,
                               const integertime_t ti_begin);

/**
 * @brief Update the properties of a particle fue to feedback effects after
 * the cooling was applied.
 *
 * Nothing to do here in the COLIBRE model.
 *
 * @param p The #part to consider.
 * @param xp The #xpart to consider.
 * @param e The #engine.
 */
__attribute__((always_inline)) INLINE static void feedback_update_part(
    struct part* restrict p, struct xpart* restrict xp,
    const struct engine* restrict e) {}

/**
 * @brief Reset the gas particle-carried fields related to SNII kinetic feedback
 * at the start of a step.
 * @param p The particle.
 * @param xp The extended data of the particle.
 */
__attribute__((always_inline)) INLINE static void feedback_reset_part(
    struct part* p, struct xpart* xp) {

  /* Reset the id of the star that is allowed to kick this gas particle */
  p->feedback_data.SNII_stellar_id_kick_allowed = -1;
  /* Reset the maximal available kinetic energy per pair of two gas particles */
  p->feedback_data.SNII_max_E_kin_pair_avail = 0.f;
}

/**
 * @brief Should this particle be doing any feedback-related operation?
 *
 * @param sp The #spart.
 * @param e The #engine.
 */
__attribute__((always_inline)) INLINE static int feedback_is_active(
    const struct spart* sp, const struct engine* e) {

  return (sp->birth_time != -1.) &&
         (e->step <= 0 || (sp->feedback_data.do_feedback));
}

/**
 * @brief Prepares a s-particle for its feedback interactions
 *
 * @param sp The particle to act upon
 */
__attribute__((always_inline)) INLINE static void feedback_init_spart(
    struct spart* sp) {

  sp->feedback_data.to_collect.enrichment_weight_inv = 0.f;
  sp->feedback_data.to_collect.ngb_mass = 0.f;
  sp->feedback_data.to_collect.ngb_N = 0;
  sp->feedback_data.to_collect.ngb_rho = 0.f;
  sp->feedback_data.to_collect.ngb_u = 0.f;
  sp->feedback_data.to_collect.ngb_Z = 0.f;

  /* Reset all ray structs carried by this star particle */
  ray_init(sp->feedback_data.SNII_rays_true, colibre_SNII_feedback_num_of_rays);
  ray_init(sp->feedback_data.SNII_rays_mirr, colibre_SNII_feedback_num_of_rays);
  ray_init(sp->feedback_data.SNIa_rays, colibre_SNIa_feedback_num_of_rays);

  /* Reset all ray_extra structs carried by this star particle */
  ray_extra_init(sp->feedback_data.SNII_rays_ext_true,
                 colibre_SNII_feedback_num_of_rays);
  ray_extra_init(sp->feedback_data.SNII_rays_ext_mirr,
                 colibre_SNII_feedback_num_of_rays);
}

/**
 * @brief Returns the length of time since the particle last did
 * enrichment/feedback.
 *
 * @param sp The #spart.
 * @param with_cosmology Are we running with cosmological time integration on?
 * @param cosmo The cosmological model.
 * @param time The current time (since the Big Bang / start of the run) in
 * internal units.
 * @param dt_star the length of this particle's time-step in internal units.
 * @return The length of the enrichment step in internal units.
 */
INLINE static double feedback_get_enrichment_timestep(
    const struct spart* sp, const int with_cosmology,
    const struct cosmology* cosmo, const double time, const double dt_star) {

  if (with_cosmology) {
    return cosmology_get_delta_time_from_scale_factors(
        cosmo, (double)sp->last_enrichment_time, cosmo->a);
  } else {
    return time - (double)sp->last_enrichment_time;
  }
}

/**
 * @brief Prepares a star's feedback field before computing what
 * needs to be distributed.
 */
__attribute__((always_inline)) INLINE static void feedback_reset_feedback(
    struct spart* sp, const struct feedback_props* feedback_props) {

  /* Zero the distribution weights */
  sp->feedback_data.to_distribute.enrichment_weight = 0.f;

  /* Zero the amount of mass that is distributed */
  sp->feedback_data.to_distribute.mass = 0.f;

  sp->feedback_data.to_distribute.mass_from_NSM = 0.f;
  sp->feedback_data.to_distribute.mass_from_CEJSN = 0.f;
  sp->feedback_data.to_distribute.mass_from_collapsar = 0.f;

  /* Zero the metal enrichment quantities */
  for (int i = 0; i < chemistry_element_count; i++) {
    sp->feedback_data.to_distribute.metal_mass[i] = 0.f;
  }
  sp->feedback_data.to_distribute.total_metal_mass = 0.f;
  sp->feedback_data.to_distribute.mass_from_AGB = 0.f;
  sp->feedback_data.to_distribute.metal_mass_from_AGB = 0.f;
  sp->feedback_data.to_distribute.mass_from_SNII = 0.f;
  sp->feedback_data.to_distribute.metal_mass_from_SNII = 0.f;
  sp->feedback_data.to_distribute.mass_from_SNIa = 0.f;
  sp->feedback_data.to_distribute.metal_mass_from_SNIa = 0.f;
  sp->feedback_data.to_distribute.Fe_mass_from_SNIa = 0.f;

  /* Zero the dust enrichment quantities */
  for (int i = 0; i < dust_grain_species_count; i++) {
    sp->feedback_data.to_distribute.dust_mass[i] = 0.f;
  }

  /* Zero the energy to inject */
  sp->feedback_data.to_distribute.energy = 0.f;

  /* Zero the SNII/SNIa feedback properties */
  sp->feedback_data.to_distribute.SNII_delta_u = 0.f;
  sp->feedback_data.to_distribute.SNIa_delta_u = 0.f;
  sp->feedback_data.to_distribute.SNII_E_kinetic = 0.f;
  sp->feedback_data.to_distribute.SNII_num_of_thermal_energy_inj = 0;
  sp->feedback_data.to_distribute.SNII_num_of_kinetic_energy_inj = 0;
  sp->feedback_data.to_distribute.SNIa_num_of_thermal_energy_inj = 0;

  /* Reset the kick velocity and probability in the early stellar feedback */
  sp->feedback_data.to_distribute.momentum_probability = -1.f;
  sp->feedback_data.to_distribute.momentum_delta_v = 0.0;

  /* Reset the HII region probability */
  sp->feedback_data.to_distribute.HIIregion_probability = -1.f;
  sp->feedback_data.to_distribute.HIIregion_endtime = -1.f;
  sp->feedback_data.to_distribute.HIIregion_starid = -1;

  /* Reset the SNII star age */
  sp->feedback_data.to_distribute.SNII_star_age_Myr = -1.;
}

/**
 * @brief Initialises the s-particles feedback props for the first time
 *
 * This function is called only once just after the ICs have been
 * read in to do some conversions.
 *
 * @param sp The particle to act upon.
 * @param feedback_props The properties of the feedback model.
 */
__attribute__((always_inline)) INLINE static void feedback_first_init_spart(
    struct spart* sp, const struct feedback_props* feedback_props) {

  feedback_init_spart(sp);

  /* Zero residual SN kinetic energy and number of injection events */
  sp->feedback_data.SNII_E_kinetic_residual = 0.f;
  sp->feedback_data.SNII_num_of_kinetic_energy_inj_residual = 0;

  sp->feedback_data.do_feedback = 0;
}

/**
 * @brief Initialises the s-particles feedback props for the first time
 *
 * This function is called only once just after the ICs have been
 * read in to do some conversions.
 *
 * @param sp The particle to act upon.
 * @param feedback_props The properties of the feedback model.
 */
__attribute__((always_inline)) INLINE static void feedback_prepare_spart(
    struct spart* sp, const struct feedback_props* feedback_props) {}

/**
 * @brief Evolve the stellar properties of a #spart.
 *
 * This function allows for example to compute the SN rate before sending
 * this information to a different MPI rank.
 *
 * @param sp The particle to act upon
 * @param feedback_props The #feedback_props structure.
 * @param cosmo The current cosmological model.
 * @param us The unit system.
 * @param phys_const The physical constants in internal units.
 * @param star_age_beg_step The age of the star at the star of the time-step in
 * internal units.
 * @param dt The time-step size of this star in internal units or in the case of
 *           downsampling, the time since the last enrichment
 * @param ti_begin The integer time at the beginning of the step.
 * @param with_cosmology Are we running with cosmological time integration on?
 */
__attribute__((always_inline)) INLINE static void feedback_prepare_feedback(
    struct spart* restrict sp, const struct feedback_props* feedback_props,
    const struct cosmology* cosmo, const struct dustevo_props* dustevo_props,
    const struct unit_system* us, const struct phys_const* phys_const,
    const double star_age_beg_step, const double dt, const double time,
    const integertime_t ti_begin, const int with_cosmology) {

#ifdef SWIFT_DEBUG_CHECKS
  if (sp->birth_time == -1.) error("Evolving a star particle that should not!");
#endif

  /* Start by finishing the loops over neighbours */
  const float h = sp->h;
  const float h_inv = 1.0f / h;                 /* 1/h */
  const float h_inv_dim = pow_dimension(h_inv); /* 1/h^d */

  sp->feedback_data.to_collect.ngb_rho *= h_inv_dim;
  const float rho_inv = 1.f / sp->feedback_data.to_collect.ngb_rho;
  sp->feedback_data.to_collect.ngb_u *= h_inv_dim * rho_inv;
  sp->feedback_data.to_collect.ngb_Z *= h_inv_dim * rho_inv;

  /* Compute amount of enrichment and feedback that needs to be done in this
   * step */
  compute_stellar_evolution(feedback_props, cosmo, dustevo_props, sp, us,
                            star_age_beg_step, dt, time - dt, ti_begin);

  /* Decrease star mass by amount of mass distributed to gas neighbours */
  sp->mass -= sp->feedback_data.to_distribute.mass;

  /* Mark this is the last time we did enrichment */
  if (with_cosmology)
    sp->last_enrichment_time = cosmo->a;
  else
    sp->last_enrichment_time = time;

#ifdef SWIFT_DEBUG_CHECKS
  if (sp->mass < 0.)
    error("Stellar mass got negative! Check the yield and DTD normalisation!");
#endif
}

/**
 * @brief Will this star particle want to do feedback during the next time-step?
 *
 * @param sp The star of interest.
 * @param feedback_props The properties of the feedback model.
 * @param with_cosmology Are we running a cosmological problem?
 * @param cosmo The cosmological model.
 * @param time The current time (since the start of the run / Big Bang).
 */
__attribute__((always_inline)) INLINE static void feedback_will_do_feedback(
    struct spart* sp, const struct feedback_props* feedback_props,
    const int with_cosmology, const struct cosmology* cosmo, const double time,
    const struct unit_system* us, const struct phys_const* phys_const,
    const integertime_t ti_begin, const double time_base) {

  /* Special case for new-born stars */
  if (with_cosmology) {
    if (sp->birth_scale_factor == (float)cosmo->a) {

      /* Flag we want to do feedback */
      sp->feedback_data.do_feedback = 1;

      /* Ok, we are done. */
      return;
    }
  } else {
    if (sp->birth_time == (float)time) {

      /* Flag we want to do feedback */
      sp->feedback_data.do_feedback = 1;

      /* Ok, we are done. */
      return;
    }
  }

  /* Calculate age of the star at current time */
  double age_of_star;
  if (with_cosmology) {
    age_of_star = cosmology_get_delta_time_from_scale_factors(
        cosmo, (double)sp->birth_scale_factor, cosmo->a);
  } else {
    age_of_star = time - (double)sp->birth_time;
  }

  /* Is the star still young? */
  if (age_of_star < feedback_props->stellar_evolution_age_cut) {

    /* Flag we want to do feedback */
    sp->feedback_data.do_feedback = 1;

  } else {

    /* Calculate age of the star at current time */
    double delta_time_since_last_enrichment;

    if (with_cosmology) {
      delta_time_since_last_enrichment =
          cosmology_get_delta_time_from_scale_factors(
              cosmo, (double)sp->last_enrichment_time, cosmo->a);
    } else {
      delta_time_since_last_enrichment =
          time - (double)sp->last_enrichment_time;
    }

    /* Is it feedback o'clock? */
    if (delta_time_since_last_enrichment >
        feedback_props->stellar_evolution_sampling_age_fraction * age_of_star) {

      sp->feedback_data.do_feedback = 1;
    } else {
      sp->feedback_data.do_feedback = 0;
    }
  }
}

void feedback_clean(struct feedback_props* feedback_props);

void feedback_struct_dump(const struct feedback_props* feedback, FILE* stream);

void feedback_struct_restore(struct feedback_props* feedback, FILE* stream);

#ifdef HAVE_HDF5
/**
 * @brief Writes the current model of feedback to the file
 *
 * @param feedback The properties of the feedback scheme.
 * @param h_grp The HDF5 group in which to write.
 */
INLINE static void feedback_write_flavour(struct feedback_props* feedback,
                                          hid_t h_grp) {

  io_write_attribute_s(h_grp, "Feedback Model", "COLIBRE");
}
#endif  // HAVE_HDF5

#endif /* SWIFT_FEEDBACK_COLIBRE_H */