/*******************************************************************************
 * This file is part of SWIFT.
 * Copyright (c) 2018 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
 *
 * 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_TRACERS_COLIBRE_H
#define SWIFT_TRACERS_COLIBRE_H

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

/* Local includes */
#include "black_holes.h"
#include "cooling.h"
#include "part.h"
#include "star_formation.h"
#include "tracers_struct.h"

/**
 * @brief Update the particle tracers just after it has been initialised at the
 * start of a step.
 *
 * Nothing to do here in the COLIBRE model.
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param with_cosmology Are we running a cosmological simulation?
 * @param cosmo The current cosmological model.
 * @param hydro_props the hydro_props struct
 * @param cooling The #cooling_function_data used in the run.
 * @param time The current time.
 */
static INLINE void tracers_after_init(
    const struct part *p, struct xpart *xp, const struct unit_system *us,
    const struct phys_const *phys_const, const int with_cosmology,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct cooling_function_data *cooling, const double time) {}

/**
 * @brief Update the particle tracers just after it has been drifted.
 *
 * Nothing to do here in the COLIBRE model.
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param with_cosmology Are we running a cosmological simulation?
 * @param cosmo The current cosmological model.
 * @param hydro_props the hydro_props struct
 * @param cooling The #cooling_function_data used in the run.
 * @param time The current time.
 */
static INLINE void tracers_after_drift(
    const struct part *p, struct xpart *xp, const struct unit_system *us,
    const struct phys_const *phys_const, const int with_cosmology,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct cooling_function_data *cooling, const double time) {}

/**
 * @brief Update the particle tracers just after its time-step has been
 * computed.
 *
 * In COLIBRE we record the highest temperature reached.
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param with_cosmology Are we running a cosmological simulation?
 * @param cosmo The current cosmological model.
 * @param hydro_props the hydro_props struct
 * @param cooling The #cooling_function_data used in the run.
 * @param time The current time.
 */
