#include <stdio.h>
#include <math.h>
#include <string.h>
#include <errno.h>
#include "boosterskel.h"

#define USAGE "Usage: don't."
#include <inttypes.h>
#include <stdlib.h>
#include <hdf5.h>


#define N_COLS 11
#define CHUNK_SIZE 501
typedef struct InRec_s {
  float chi2_opt;
  float dec;
  float feh_confidence;
  int64_t gdr3_source_id;
  float ln_prior;
  float logg_confidence;
  unsigned char quality_flags;
  float ra;
  float stellar_params_err[6];
  float stellar_params_est[6];
  float teff_confidence;
} InRec;
#define HDFGUARD(x) if ((x)<0) abort();
#define QUERY_N_PARS 23

enum outputFields {
	fi_source_id,    /* Source Id, bigint */
	fi_ra,           /* RA, double precision */
	fi_dec,          /* Dec, double precision */
	fi_teff,         /* Teff, real */
	fi_fe_h,         /* [Fe/H], real */
	fi_logg,         /* log(g), real */
	fi_xi,           /* ξ, real */
	fi_r55,          /* R(55), real */
	fi_ext,          /* Ext., real */
	fi_mod_parallax,  /* Est. Par., real */
	fi_err_teff,     /* Err. Teff, real */
	fi_err_fe_h,     /* Err. [Fe/H], real */
	fi_err_logg,     /* Err. log(g), real */
	fi_err_xi,       /* Err. ξ, real */
	fi_err_r55,      /* Err. R(55), real */
	fi_err_ext,      /* Err. Ext., real */
	fi_err_mod_parallax,  /* Err. Par., real */
	fi_chi2_opt,     /* χ², real */
	fi_ln_prior,     /* ln(prior), real */
	fi_teff_confidence,  /* Conf. T, real */
	fi_feh_confidence,  /* Conf. [Fe/H], real */
	fi_logg_confidence,  /* Conf. logg, real */
	fi_quality_flags,  /* Quality Flags, smallint */
};

/* We don't want to complicate compiling and linking this stuff and just
include Gregory's code literally, with a few trivial changes such as
dropping/fixing includes and removing the main function. */

/* ********************************************************* */



// Boilerplate code for extinction_curve.c. I will need to
// use basic math function and standard I/O functions.
// Finally, I need a basic main function to test the code.


double sech2(double x) {
    // Hyperbolic secant squared function:
    // sech^2(x) = 1 / cosh^2(x) = 1 - tanh^2(x).
    // double cosh_x = cosh(x);
    // return 1.0 / (cosh_x * cosh_x);
    double tanh_x = tanh(x);
    return 1.0 - tanh_x * tanh_x;
}

/* Single-wavelength calculations of A, R and their derivatives */

double A(double a, double b, double xi) {
    // Extinction at one wavelength, as a function of
    // xi, which is a parameter equivalent to R(V).
    // There is a different set of a, b for each wavelength.
    return exp(a + b * tanh(xi));
}

double dA_dxi(double a, double b, double xi) {
    // Derivative of extinction with respect to xi.
    return A(a, b, xi) * (b * sech2(xi));
}

double R(double a0, double b0, double a1, double b1, double xi) {
    // R := A(lambda_1) / [A(lambda_0) - A(lambda_1)]
    double A0 = A(a0, b0, xi);
    double A1 = A(a1, b1, xi);
    return A1 / (A0 - A1);
}

double dR_dxi(double a0, double b0, double a1, double b1, double xi) {
    // Derivative of R with respect to xi:
    // dR/dxi = [A0*dA1/dxi - A1*dA0/dxi] / (A1 - A0)^2
    double A0 = A(a0, b0, xi);
    double A1 = A(a1, b1, xi);
    double dA0_dxi = dA_dxi(a0, b0, xi);
    double dA1_dxi = dA_dxi(a1, b1, xi);

    double num = A0 * dA1_dxi - A1 * dA0_dxi;
    double denom = (A1 - A0) * (A1 - A0);

    return num / denom;
}

/* Same as above, but interpolating between different wavelengths.
   For example, 440 and 550 nm do not lie at anchor points of
   the extinction model. */

