/*******************************************************************************
 * This file is part of SWIFT.
 * Copyright (c) 2017 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/>.
 *
 ******************************************************************************/

/**
 * @file src/cooling/CHIMES/colibre_tables.c
 * @brief Functions to read COLIBRE tables
 */

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

/* Standard includes */
#include <hdf5.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>

/* Local includes. */
#include "chemistry_struct.h"
#include "cooling_properties.h"
#include "dust.h"
#include "error.h"
#include "exp10.h"

/**
 * @brief Reads in COLIBRE table of redshift values
 *
 * @param cooling #cooling_function_data structure
 */
void get_cooling_redshifts(struct colibre_cooling_tables *table) {

  /* Read the list of table redshifts */
  char redshift_filename[colibre_table_path_name_length + 16];
  sprintf(redshift_filename, "%s/redshifts.dat", table->cooling_table_path);

  FILE *infile = fopen(redshift_filename, "r");
  if (infile == NULL) {
    error("Cannot open the list of cooling table redshifts (%s)",
          redshift_filename);
  }

  int N_Redshifts = -1;

  /* Read the file */
  if (!feof(infile)) {

    char buffer[50];

    /* Read the number of redshifts (1st line in the file) */
    if (fgets(buffer, 50, infile) != NULL)
      N_Redshifts = atoi(buffer);
    else
      error("Impossible to read the number of redshifts");

    /* Be verbose about it */
    if (engine_rank == 0)
      message("Found cooling tables at %d redshifts", N_Redshifts);

    table->number_of_redshifts = N_Redshifts;

    /* Allocate the list of redshifts */
    if (swift_memalign("cooling", (void **)&table->Redshifts,
                       SWIFT_STRUCT_ALIGNMENT,
                       N_Redshifts * sizeof(float)) != 0)
      error("Failed to allocate redshift table");

    /* Read all the redshift values */
    int count = 0;
    while (!feof(infile)) {
      if (fgets(buffer, 50, infile) != NULL) {
        table->Redshifts[count] = atof(buffer);
        count++;
      }
    }

    /* Verify that the file was self-consistent */
    if (count != N_Redshifts) {
      error(
          "Redshift file (%s) does not contain the correct number of redshifts "
          "(%d vs. %d)",
          redshift_filename, count, N_Redshifts);
    }
  } else {
    error("Redshift file (%s) is empty!", redshift_filename);
  }

  /* We are done with this file */
  fclose(infile);

  /* COLIBRE cooling assumes cooling->Redshifts table is in increasing order.
   * Test this. */
  for (int i = 0; i < N_Redshifts - 1; i++) {
    if (table->Redshifts[i + 1] < table->Redshifts[i]) {
      error("table should be in increasing order\n");
    }
  }
}

/**
 * @brief Reads in COLIBRE cooling table header. Consists of tables
 * of values for temperature, hydrogen number density, metallicity,
 * abundance ratios, and elements used to index the cooling tables.
 *
 * @param fname Filepath for cooling table from which to read header
 * @param table Colibre cooling table structure.
 */