static INLINE void tracers_after_timestep_part(
    const struct part *p, struct xpart *xp, const struct unit_system *us,
    const struct phys_const *phys_const, const int with_cosmology,
    const struct cosmology *cosmo, const struct hydro_props *hydro_props,
    const struct cooling_function_data *cooling,
    const struct entropy_floor_properties *floor_props, const double time,
    const double time_step_length, const int *const tracers_triggers_started) {

  /* Current temperature */
  const float temperature = cooling_get_temperature(phys_const, hydro_props, us,
                                                    cosmo, cooling, p, xp);

  /* New record? */
  if (temperature > xp->tracers_data.maximum_temperature) {

    /* Throw an error if there is a particle(s) with a too high T */
    const double T_extreme_cgs = 1e12;
    const double T_extreme = T_extreme_cgs / us->UnitTemperature_in_cgs;

    if (temperature > T_extreme)
      error("High temperature! [T=%.5e, pj->id=%lld, pj->h=%.5e, pj->rho=%.5e]",
            temperature, p->id, p->h, p->rho);

    xp->tracers_data.maximum_temperature = temperature;

    if (with_cosmology) {
      xp->tracers_data.maximum_temperature_scale_factor = cosmo->a;
    } else {
      xp->tracers_data.maximum_temperature_time = time;
    }
  }

  /* New h record? */
  if (p->h < xp->tracers_data.min_h) {

    xp->tracers_data.min_h = p->h;

    if (with_cosmology) {
      xp->tracers_data.min_h_scale_factor = cosmo->a;
    } else {
      xp->tracers_data.min_h_time = time;
    }
  }

  /* Check for last time in ISM
   * Gas must be above a set over-density AND
   * (be an HII region OR have a atomic + molecular fraction < 0.5) */
  if (p->rho > 100.f * cosmo->mean_density_Omega_m) {

    const float fHI = cooling_get_particle_subgrid_HI_fraction(
        us, phys_const, cosmo, hydro_props, floor_props, cooling, p, xp);
    const float fH2 = cooling_get_particle_subgrid_H2_fraction(
        us, phys_const, cosmo, hydro_props, floor_props, cooling, p, xp);

    const int is_HII_region = xp->tracers_data.HIIregion_timer_gas > 0.;

    if (fHI + 2.f * fH2 > 0.5f || is_HII_region) {
      if (with_cosmology) {
        xp->tracers_data.last_ISM_scale_factor = cosmo->a;
      } else {
        xp->tracers_data.last_ISM_time = time;
      }
    }
  }

  /* Accumulate average SFR */
  for (int i = 0; i < num_snapshot_triggers_part; ++i) {
    if (tracers_triggers_started[i])
      xp->tracers_data.averaged_SFR[i] +=
          star_formation_get_SFR(p, xp) * time_step_length;
  }

  /* Current particle volume, convert to physical frame */
  const float volume = (p->mass / p->rho) * cosmo->a * cosmo->a * cosmo->a;

  /* Rate at which particle is converting its kinetic energy to internal
   * energy, on account of shocks (artificial viscosity). Also convert to
   * total across particle (not per unit mass) and in physical frame. */
  const float shocking_rate = hydro_get_comoving_viscous_internal_energy_dt(p) *
                              p->mass * cosmo->a2_inv;

  /* CMB energy density (in physical frame and in internal units) */
  const float CMB_energy_density =
      (4.f / phys_const->const_speed_light_c) *
      phys_const->const_stefan_boltzmann * phys_const->const_T_CMB_0 *
      phys_const->const_T_CMB_0 * phys_const->const_T_CMB_0 *
      phys_const->const_T_CMB_0;

  /* Slope of the distribution of number of electrons versus Lorentz factor.
   * This is assuming that shocks where they are injected are relativistic
   * and strong (Mach number > 10) */
  const float electron_slope = 2.2f;

  /* Adiabatic indeces of the relativistic electrons and (tangled) magnetic
   * fields */
  const float electron_adiabatic_index = 4.f / 3.f;
  const float magnetic_adiabatic_index = 4.f / 3.f;

  /* Combination of above numbers that will appear often */
  const float xi = (1.f - electron_slope) * (electron_adiabatic_index - 1.f);

  /* Increment all quantities we need for radio modeling. These are only
   * incremented if 'last_jet_kick_initial_volume' is > 0 because this will
   * only be true if this particle has been kicked before. The other condition
   * is that the shocking rate is larger than 0; otherwise no injection
   * occurs. */
  const float initial_volume =
      xp->tracers_data.jet_radio_emission.last_jet_kick_initial_volume;

  if (initial_volume > 0.f && shocking_rate > 0.f) {

    /* Ratio of volumes the particle has now and had when it was kicked. */
    const float volume_ratio = volume / initial_volume;

    /* Various integrals over shocking rate (modulo volume ratio factors raised
     * to different powers) */
    xp->tracers_data.jet_radio_emission.shocking_integrals[0] +=
        shocking_rate * powf(volume_ratio, -xi) * time_step_length;
    xp->tracers_data.jet_radio_emission.shocking_integrals[1] +=
        shocking_rate * powf(volume_ratio, -(magnetic_adiabatic_index - 1.f)) *
        time_step_length;
    xp->tracers_data.jet_radio_emission.shocking_integrals[2] +=
        shocking_rate *
        powf(volume_ratio, -xi - (electron_adiabatic_index - 1.f)) *
        time_step_length;

    /* Integral used to compute when most of the shocking ocurred */
    xp->tracers_data.jet_radio_emission.weighted_injection_scale_factor +=
        shocking_rate * cosmo->a *
        powf(volume_ratio, -xi - (electron_adiabatic_index - 1.f)) *
        time_step_length;

    /* Integrals used to compute the minimum Lorentz factors gamma of the
       electron distribution */
    xp->tracers_data.jet_radio_emission.gamma_integrals[0] +=
        (xp->tracers_data.jet_radio_emission.shocking_integrals[1] / volume) *
        time_step_length;
    xp->tracers_data.jet_radio_emission.gamma_integrals[1] +=
        powf(volume_ratio, (electron_adiabatic_index - 1.f)) *
        CMB_energy_density * time_step_length;

    /* Integrals used to compute the effective minimum Lorentz factors gamma
       of the electron distribution */
    xp->tracers_data.jet_radio_emission.gamma_effective_integrals[0] +=
        shocking_rate * xp->tracers_data.jet_radio_emission.gamma_integrals[0] *
        powf(volume_ratio, -xi) * time_step_length;
    xp->tracers_data.jet_radio_emission.gamma_effective_integrals[1] +=
        shocking_rate * xp->tracers_data.jet_radio_emission.gamma_integrals[1] *
        powf(volume_ratio, -xi) * time_step_length;
  }
}