double A_interp(double a_low, double b_low,
                double a_high, double b_high,
                double frac_high,
                double xi) {
    // Extinction at one wavelength, as a function of
    // xi, which is a parameter equivalent to R(V).
    // There is a different set of a, b for each wavelength.
    double a = a_low * (1.0 - frac_high) + a_high * frac_high;
    double b = b_low * (1.0 - frac_high) + b_high * frac_high;
    return exp(a + b * tanh(xi));
}

double dA_interp_dxi(double a_low, double b_low,
                     double a_high, double b_high,
                     double frac_high,
                     double xi) {
    // Derivative of extinction with respect to xi.
    double dA_dxi_low = dA_dxi(a_low, b_low, xi);
    double dA_dxi_high = dA_dxi(a_high, b_high, xi);
    return dA_dxi_low * (1.0 - frac_high) + dA_dxi_high * frac_high;
}

double R_interp(double a0_low, double b0_low,
                double a0_high, double b0_high,
                double a1_low, double b1_low,
                double a1_high, double b1_high,
                double frac_high_0, double frac_high_1,
                double xi) {
    // R := A(lambda_1) / [A(lambda_0) - A(lambda_1)]
    double A0 = A_interp(a0_low, b0_low, a0_high, b0_high, frac_high_0, xi);
    double A1 = A_interp(a1_low, b1_low, a1_high, b1_high, frac_high_1, xi);
    return A1 / (A0 - A1);
}

double dR_interp_dxi(double a0_low, double b0_low,
                     double a0_high, double b0_high,
                     double a1_low, double b1_low,
                     double a1_high, double b1_high,
                     double frac_high_0, double frac_high_1,
                     double xi) {
    // Derivative of R with respect to xi:
    // dR/dxi = [A0*dA1/dxi - A1*dA0/dxi] / (A1 - A0)^2
    double A0 = A_interp(a0_low, b0_low, a0_high, b0_high, frac_high_0, xi);
    double A1 = A_interp(a1_low, b1_low, a1_high, b1_high, frac_high_1, xi);
    double dA0_dxi = dA_interp_dxi(a0_low, b0_low,
                                   a0_high, b0_high,
                                   frac_high_0, xi);
    double dA1_dxi = dA_interp_dxi(a1_low, b1_low,
                                   a1_high, b1_high,
                                   frac_high_1, xi);

    double num = A0 * dA1_dxi - A1 * dA0_dxi;
    double denom = (A1 - A0) * (A1 - A0);

    return num / denom;
}

void xi_to_R(double a0_low, double b0_low,
             double a0_high, double b0_high,
             double a1_low, double b1_low,
             double a1_high, double b1_high,
             double frac_high_0, double frac_high_1,
             double xi, double xi_err,
             double *R, double *R_err) {
    // Function that transforms (xi, xi_err) to (R, R_err) using linear
    // propagation of errors: R_err = |dR/dxi| * xi_err
    *R = R_interp(a0_low, b0_low,
                  a0_high, b0_high,
                  a1_low, b1_low,
                  a1_high, b1_high,
                  frac_high_0, frac_high_1,
                  xi);
    double dR_dxi_val = dR_interp_dxi(a0_low, b0_low,
                                      a0_high, b0_high,
                                      a1_low, b1_low,
                                      a1_high, b1_high,
                                      frac_high_0, frac_high_1,
                                      xi);
    *R_err = fabs(dR_dxi_val) * xi_err;
}

void xi_to_R55(double xi, double xi_err,
               double *R55, double *R55_err) {
    // Convenience function to calculate the conversion
    // to R(550 nm).

    // Coefficients for 440 nm
    double a_440[] = {1.3581201, 1.3495619};
    double b_440[] = {-0.0557438, -0.02957537};
    double frac_high_440 = 0.8;

    // Coefficients for 550 nm
    double a_550[] = {1.0725757, 1.0338929};
    double b_550[] = {0.22837315, 0.25460893};
    double frac_high_550 = 0.8;

    xi_to_R(a_440[0], b_440[0],
            a_440[1], b_440[1],
            a_550[0], b_550[0],
            a_550[1], b_550[1],
            frac_high_440, frac_high_550,
            xi, xi_err,
            R55, R55_err);
}

/* ****************************************************** */
/* (End gregory's code) */