void read_cooling_header(const char *fname,
                         struct colibre_cooling_tables *table) {

#ifdef HAVE_HDF5

  hid_t dataset;
  herr_t status;

  /* read sizes of array dimensions */
  hid_t tempfile_id = H5Fopen(fname, H5F_ACC_RDONLY, H5P_DEFAULT);
  if (tempfile_id < 0)
    error("unable to open file %s\n", table->cooling_table_path);

  /* allocate arrays of bins */
  if (posix_memalign((void **)&table->Temp, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_temperature * sizeof(float)) != 0)
    error("Failed to allocate temperature table\n");

  if (posix_memalign((void **)&table->nH, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_density * sizeof(float)) != 0)
    error("Failed to allocate density table\n");

  if (posix_memalign((void **)&table->Metallicity, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity * sizeof(float)) != 0)
    error("Failed to allocate metallicity table\n");

  if (posix_memalign((void **)&table->LogAbundances, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity *
                         colibre_cooling_N_elementtypes * sizeof(float)) != 0)
    error("Failed to allocate abundance array\n");

  if (posix_memalign((void **)&table->Abundances, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity *
                         colibre_cooling_N_elementtypes * sizeof(float)) != 0)
    error("Failed to allocate abundance array\n");

  if (posix_memalign((void **)&table->Abundances_inv, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity *
                         colibre_cooling_N_elementtypes * sizeof(float)) != 0)
    error("Failed to allocate abundance array\n");

  if (posix_memalign((void **)&table->atomicmass, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_elementtypes * sizeof(float)) != 0)
    error("Failed to allocate atomic masses array\n");

  if (posix_memalign((void **)&table->atomicmass_inv, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_elementtypes * sizeof(float)) != 0)
    error("Failed to allocate inverse atomic masses array\n");

  if (posix_memalign((void **)&table->Zsol, SWIFT_STRUCT_ALIGNMENT,
                     1 * sizeof(float)) != 0)
    error("Failed to allocate solar metallicity array\n");

  if (posix_memalign((void **)&table->Zsol_inv, SWIFT_STRUCT_ALIGNMENT,
                     1 * sizeof(float)) != 0)
    error("Failed to allocate inverse solar metallicity array\n");

  if (posix_memalign((void **)&table->LogMassFractions, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity *
                         colibre_cooling_N_elementtypes * sizeof(float)) != 0)
    error("Failed to allocate log mass fraction array\n");

  if (posix_memalign((void **)&table->MassFractions, SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity *
                         colibre_cooling_N_elementtypes * sizeof(float)) != 0)
    error("Failed to allocate mass fraction array\n");

  /* read in bins and misc information */
  dataset = H5Dopen(tempfile_id, "/TableBins/TemperatureBins", H5P_DEFAULT);
  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                   table->Temp);
  if (status < 0) error("error reading temperature bins\n");
  status = H5Dclose(dataset);
  if (status < 0) error("error closing cooling dataset");

  dataset = H5Dopen(tempfile_id, "/TableBins/DensityBins", H5P_DEFAULT);
  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                   table->nH);
  if (status < 0) error("error reading density bins\n");
  status = H5Dclose(dataset);
  if (status < 0) error("error closing cooling dataset");

  dataset = H5Dopen(tempfile_id, "/TableBins/MetallicityBins", H5P_DEFAULT);
  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                   table->Metallicity);
  if (status < 0) error("error reading metallicity bins\n");
  status = H5Dclose(dataset);
  if (status < 0) error("error closing cooling dataset");

  dataset = H5Dopen(tempfile_id, "/TotalAbundances", H5P_DEFAULT);
  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                   table->LogAbundances);
  if (status < 0) error("error reading total abundances\n");
  status = H5Dclose(dataset);
  if (status < 0) error("error closing cooling dataset");

  dataset = H5Dopen(tempfile_id, "/TotalMassFractions", H5P_DEFAULT);
  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                   table->LogMassFractions);
  if (status < 0) error("error reading total mass fractions\n");
  status = H5Dclose(dataset);
  if (status < 0) error("error closing cooling dataset");

  dataset = H5Dopen(tempfile_id, "/ElementMasses", H5P_DEFAULT);
  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                   table->atomicmass);
  if (status < 0) error("error reading element masses\n");
  status = H5Dclose(dataset);
  if (status < 0) error("error closing cooling dataset");

  dataset = H5Dopen(tempfile_id, "/SolarMetallicity", H5P_DEFAULT);
  status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                   table->Zsol);
  if (status < 0) error("error reading solar metallicity \n");
  status = H5Dclose(dataset);
  if (status < 0) error("error closing cooling dataset");

  /* Close the file */
  H5Fclose(tempfile_id);

  table->Zsol_inv[0] = 1.f / table->Zsol[0];

  /* find the metallicity bin that refers to solar metallicity */
  const float tol = 1.e-3;
  for (int i = 0; i < colibre_cooling_N_metallicity; i++) {
    if (fabsf(table->Metallicity[i]) < tol) {
      table->indxZsol = i;
    }
  }

#if defined(__ICC)
#pragma novector
#endif
  for (int i = 0; i < colibre_cooling_N_elementtypes; i++) {
    table->atomicmass_inv[i] = 1.f / table->atomicmass[i];
  }

  /* set some additional useful abundance arrays */
  for (int i = 0; i < colibre_cooling_N_metallicity; i++) {

#if defined(__ICC)
#pragma novector
#endif
    for (int j = 0; j < colibre_cooling_N_elementtypes; j++) {
      const int indx1d = cooling_row_major_index_2d(
          i, j, colibre_cooling_N_metallicity, colibre_cooling_N_elementtypes);
      table->Abundances[indx1d] = exp10f(table->LogAbundances[indx1d]);
      table->Abundances_inv[indx1d] = 1.f / table->Abundances[indx1d];
      table->MassFractions[indx1d] = exp10f(table->LogMassFractions[indx1d]);
    }
  }

#else
  error("Need HDF5 to read cooling tables");
#endif
}

/**
 * @brief Allocate space for cooling tables
 *
 * @param table Colibre cooling table structure.
 */