/**
 * @brief Update the star particle tracers just after its time-step has been
 * computed.
 *
 * In EAGLE, nothing to do.
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param with_cosmology Are we running a cosmological simulation?
 * @param cosmo The current cosmological model.
 * @param time_step_length The length of the step that just finished
 * @param tracers_triggers_started Which triggers have started? (array of size
 * num_snapshot_triggers_spart)
 */
static INLINE void tracers_after_timestep_spart(
    struct spart *sp, const struct unit_system *us,
    const struct phys_const *phys_const, const int with_cosmology,
    const struct cosmology *cosmo, const double time_step_length,
    const int *const tracers_triggers_started) {}

/**
 * @brief Update the black hole particle tracers just after its time-step has
 * been computed.
 *
 * In EAGLE, we record the average accr. rate.
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param with_cosmology Are we running a cosmological simulation?
 * @param cosmo The current cosmological model.
 * @param time_step_length The length of the step that just finished
 * @param tracers_triggers_started Which triggers have started? (array of size
 * num_snapshot_triggers_bpart)
 */
static INLINE void tracers_after_timestep_bpart(
    struct bpart *bp, const struct unit_system *us,
    const struct phys_const *phys_const, const int with_cosmology,
    const struct cosmology *cosmo, const double time_step_length,
    const int *const tracers_triggers_started) {

  const float accr_rate = black_holes_get_accretion_rate(bp);

  /* Accumulate average accretion rate */
  for (int i = 0; i < num_snapshot_triggers_part; ++i) {
    if (tracers_triggers_started[i])
      bp->tracers_data.averaged_accretion_rate[i] +=
          accr_rate * time_step_length;
  }

  if (bp->time_bin < bp->tracers_data.min_time_bin) {

    bp->tracers_data.min_time_bin = bp->time_bin;

    if (with_cosmology) {
      bp->tracers_data.min_time_bin_scale_factor = cosmo->a;
    } else {
      bp->tracers_data.min_time_bin_time = cosmo->time;
    }
  }
}

/**
 * @brief Update the particle tracers just after its time-step has been
 * computed.
 *
 * Set the maximal temperature to a valid initial state
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param cosmo The current cosmological model.
 * @param hydro_props the hydro_props struct
 * @param cooling The #cooling_function_data used in the run.
 */
static INLINE void tracers_first_init_xpart(
    const struct part *p, struct xpart *xp, const struct unit_system *us,
    const struct phys_const *phys_const, const struct cosmology *cosmo,
    const struct hydro_props *hydro_props,
    const struct cooling_function_data *cooling) {

  xp->tracers_data.maximum_temperature = -1.f;
  xp->tracers_data.maximum_temperature_scale_factor = -1.f;
  xp->tracers_data.stellar_wind_momentum_received = 0.f;
  xp->tracers_data.HIIregion_timer_gas = -1.f;
  xp->tracers_data.HIIregion_starid = -1;
  xp->tracers_data.last_stellar_wind_kick_scale_factor = -1.f;
  xp->tracers_data.last_SNII_thermal_injection_scale_factor = -1.f;
  xp->tracers_data.last_SNII_kick_scale_factor = -1.f;
  xp->tracers_data.last_SNIa_thermal_injection_scale_factor = -1.f;
  xp->tracers_data.last_AGN_injection_scale_factor = -1.f;
  xp->tracers_data.last_SNII_kick_velocity = -1.f;
  xp->tracers_data.max_SNII_kick_velocity = -1.f;
  xp->tracers_data.density_at_last_SN_thermal_feedback_event = -1.f;
  xp->tracers_data.density_at_last_AGN_feedback_event = -1.f;
  xp->tracers_data.AGN_feedback_energy = 0.f;
  xp->tracers_data.last_AGN_feedback_energy = 0.f;
  xp->tracers_data.hit_by_jet_feedback = 0;
  xp->tracers_data.jet_feedback_energy = 0.f;
  xp->tracers_data.last_AGN_jet_feedback_scale_factor = 0.f;
  xp->tracers_data.last_AGN_jet_feedback_time = 0.f;
  xp->tracers_data.last_jet_kick_velocity = 0.f;
  xp->tracers_data.last_jet_kick_accretion_mode = BH_thick_disc;
  xp->tracers_data.last_jet_kick_BH_id = 0;
  xp->tracers_data.jet_radio_emission.last_jet_kick_initial_volume = 0.f;
  for (int i = 0; i < 3; ++i)
    xp->tracers_data.jet_radio_emission.shocking_integrals[i] = 0.f;
  for (int i = 0; i < 2; ++i)
    xp->tracers_data.jet_radio_emission.gamma_integrals[i] = 0.f;
  for (int i = 0; i < 2; ++i)
    xp->tracers_data.jet_radio_emission.gamma_effective_integrals[i] = 0.f;
  xp->tracers_data.jet_radio_emission.weighted_injection_scale_factor = 0.f;
  xp->tracers_data.delta_T_at_last_SN_thermal_feedback_event = -1.f;
  xp->tracers_data.delta_T_at_last_AGN_feedback_event = -1.f;
  xp->tracers_data.last_ISM_scale_factor = -1.f;
  xp->tracers_data.min_h = FLT_MAX;
  xp->tracers_data.min_h_scale_factor = -1.f;
  xp->tracers_data.last_halo_mass = -1.f;
  xp->tracers_data.last_in_halo_scale_factor = -1.f;
}