Field *getTuple(InRec *data, int rowIndex)
{
	static Field vals[QUERY_N_PARS];
	double r55, err_r55;

	xi_to_R55(
		(double)(data->stellar_params_est[3]), (double)(data->stellar_params_err[3]),
		&r55, &err_r55);
	
	/* Gaia DR3 unique source identifier. You can match this against gaia.dr3lite on this TAP service. (bigint) */
	MAKE_BIGINT(fi_source_id, ((uint64_t)(data->gdr3_source_id)));
	/* Gaia ICRS right ascension for this object. (double precision) */
	MAKE_DOUBLE(fi_ra, ((float)(data->ra)));
	/* Gaia ICRS declination for this object. (double precision) */
	MAKE_DOUBLE(fi_dec, ((float)(data->dec)));

	/* Estimated effective Temperature (real) */
	MAKE_FLOAT(fi_teff, data->stellar_params_est[0]*1000);
	/* Log of Fe/H in solar units (real) */
	MAKE_FLOAT(fi_fe_h, data->stellar_params_est[1]);
	/* Log of surface gravity in solar units (real) */
	MAKE_FLOAT(fi_logg, data->stellar_params_est[2]);
	/* xi (real) */
	MAKE_FLOAT(fi_xi, data->stellar_params_est[3]);
	MAKE_FLOAT(fi_r55, r55);

	/* ext (real) */
	MAKE_FLOAT(fi_ext, data->stellar_params_est[4]);
	/* Parallax estimated from the model. (real) */
	MAKE_FLOAT(fi_mod_parallax, data->stellar_params_est[5]);
	/* Error in estimated effective Temperature (real) */
	MAKE_FLOAT(fi_err_teff, data->stellar_params_err[0]*1000);
	/* Error in fe_h (real) */
	MAKE_FLOAT(fi_err_fe_h, data->stellar_params_err[1]);
	/* Error in log_g (real) */
	MAKE_FLOAT(fi_err_logg, data->stellar_params_err[2]);
	/* Error in err_xi (real) */
	MAKE_FLOAT(fi_err_xi, data->stellar_params_err[3]);
	MAKE_FLOAT(fi_err_r55, err_r55);

	/* Error in ext (real) */
	MAKE_FLOAT(fi_err_ext, data->stellar_params_err[4]);
	/* Error in the parallax estimated from the model. (real) */
	MAKE_FLOAT(fi_err_mod_parallax, data->stellar_params_err[5]);

	/* χ² of the best-fit solution (real) */
	MAKE_FLOAT(fi_chi2_opt, ((float)(data->chi2_opt)));
	/* Natural log of the GMM prior on stellar type, at the location of the optimal solution. (real) */
	MAKE_FLOAT(fi_ln_prior, ((float)(data->ln_prior)));
	/* A neural-network-based estimate of the confidence in the effective temperature estimate, on a scale of 0 (no confidence) to 1 (high confidence). (real) */
	MAKE_FLOAT(fi_teff_confidence, ((float)(data->teff_confidence)));
	/* A neural-network-based estimate of the confidence in the [Fe/H] estimate, on a scale of 0 (no confidence) to 1 (high confidence). (real) */
	MAKE_FLOAT(fi_feh_confidence, ((float)(data->feh_confidence)));
	/* A neural-network-based estimate of the confidence in the log(g) estimate, on a scale of 0 (no confidence) to 1 (high confidence). (real) */
	MAKE_FLOAT(fi_logg_confidence, ((float)(data->logg_confidence)));
	/* The three least significant bits represent whether the confidence in effective temperature, [Fe/H] and log(g) is less than 0.5, respectively. The 4th bit is set if chi2_opt/61 > 2. The 5th bit is set if ln_prior < -7.43. The 6th bit is set if our parallax estimate is more than 10 sigma from the GDR3 measurement (using reported parallax uncertainties from GDR3). The two most significant bits are always unset. We recommend a cut of quality_flags < 8 (the "basic reliability cut"), although a stricter cut of quality_flags == 0 ensures higher reliability at the cost of lower completeness. (smallint) */
	MAKE_SHORT(fi_quality_flags, ((unsigned char)(data->quality_flags)));
	return vals;
}