void allocate_cooling_tables(struct colibre_cooling_tables *table) {

  /* Allocate arrays to store cooling tables. Arrays contain two tables of
   * cooling rates with one table being for the redshift above current redshift
   * and one below. */

  /* Cooling reduced (temperature) */
  if (swift_memalign(
          "cooling_table.Tcooling_reduced", (void **)&table->Tcooling_reduced,
          SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_loaded_redshifts * colibre_cooling_N_temperature *
              colibre_cooling_N_metallicity * colibre_cooling_N_density *
              colibre_reduced_cooling_N_cooltypes * sizeof(float)) != 0)
    error("Failed to allocate Tcooling_reduced array\n");

  /* Heating reduced (temperature) */
  if (swift_memalign(
          "cooling_table.Theating_reduced", (void **)&table->Theating_reduced,
          SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_loaded_redshifts * colibre_cooling_N_temperature *
              colibre_cooling_N_metallicity * colibre_cooling_N_density *
              colibre_reduced_cooling_N_heattypes * sizeof(float)) != 0)
    error("Failed to allocate Theating_reduced array\n");

  /* Electron fraction (temperature) */
  if (swift_memalign(
          "cooling_table.Tefrac", (void **)&table->Telectron_fraction,
          SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_loaded_redshifts * colibre_cooling_N_temperature *
              colibre_cooling_N_metallicity * colibre_cooling_N_density *
              colibre_cooling_N_electrontypes * sizeof(float)) != 0)
    error("Failed to allocate Telectron_fraction array\n");

  /* Thermal equilibrium temperature */
  if (swift_memalign(
          "cooling_table.Teq", (void **)&table->logTeq, SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_loaded_redshifts * colibre_cooling_N_metallicity *
              colibre_cooling_N_density * sizeof(float)) != 0)
    error("Failed to allocate logTeq array\n");

  /* Mean particle mass at thermal equilibrium temperature */
  if (swift_memalign("cooling_table.mueq", (void **)&table->meanpartmass_Teq,
                     SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_loaded_redshifts *
                         colibre_cooling_N_metallicity *
                         colibre_cooling_N_density * sizeof(float)) != 0)
    error("Failed to allocate mu array\n");

  /* Pressure at thermal equilibrium temperature */
  if (swift_memalign(
          "cooling_table.Peq", (void **)&table->logPeq, SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_loaded_redshifts * colibre_cooling_N_metallicity *
              colibre_cooling_N_density * sizeof(float)) != 0)
    error("Failed to allocate logPeq array\n");

#ifdef SWIFT_DEBUG_CHECKS
  message("Allocated all cooling arrays for tables");
#endif
}

/**
 * @brief Get redshift dependent table of cooling rates.
 * Reads in cooling tables (Ploeckinger & Schaye 2020) split in redshift
 *
 * @param cooling #cooling_function_data structure
 * @param low_z_index Index of the lowest redshift table to load.
 * @param high_z_index Index of the highest redshift table to load.
 */