/**
 * @brief Initialise the star tracer data at the start of a calculation.
 *
 * In EAGLE, nothing to do.
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param cosmo The current cosmological model.
 */
static INLINE void tracers_first_init_spart(struct spart *sp,
                                            const struct unit_system *us,
                                            const struct phys_const *phys_const,
                                            const struct cosmology *cosmo) {}

/**
 * @brief Initialise the black hole tracer data at the start of a calculation.
 *
 * @param p Pointer to the particle data.
 * @param xp Pointer to the extended particle data (containing the tracers
 * struct).
 * @param us The internal system of units.
 * @param phys_const The physical constants in internal units.
 * @param cosmo The current cosmological model.
 */
static INLINE void tracers_first_init_bpart(struct bpart *bp,
                                            const struct unit_system *us,
                                            const struct phys_const *phys_const,
                                            const struct cosmology *cosmo) {
  for (int i = 0; i < num_snapshot_triggers_bpart; ++i)
    bp->tracers_data.averaged_accretion_rate[i] = 0.f;

  bp->tracers_data.min_time_bin = num_time_bins;
  bp->tracers_data.min_time_bin_scale_factor = -1.f;
}

static INLINE void tracers_after_SNII_thermal_feedback(
    const struct part *p, struct xpart *xp, const int with_cosmology,
    const float scale_factor, const double time, const float delta_T) {

  if (with_cosmology)
    xp->tracers_data.last_SNII_thermal_injection_scale_factor = scale_factor;
  else
    xp->tracers_data.last_SNII_thermal_injection_time = time;

  xp->tracers_data.density_at_last_SN_thermal_feedback_event =
      hydro_get_comoving_density(p) /
      (scale_factor * scale_factor * scale_factor);

  xp->tracers_data.delta_T_at_last_SN_thermal_feedback_event = delta_T;
}

static INLINE void tracers_after_SNII_kinetic_feedback(struct xpart *xp,
                                                       const int with_cosmology,
                                                       const float scale_factor,
                                                       const double time,
                                                       const float v_kick) {

  if (with_cosmology)
    xp->tracers_data.last_SNII_kick_scale_factor = scale_factor;
  else
    xp->tracers_data.last_SNII_kick_time = time;

  xp->tracers_data.last_SNII_kick_velocity = v_kick;

  if (v_kick > xp->tracers_data.max_SNII_kick_velocity) {
    xp->tracers_data.max_SNII_kick_velocity = v_kick;
  }
}

static INLINE void tracers_after_SNIa_thermal_feedback(
    const struct part *p, struct xpart *xp, const int with_cosmology,
    const float scale_factor, const double time, const float delta_T) {

  if (with_cosmology)
    xp->tracers_data.last_SNIa_thermal_injection_scale_factor = scale_factor;
  else
    xp->tracers_data.last_SNIa_thermal_injection_time = time;

  xp->tracers_data.density_at_last_SN_thermal_feedback_event =
      hydro_get_comoving_density(p) /
      (scale_factor * scale_factor * scale_factor);

  xp->tracers_data.delta_T_at_last_SN_thermal_feedback_event = delta_T;
}