/* Common main header */
void createDumpfile(int argc, char **argv)
{
	FILE *destination=stdout;
	char inputLine[INPUT_LINE_MAX];
	int recCount = 0;
/* /Common main header */

	hid_t input_id;	hid_t datasets[N_COLS];
	hid_t colspace;
	hid_t arr_colspace;
	size_t nrecs;
	int rank;
	hsize_t dims[1], offsets[1];
	hsize_t null_offset[1] = {0};
	hsize_t arr_null_offsets[2] = {0,0};
	hsize_t chunk_size[1] = {CHUNK_SIZE};
	hid_t memspace = H5Screate_simple(1, chunk_size, NULL);
	size_t index_in_chunk = CHUNK_SIZE;
  // I admit I've not bothered to figure out nD-arrays in HDF5.  Oh boy.
  hsize_t arr_chunk_size[2] = {CHUNK_SIZE, 6};
  hsize_t arr_offsets[2] = {0, 0};
  hid_t arr_memspace = H5Screate_simple(2, arr_chunk_size, NULL);
	size_t total_read = 0;
	InRec cur_rec;


	/* common main intro */
	writeHeader(destination);
	/* /common main intro */

	input_id = H5Fopen(argv[1], H5F_ACC_RDONLY, H5P_DEFAULT);
	if (input_id<0) abort();
	HDFGUARD(H5Sselect_hyperslab(memspace, H5S_SELECT_SET, 		null_offset, NULL, chunk_size, NULL));
  float arr_chi2_opt[CHUNK_SIZE];
  float arr_dec[CHUNK_SIZE];
  float arr_feh_confidence[CHUNK_SIZE];
  int64_t arr_gdr3_source_id[CHUNK_SIZE];
  float arr_ln_prior[CHUNK_SIZE];
  float arr_logg_confidence[CHUNK_SIZE];
  unsigned char arr_quality_flags[CHUNK_SIZE];
  float arr_ra[CHUNK_SIZE];
  float arr_stellar_params_err[CHUNK_SIZE*6];
  float arr_stellar_params_est[CHUNK_SIZE*6];
  float arr_teff_confidence[CHUNK_SIZE];
  datasets[0] = H5Dopen2(input_id, "chi2_opt", H5P_DEFAULT);
  if (datasets[0]<0) abort();
  datasets[1] = H5Dopen2(input_id, "dec", H5P_DEFAULT);
  if (datasets[1]<0) abort();
  datasets[2] = H5Dopen2(input_id, "feh_confidence", H5P_DEFAULT);
  if (datasets[2]<0) abort();
  datasets[3] = H5Dopen2(input_id, "gdr3_source_id", H5P_DEFAULT);
  if (datasets[3]<0) abort();
  datasets[4] = H5Dopen2(input_id, "ln_prior", H5P_DEFAULT);
  if (datasets[4]<0) abort();
  datasets[5] = H5Dopen2(input_id, "logg_confidence", H5P_DEFAULT);
  if (datasets[5]<0) abort();
  datasets[6] = H5Dopen2(input_id, "quality_flags", H5P_DEFAULT);
  if (datasets[6]<0) abort();
  datasets[7] = H5Dopen2(input_id, "ra", H5P_DEFAULT);
  if (datasets[7]<0) abort();
  datasets[8] = H5Dopen2(input_id, "stellar_params_err", H5P_DEFAULT);
  if (datasets[8]<0) abort();
  datasets[9] = H5Dopen2(input_id, "stellar_params_est", H5P_DEFAULT);
  if (datasets[9]<0) abort();
  datasets[10] = H5Dopen2(input_id, "teff_confidence", H5P_DEFAULT);
  if (datasets[10]<0) abort();
	colspace = H5Dget_space(datasets[0]);
	rank = H5Sget_simple_extent_ndims(colspace);
	if (rank>1) abort();

	arr_colspace = H5Dget_space(datasets[8]);

	H5Sget_simple_extent_dims(colspace, dims, NULL);
	nrecs = dims[0];
	while (total_read<nrecs) {

		Field *tuple;
		context = inputLine;
		if (!setjmp(ignoreRecord)) {

			if (index_in_chunk>=CHUNK_SIZE) {
				offsets[0] = total_read;
				arr_offsets[0] = total_read;
				if (total_read+chunk_size[0]>nrecs) {
          chunk_size[0] = nrecs-total_read;
          // see above on me not having worked out HDF5 arrays.
          arr_chunk_size[0] = chunk_size[0];
          HDFGUARD(H5Sselect_hyperslab(memspace, H5S_SELECT_SET,
          	null_offset, NULL, chunk_size, NULL));
          HDFGUARD(H5Sselect_hyperslab(arr_memspace, H5S_SELECT_SET,
          	arr_null_offsets, NULL, arr_chunk_size, NULL));
        }
				HDFGUARD(H5Sselect_hyperslab(colspace, H5S_SELECT_SET,  offsets, NULL, chunk_size, NULL));
        HDFGUARD(H5Sselect_hyperslab(arr_colspace, H5S_SELECT_SET,
          	arr_offsets, NULL, arr_chunk_size, NULL));

				HDFGUARD(H5Dread(datasets[0], H5T_NATIVE_FLOAT, memspace, colspace, H5P_DEFAULT, arr_chi2_opt));
				HDFGUARD(H5Dread(datasets[1], H5T_NATIVE_FLOAT, memspace, colspace, H5P_DEFAULT, arr_dec));
				HDFGUARD(H5Dread(datasets[2], H5T_NATIVE_FLOAT, memspace, colspace, H5P_DEFAULT, arr_feh_confidence));
				HDFGUARD(H5Dread(datasets[3], H5T_NATIVE_LLONG, memspace, colspace, H5P_DEFAULT, arr_gdr3_source_id));
				HDFGUARD(H5Dread(datasets[4], H5T_NATIVE_FLOAT, memspace, colspace, H5P_DEFAULT, arr_ln_prior));
				HDFGUARD(H5Dread(datasets[5], H5T_NATIVE_FLOAT, memspace, colspace, H5P_DEFAULT, arr_logg_confidence));
				HDFGUARD(H5Dread(datasets[6], H5T_NATIVE_UCHAR, memspace, colspace, H5P_DEFAULT, arr_quality_flags));
				HDFGUARD(H5Dread(datasets[7], H5T_NATIVE_FLOAT, memspace, colspace, H5P_DEFAULT, arr_ra));
				HDFGUARD(H5Dread(datasets[8], H5T_NATIVE_FLOAT, arr_memspace, arr_colspace, H5P_DEFAULT, arr_stellar_params_err));
				HDFGUARD(H5Dread(datasets[9], H5T_NATIVE_FLOAT, arr_memspace, arr_colspace, H5P_DEFAULT, arr_stellar_params_est));
				HDFGUARD(H5Dread(datasets[10], H5T_NATIVE_FLOAT, memspace, colspace, H5P_DEFAULT, arr_teff_confidence));
			index_in_chunk = 0;
		}
		cur_rec.chi2_opt = arr_chi2_opt[index_in_chunk];
		cur_rec.dec = arr_dec[index_in_chunk];
		cur_rec.feh_confidence = arr_feh_confidence[index_in_chunk];
		cur_rec.gdr3_source_id = arr_gdr3_source_id[index_in_chunk];
		cur_rec.ln_prior = arr_ln_prior[index_in_chunk];
		cur_rec.logg_confidence = arr_logg_confidence[index_in_chunk];
		cur_rec.quality_flags = arr_quality_flags[index_in_chunk];
		cur_rec.ra = arr_ra[index_in_chunk];
    memcpy(cur_rec.stellar_params_err,
            &(arr_stellar_params_err[index_in_chunk*6]),
            6*sizeof(float));
    memcpy(cur_rec.stellar_params_est,
            &(arr_stellar_params_est[index_in_chunk*6]),
            6*sizeof(float));

		cur_rec.teff_confidence = arr_teff_confidence[index_in_chunk];
		total_read += 1;
		index_in_chunk += 1;
		tuple = getTuple(&cur_rec, total_read-1);

			if (!tuple) {
				handleBadRecord("Bad input line at record %d", recCount);
			}
			writeTuple(tuple, QUERY_N_PARS, destination);
			context = NULL;
			recCount ++;
			if (!(recCount%1000)) {
				fprintf(stderr, "%08d\r", recCount);
				fflush(stderr);
			}
		}

	}

	writeEndMarker(destination);
	fprintf(stderr, "%08d records done.\n", recCount);
}