void get_cooling_table(struct colibre_cooling_tables *table,
                       const int low_z_index, const int high_z_index,
                       float *log_depletion_fractions,
                       const int do_dust_pairing) {

#ifdef HAVE_HDF5

  hid_t dataset;
  herr_t status;

  /* Temporary tables */
  float *cooling_T = NULL;
  float *heating_T = NULL;
  float *electron_fraction_T = NULL;
  float *equilibrium_temperature = NULL;
  float *mean_particle_mass_Teq = NULL;

  /* Allocate temporary tables */
  /* Cooling (temperature) */
  if (swift_memalign(
          "cooling_temp", (void **)&cooling_T, SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_temperature * colibre_cooling_N_metallicity *
              colibre_cooling_N_density * colibre_cooling_N_cooltypes *
              sizeof(float)) != 0)
    error("Failed to allocate temporary Tcooling array\n");

  /* Heating (temperature) */
  if (swift_memalign(
          "cooling_temp", (void **)&heating_T, SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_temperature * colibre_cooling_N_metallicity *
              colibre_cooling_N_density * colibre_cooling_N_heattypes *
              sizeof(float)) != 0)
    error("Failed to allocate temporary Theating array\n");

  /* Cooling, all channels, only loaded redshifts. */
  if (swift_memalign(
          "cooling_table.Tcooling", (void **)&table->Tcooling,
          SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_loaded_redshifts * colibre_cooling_N_temperature *
              colibre_cooling_N_metallicity * colibre_cooling_N_density *
              colibre_cooling_N_cooltypes * sizeof(float)) != 0)
    error("Failed to allocate Tcooling array\n");

  /* Heating, all channels, only loaded redshifts. */
  if (swift_memalign(
          "cooling_table.Theating", (void **)&table->Theating,
          SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_loaded_redshifts * colibre_cooling_N_temperature *
              colibre_cooling_N_metallicity * colibre_cooling_N_density *
              colibre_cooling_N_heattypes * sizeof(float)) != 0)
    error("Failed to allocate Theating array\n");

  /* Electron fraction (temperature) */
  if (swift_memalign(
          "cooling_temp", (void **)&electron_fraction_T, SWIFT_STRUCT_ALIGNMENT,
          colibre_cooling_N_temperature * colibre_cooling_N_metallicity *
              colibre_cooling_N_density * colibre_cooling_N_electrontypes *
              sizeof(float)) != 0)
    error("Failed to allocate temporary Telectron_fraction array\n");

  /* Thermal equilibrium temperature */
  if (swift_memalign("cooling_temp", (void **)&equilibrium_temperature,
                     SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity * colibre_cooling_N_density *
                         sizeof(float)) != 0)
    error("Failed to allocate temporary logTeq array\n");

  /* Mean particle mass at thermal equilibrium temperature */
  if (swift_memalign("cooling_temp", (void **)&mean_particle_mass_Teq,
                     SWIFT_STRUCT_ALIGNMENT,
                     colibre_cooling_N_metallicity * colibre_cooling_N_density *
                         sizeof(float)) != 0)
    error("Failed to allocate temporary mu array\n");

  /* Read in tables for redshift above and redshift below current value.  */
  for (int z_index = low_z_index; z_index <= high_z_index; z_index++) {

    /* Index along redhsift dimension for the subset of tables we read */
    const int local_z_index = z_index - low_z_index;

#ifdef SWIFT_DEBUG_CHECKS
    if (local_z_index >= colibre_cooling_N_loaded_redshifts)
      error("Reading invalid number of tables along z axis.");
#endif

    /* Open table for this redshift index */
    char fname[colibre_table_path_name_length + colibre_table_filebase_length +
               12];

    sprintf(fname, "%s/%s_z%4.2f.hdf5", table->cooling_table_path,
            table->cooling_table_filebase, table->Redshifts[z_index]);

    if (engine_rank == 0) message("Reading cooling table %s", fname);

    hid_t file_id = H5Fopen(fname, H5F_ACC_RDONLY, H5P_DEFAULT);
    if (file_id < 0) error("unable to open file %s", fname);

    /* Read Tdep/Cooling into cooling->table.Tcooling */
    dataset = H5Dopen(file_id, "/Tdep/Cooling", H5P_DEFAULT);
    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                     cooling_T);
    if (status < 0) error("error reading Tcooling\n");
    status = H5Dclose(dataset);
    if (status < 0) error("error closing cooling dataset");

    for (int i = 0; i < colibre_cooling_N_temperature; i++) {
      for (int j = 0; j < colibre_cooling_N_metallicity; j++) {
        for (int k = 0; k < colibre_cooling_N_density; k++) {
          for (int l = 0; l < colibre_cooling_N_cooltypes; l++) {

            /* Index in the HDF5 table */
            const int hdf5_index = cooling_row_major_index_4d(
                i, j, k, l, colibre_cooling_N_temperature,
                colibre_cooling_N_metallicity, colibre_cooling_N_density,
                colibre_cooling_N_cooltypes);

            /* Index in the internal table */
            const int internal_index = cooling_row_major_index_5d(
                local_z_index, i, j, k, l, colibre_cooling_N_loaded_redshifts,
                colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
                colibre_cooling_N_density, colibre_cooling_N_cooltypes);

            table->Tcooling[internal_index] = cooling_T[hdf5_index];
          }
        }
      }
    }

    /* Read Tdep/Heating into cooling->table.Theating */
    dataset = H5Dopen(file_id, "/Tdep/Heating", H5P_DEFAULT);
    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                     heating_T);
    if (status < 0) error("error reading Theating\n");
    status = H5Dclose(dataset);
    if (status < 0) error("error closing cooling dataset");

    for (int i = 0; i < colibre_cooling_N_temperature; i++) {
      for (int j = 0; j < colibre_cooling_N_metallicity; j++) {
        for (int k = 0; k < colibre_cooling_N_density; k++) {
          for (int l = 0; l < colibre_cooling_N_heattypes; l++) {

            /* Index in the HDF5 table */
            const int hdf5_index = cooling_row_major_index_4d(
                i, j, k, l, colibre_cooling_N_temperature,
                colibre_cooling_N_metallicity, colibre_cooling_N_density,
                colibre_cooling_N_heattypes);

            /* Index in the internal table */
            const int internal_index = cooling_row_major_index_5d(
                local_z_index, i, j, k, l, colibre_cooling_N_loaded_redshifts,
                colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
                colibre_cooling_N_density, colibre_cooling_N_heattypes);

            table->Theating[internal_index] = heating_T[hdf5_index];
          }
        }
      }
    }

    /* Read Tdep/ElectronFractions into cooling->table.Telectron_fraction */

    /* Dataset is named /Tdep/ElectronFractions in the published version of the
     * tables and for historical reasons /Tdep/ElectronFractionsVol in the
     * version used in the COLIBRE repository. Content is identical but we deal
     * here with both names */
    if (H5Lexists(file_id, "/Tdep/ElectronFractionsVol", H5P_DEFAULT) > 0) {
      dataset = H5Dopen(file_id, "/Tdep/ElectronFractionsVol", H5P_DEFAULT);
    } else if (H5Lexists(file_id, "/Tdep/ElectronFractions", H5P_DEFAULT) > 0) {
      dataset = H5Dopen(file_id, "/Tdep/ElectronFractions", H5P_DEFAULT);
    } else {
      error("Could not find the electron_fraction (temperature)!");
    }

    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                     electron_fraction_T);
    if (status < 0) error("error reading electron_fraction (temperature)\n");
    status = H5Dclose(dataset);
    if (status < 0) error("error closing cooling dataset");

    for (int i = 0; i < colibre_cooling_N_temperature; i++) {
      for (int j = 0; j < colibre_cooling_N_metallicity; j++) {
        for (int k = 0; k < colibre_cooling_N_density; k++) {
          for (int l = 0; l < colibre_cooling_N_electrontypes; l++) {

            /* Index in the HDF5 table */
            const int hdf5_index = cooling_row_major_index_4d(
                i, j, k, l, colibre_cooling_N_temperature,
                colibre_cooling_N_metallicity, colibre_cooling_N_density,
                colibre_cooling_N_electrontypes);

            /* Index in the internal table */
            const int internal_index = cooling_row_major_index_5d(
                local_z_index, i, j, k, l, colibre_cooling_N_loaded_redshifts,
                colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
                colibre_cooling_N_density, colibre_cooling_N_electrontypes);

            table->Telectron_fraction[internal_index] =
                electron_fraction_T[hdf5_index];
          }
        }
      }
    }

    /* Read ThermEq/Temperature into cooling->table.logTeq */
    dataset = H5Dopen(file_id, "/ThermEq/Temperature", H5P_DEFAULT);
    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                     equilibrium_temperature);
    if (status < 0) error("error reading Teq array\n");
    status = H5Dclose(dataset);
    if (status < 0) error("error closing logTeq dataset");

    for (int j = 0; j < colibre_cooling_N_metallicity; j++) {
      for (int k = 0; k < colibre_cooling_N_density; k++) {

        /* Index in the HDF5 table */
        const int hdf5_index = cooling_row_major_index_2d(
            j, k, colibre_cooling_N_metallicity, colibre_cooling_N_density);

        /* Index in the internal table */
        const int internal_index = cooling_row_major_index_3d(
            local_z_index, j, k, colibre_cooling_N_loaded_redshifts,
            colibre_cooling_N_metallicity, colibre_cooling_N_density);

        table->logTeq[internal_index] = equilibrium_temperature[hdf5_index];
      }
    }

    /* Read ThermEq/MeanParticleMass into cooling->table.meanpartmass_Teq */
    dataset = H5Dopen(file_id, "/ThermEq/MeanParticleMass", H5P_DEFAULT);
    status = H5Dread(dataset, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                     mean_particle_mass_Teq);
    if (status < 0) error("error reading mu array\n");
    status = H5Dclose(dataset);
    if (status < 0) error("error closing mu dataset");

    for (int j = 0; j < colibre_cooling_N_metallicity; j++) {
      for (int k = 0; k < colibre_cooling_N_density; k++) {

        /* Index in the HDF5 table */
        const int hdf5_index = cooling_row_major_index_2d(
            j, k, colibre_cooling_N_metallicity, colibre_cooling_N_density);

        /* Index in the internal table */
        const int internal_index = cooling_row_major_index_3d(
            local_z_index, j, k, colibre_cooling_N_loaded_redshifts,
            colibre_cooling_N_metallicity, colibre_cooling_N_density);

        table->meanpartmass_Teq[internal_index] =
            mean_particle_mass_Teq[hdf5_index];
      }
    }

    /* Compute the pressures at thermal eq. for subgrid properties (if entropy
     * floor is used) */
    const float log10_kB_cgs = table->log10_kB_cgs;
    for (int j = 0; j < colibre_cooling_N_metallicity; j++) {

      const int index_XH = cooling_row_major_index_2d(
          j, 0, colibre_cooling_N_metallicity, colibre_cooling_N_elementtypes);

      const float log10_XH = table->LogMassFractions[index_XH];

      for (int k = 0; k < colibre_cooling_N_density; k++) {

        const int index_Peq = cooling_row_major_index_3d(
            local_z_index, j, k, colibre_cooling_N_loaded_redshifts,
            colibre_cooling_N_metallicity, colibre_cooling_N_density);

        table->logPeq[index_Peq] =
            table->nH[k] + table->logTeq[index_Peq] - log10_XH -
            log10(table->meanpartmass_Teq[index_Peq]) + log10_kB_cgs;
      }
    }

    /* read depletion table if running with dust evolution modelling */
    dust_read_depletion(
        file_id, &log_depletion_fractions,
        1 /* number of redshifts in hdf5 file */, colibre_cooling_N_temperature,
        colibre_cooling_N_metallicity, colibre_cooling_N_density,
        colibre_cooling_N_elementtypes);

    H5Fclose(file_id);
  }

  if (do_dust_pairing) {
    /** Modify COLIBRE heating and cooling rates to remove the effects
     * of implicit depletion when running with dust evolution **/
    depletion_correct_rates_split(
        table->Theating, table->Tcooling, log_depletion_fractions,
        colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
        colibre_cooling_N_metallicity, colibre_cooling_N_density,
        colibre_cooling_N_elementtypes, colibre_cooling_N_cooltypes,
        colibre_cooling_N_heattypes);
  }

  /** Reduce the cooling and heating tables to include only
   *  the channels required by CHIMES. **/
  int idx_red, idx_T, idx_met, idx_dens, idx_elem, idx_table, idx_reduced_table;
  float result;
  for (idx_red = 0; idx_red < colibre_cooling_N_loaded_redshifts; idx_red++) {
    for (idx_T = 0; idx_T < colibre_cooling_N_temperature; idx_T++) {
      for (idx_met = 0; idx_met < colibre_cooling_N_metallicity; idx_met++) {
        for (idx_dens = 0; idx_dens < colibre_cooling_N_density; idx_dens++) {

          /* Copy the metal elements to the reduced table. */
          for (idx_elem = 0; idx_elem < colibre_reduced_cooling_N_elementtypes;
               idx_elem++) {
            idx_table = cooling_row_major_index_5d(
                idx_red, idx_T, idx_met, idx_dens, idx_elem + element_C,
                colibre_cooling_N_loaded_redshifts,
                colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
                colibre_cooling_N_density, colibre_cooling_N_cooltypes);
            idx_reduced_table = cooling_row_major_index_5d(
                idx_red, idx_T, idx_met, idx_dens, idx_elem,
                colibre_cooling_N_loaded_redshifts,
                colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
                colibre_cooling_N_density, colibre_reduced_cooling_N_cooltypes);

            table->Tcooling_reduced[idx_reduced_table] =
                table->Tcooling[idx_table];
          }

          /* Copy eeBrems to the reduced table. */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, cooltype_eeBrems,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_cooltypes);
          idx_reduced_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, reduced_cooltype_eeBrems,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_reduced_cooling_N_cooltypes);

          table->Tcooling_reduced[idx_reduced_table] =
              table->Tcooling[idx_table];

          /* Combine remaining channels if they are
           * not already included in CHIMES. */
          /* element_OA */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, element_OA,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_cooltypes);
          result = exp10f(table->Tcooling[idx_table]);

          /* molecules */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, cooltype_molecules,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_cooltypes);
          result += exp10f(table->Tcooling[idx_table]);

          /* HD */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, cooltype_HD,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_cooltypes);
          result += exp10f(table->Tcooling[idx_table]);

          /* NetFFM */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, cooltype_NetFFM,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_cooltypes);
          result += exp10f(table->Tcooling[idx_table]);

          idx_reduced_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, reduced_cooltype_other,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_reduced_cooling_N_cooltypes);
          table->Tcooling_reduced[idx_reduced_table] =
              log10f(max(result, FLT_MIN));
        }
      }
    }
  }

  for (idx_red = 0; idx_red < colibre_cooling_N_loaded_redshifts; idx_red++) {
    for (idx_T = 0; idx_T < colibre_cooling_N_temperature; idx_T++) {
      for (idx_met = 0; idx_met < colibre_cooling_N_metallicity; idx_met++) {
        for (idx_dens = 0; idx_dens < colibre_cooling_N_density; idx_dens++) {

          /* Copy the metal elements to the reduced table. */
          for (idx_elem = 0; idx_elem < colibre_reduced_cooling_N_elementtypes;
               idx_elem++) {
            idx_table = cooling_row_major_index_5d(
                idx_red, idx_T, idx_met, idx_dens, idx_elem + element_C,
                colibre_cooling_N_loaded_redshifts,
                colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
                colibre_cooling_N_density, colibre_cooling_N_heattypes);
            idx_reduced_table = cooling_row_major_index_5d(
                idx_red, idx_T, idx_met, idx_dens, idx_elem,
                colibre_cooling_N_loaded_redshifts,
                colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
                colibre_cooling_N_density, colibre_reduced_cooling_N_heattypes);

            table->Theating_reduced[idx_reduced_table] =
                table->Theating[idx_table];
          }

          /* Combine remaining channels if they are
           * not already included in CHIMES. */
          /* element_OA */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, element_OA,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_heattypes);
          result = exp10f(table->Theating[idx_table]);

          /* COdiss */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, heattype_COdiss,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_heattypes);
          result += exp10f(table->Theating[idx_table]);

          /* UTA */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, heattype_UTA,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_heattypes);
          result += exp10f(table->Theating[idx_table]);

          /* line */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, heattype_line,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_heattypes);
          result += exp10f(table->Theating[idx_table]);

          /* Hlin */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, heattype_Hlin,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_heattypes);
          result += exp10f(table->Theating[idx_table]);

          /* ChaT */
          idx_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, heattype_ChaT,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_cooling_N_heattypes);
          result += exp10f(table->Theating[idx_table]);

          idx_reduced_table = cooling_row_major_index_5d(
              idx_red, idx_T, idx_met, idx_dens, reduced_heattype_other,
              colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
              colibre_cooling_N_metallicity, colibre_cooling_N_density,
              colibre_reduced_cooling_N_heattypes);
          table->Theating_reduced[idx_reduced_table] =
              log10f(max(result, FLT_MIN));
        }
      }
    }
  }

  swift_free("cooling_table.Tcooling", table->Tcooling);
  swift_free("cooling_table.Theating", table->Theating);
  swift_free("cooling-temp", cooling_T);
  swift_free("cooling-temp", heating_T);
  swift_free("cooling-temp", electron_fraction_T);
  swift_free("cooling-temp", equilibrium_temperature);
  swift_free("cooling-temp", mean_particle_mass_Teq);