static INLINE void tracers_after_black_holes_feedback(
    const struct part *p, struct xpart *xp, const int with_cosmology,
    const float scale_factor, const double time, const double delta_energy,
    const float delta_T) {

  if (with_cosmology)
    xp->tracers_data.last_AGN_injection_scale_factor = scale_factor;
  else
    xp->tracers_data.last_AGN_injection_time = time;

  xp->tracers_data.density_at_last_AGN_feedback_event =
      hydro_get_comoving_density(p) /
      (scale_factor * scale_factor * scale_factor);

  xp->tracers_data.AGN_feedback_energy += delta_energy;
  xp->tracers_data.last_AGN_feedback_energy = delta_energy;
  xp->tracers_data.delta_T_at_last_AGN_feedback_event = delta_T;
}

static INLINE void tracers_after_jet_feedback(
    const struct part *p, struct xpart *xp, const int with_cosmology,
    const float scale_factor, const double time, const double delta_energy,
    const double vel_kick, const enum BH_accretion_modes accretion_mode,
    const long long id) {

  if (with_cosmology)
    xp->tracers_data.last_AGN_jet_feedback_scale_factor = scale_factor;
  else
    xp->tracers_data.last_AGN_jet_feedback_time = time;
  xp->tracers_data.hit_by_jet_feedback++;
  xp->tracers_data.jet_feedback_energy += delta_energy;
  xp->tracers_data.last_jet_kick_velocity = vel_kick;
  xp->tracers_data.last_jet_kick_accretion_mode = accretion_mode;
  xp->tracers_data.last_jet_kick_BH_id = id;

  /* Volume of this particle at time of kick (physical frame) */
  xp->tracers_data.jet_radio_emission.last_jet_kick_initial_volume =
      (p->mass / p->rho) * scale_factor * scale_factor * scale_factor;
}

static INLINE void tracers_after_stellar_wind_feedback(struct xpart *xp,
                                                       const int with_cosmology,
                                                       const float scale_factor,
                                                       const double time) {
  if (with_cosmology)
    xp->tracers_data.last_stellar_wind_kick_scale_factor = scale_factor;
  else
    xp->tracers_data.last_stellar_wind_kick_time = time;
}

/**
 * @brief Tracer event called after a snapshot was written.
 *
 * @param p the #part.
 * @param xp the #xpart.
 */
static INLINE void tracers_after_snapshot_part(const struct part *p,
                                               struct xpart *xp) {

  for (int i = 0; i < num_snapshot_triggers_part; ++i)
    xp->tracers_data.averaged_SFR[i] = 0.f;
}

/**
 * @brief Tracer event called after a snapshot was written.
 *
 * @param sp the #spart.
 */
static INLINE void tracers_after_snapshot_spart(struct spart *sp) {

  for (int i = 0; i < num_snapshot_triggers_part; ++i)
    sp->tracers_data.averaged_SFR[i] = 0.f;
}

/**
 * @brief Tracer event called after a snapshot was written.
 *
 * @param bp the #bpart.
 */
static INLINE void tracers_after_snapshot_bpart(struct bpart *bp) {

  for (int i = 0; i < num_snapshot_triggers_bpart; ++i)
    bp->tracers_data.averaged_accretion_rate[i] = 0.f;
}

/**
 * @brief Tracer event called after a fof call.
 *
 * @param xp the #xpart.
 * @param fof_mass the mass of the FOF group this particle belongs to.
 * @param with_cosmology Are we running with comoving integration?
 * @param scale_factore The current scale-factor.
 * @param time The current time.
 */
static INLINE void tracers_after_fof_part(struct xpart *xp,
                                          const float fof_mass,
                                          const int with_cosmology,
                                          const float scale_factor,
                                          const double time) {

  xp->tracers_data.last_halo_mass = fof_mass;
  if (with_cosmology)
    xp->tracers_data.last_in_halo_scale_factor = scale_factor;
  else
    xp->tracers_data.last_in_halo_time = time;
}

/**
 * @brief Split the tracer content of a particle into n pieces
 *
 * @param p The #part.
 * @param xp The #xpart.
 * @param n The number of pieces to split into.
 */
__attribute__((always_inline)) INLINE static void tracers_split_part(
    struct part *p, struct xpart *xp, const double n) {

  xp->tracers_data.AGN_feedback_energy /= n;
  xp->tracers_data.stellar_wind_momentum_received /= n;
}

#endif /* SWIFT_TRACERS_COLIBRE_H */