#ifdef SWIFT_DEBUG_CHECKS
  message("Done reading in general cooling table");
#endif

#else
  error("Need HDF5 to read cooling tables");
#endif
}

/**
 * @brief Computes net metal heating rate from Colibre tables.
 *
 * Computes the net heating rate (heating - cooling) for a given element
 * abundance ratio, temperature, redshift, and density. The unit of the net
 * cooling rate is Lambda / nH**2 [erg cm^3 s-1] and all input values are in
 * cgs.
 *
 * @param myGasVars The #gasVariables struct.
 * @param myGlobalVars The #globalVariables struct.
 */
ChimesFloat colibre_metal_cooling_rate_temperature(
    struct gasVariables *myGasVars,
    const struct globalVariables *myGlobalVars) {
  struct global_hybrid_data_struct *myGlobalData;
  myGlobalData = (struct global_hybrid_data_struct *)myGlobalVars->hybrid_data;
  struct colibre_cooling_tables *table = myGlobalData->table;

  float log_T_cgs = log10f(myGasVars->temperature);
  float n_H_cgs = myGasVars->nH_tot;
  float Z_absolute = myGasVars->metallicity * myGlobalData->Zsol;

  /* Prevent electron fraction from
   * becoming negative. */
  float noneq_electron_fraction = max(myGasVars->abundances[sp_elec], 0.0f);

  const float *abundance_ratio;
  struct gas_hybrid_data_struct *myGasData;
  myGasData = (struct gas_hybrid_data_struct *)myGasVars->hybrid_data;
  abundance_ratio = myGasData->abundance_ratio;

  /* Set weights for cooling rates */
  float weights_cooling[colibre_reduced_cooling_N_cooltypes];
  for (int i = 0; i < colibre_reduced_cooling_N_cooltypes; i++) {

    if (i < colibre_reduced_cooling_N_elementtypes) {
      /* Only include metals that aren't
       * already included in CHIMES.
       * Note that the abundance_ratio array
       * includes H and He, so add element_C
       * to the index. */
      if (myGlobalVars->element_included[i] == 0)
        weights_cooling[i] = abundance_ratio[i + element_C];
      else
        weights_cooling[i] = 0.f;
    } else {
      /* use same abundances as in the tables */
      weights_cooling[i] = 1.f;
    }
  }

  /* Set weights for heating rates */
  float weights_heating[colibre_reduced_cooling_N_heattypes];
  for (int i = 0; i < colibre_reduced_cooling_N_heattypes; i++) {
    if (i < colibre_reduced_cooling_N_elementtypes) {
      /* Only include metals that aren't
       * already included in CHIMES.
       * Note that the abundance_ratio array
       * includes H and He, so add element_C
       * to the index. */
      if (myGlobalVars->element_included[i] == 0)
        weights_heating[i] = abundance_ratio[i + element_C];
      else
        weights_heating[i] = 0.f;
    } else {
      weights_heating[i] = 1.f; /* use same abundances as in the tables */
    }
  }

  // Set weights for electron densities.
  float weights_electron[colibre_cooling_N_electrontypes];
  for (int i = 0; i < colibre_cooling_N_electrontypes; i++) {
    if (i < colibre_cooling_N_elementtypes - 1)
      weights_electron[i] = abundance_ratio[i];
    else
      weights_electron[i] = 1.f;
  }

  /* Get indices of T, nH, and metallicity*/
  int T_index, n_H_index, met_index, red_index;
  float d_T, d_n_H, d_met;
  float logZZsol = log10f((Z_absolute / table->Zsol[0]) + FLT_MIN);

  cooling_get_index_1d(table->Temp, colibre_cooling_N_temperature, log_T_cgs,
                       &T_index, &d_T);
  cooling_get_index_1d(table->Metallicity, colibre_cooling_N_metallicity,
                       logZZsol, &met_index, &d_met);
  cooling_get_index_1d(table->nH, colibre_cooling_N_density, log10(n_H_cgs),
                       &n_H_index, &d_n_H);

  if (table->dz > 0.5)
    red_index = 1;
  else
    red_index = 0;

  // n_e / n_H from Colibre

  // Electron fraction from H + He + metals, with table metal ratios
  float colibre_electron_fraction_table = interpolation3d_fix_x_plus_summation(
      table->Telectron_fraction, weights_electron,
      colibre_cooling_N_electrontypes - 1, colibre_cooling_N_electrontypes - 1,
      red_index, T_index, met_index, n_H_index, d_T, d_met, d_n_H,
      colibre_cooling_N_loaded_redshifts, colibre_cooling_N_temperature,
      colibre_cooling_N_metallicity, colibre_cooling_N_density,
      colibre_cooling_N_electrontypes);

  float colibre_electron_fraction_metal_actual;
  if (met_index == 0) {
    /* Primordial abundances; metal
     * contribution is zero. */
    colibre_electron_fraction_metal_actual = 0.f;
  } else {
    /* From metals, with actual metal ratios.
     * We also need to exclude any metals
     * that are already included in CHIMES. */
    for (int i = element_C; i < element_OA; i++) {
      if (myGlobalVars->element_included[i - 2] == 1) weights_electron[i] = 0.f;
    }

    colibre_electron_fraction_metal_actual =
        interpolation3d_fix_x_plus_summation(
            table->Telectron_fraction, weights_electron, element_C,
            colibre_cooling_N_electrontypes - 4, red_index, T_index, met_index,
            n_H_index, d_T, d_met, d_n_H, colibre_cooling_N_loaded_redshifts,
            colibre_cooling_N_temperature, colibre_cooling_N_metallicity,
            colibre_cooling_N_density, colibre_cooling_N_electrontypes);
  }

  /* Adjust the weights based on the non-eq
   * electron fraction from CHIMES, compared
   * to the electron fraction that was used
   * in the Colibre tables. */
  float electron_fraction_ratio =
      (noneq_electron_fraction + colibre_electron_fraction_metal_actual) /
      (colibre_electron_fraction_table + FLT_MIN);

  for (int i = 0; i < colibre_reduced_cooling_N_cooltypes; i++)
    weights_cooling[i] *= electron_fraction_ratio;

  /* eeBrems needs an extra factor of ne */
  weights_cooling[reduced_cooltype_eeBrems] *= electron_fraction_ratio;

  int cool_start, heat_start;
  if (met_index == 0) {
    /* Primordial abundances; metal
     * contribution is zero. */
    cool_start = reduced_cooltype_eeBrems;
    heat_start = reduced_heattype_other;
  } else {
    cool_start = reduced_cooltype_C;
    heat_start = reduced_heattype_C;
  }

  /* Lambda / n_H**2 */
  const float cooling_rate = interpolation3d_fix_x_plus_summation(
      table->Tcooling_reduced, weights_cooling,            /* */
      cool_start, colibre_reduced_cooling_N_cooltypes - 1, /* */
      red_index, T_index, met_index, n_H_index,            /* */
      d_T, d_met, d_n_H,                                   /* */
      colibre_cooling_N_loaded_redshifts,                  /* */
      colibre_cooling_N_temperature,                       /* */
      colibre_cooling_N_metallicity,                       /* */
      colibre_cooling_N_density,                           /* */
      colibre_reduced_cooling_N_cooltypes);                /* */

  /* Gamma / n_H**2 */
  const float heating_rate = interpolation3d_fix_x_plus_summation(
      table->Theating_reduced, weights_heating,            /* */
      heat_start, colibre_reduced_cooling_N_heattypes - 1, /* */
      red_index, T_index, met_index, n_H_index,            /* */
      d_T, d_met, d_n_H,                                   /* */
      colibre_cooling_N_loaded_redshifts,                  /* */
      colibre_cooling_N_temperature,                       /* */
      colibre_cooling_N_metallicity,                       /* */
      colibre_cooling_N_density,                           /* */
      colibre_reduced_cooling_N_heattypes);                /* */

  /* Return the net heating rate (Lambda_heat - Lambda_cool) */
  return heating_rate - cooling_rate;
}

/**
 * @brief Allocate gas_hybrid_data struct in gasVars.
 *
 * @param myGasVars The #gasVariables struct.
 */
void chimes_allocate_gas_hybrid_data(struct gasVariables *myGasVars) {
  myGasVars->hybrid_data =
      (void *)malloc(sizeof(struct gas_hybrid_data_struct));
}

/**
 * @brief Allocate gas_hybrid_data struct in gasVars.
 *
 * @param myGasVars The #gasVariables struct.
 */
void chimes_free_gas_hybrid_data(struct gasVariables *myGasVars) {
  free(myGasVars->hybrid_data);
}
