#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <math.h>

#include "netcdf.h"

/*

 * Loop through all history tapes from the Community Land Model adding new

 * fields that are totals of the multi-level fields included in total_vars[]

 * below.

 */

static char *total_vars[] = {

	"SOILICE",

	"SOILLIQ"

};

static int nmtotal_vars = sizeof(total_vars)/sizeof(*total_vars);

static char *combo_vars[] = {

	"LATENT",		/* FCTR + FCEV + FGEV */

	"ET",			/* QVEGE + QVEGT + QSOIL */

	"ALL_SKY_ALBEDO",	/* FSR / FSDS */

	"BLACK_SKY_ALBEDO",	/* (FSRNDLN + FSRVDLN)/(FSDSNDLN + FSDSVDLN) */

	"NETRAD",		/* FSA - FIRA */

};

static char *combo_long_names[] = {

	"latent heat flux",

	"evapotranspiration",

	"surface all-sky albedo",

	"surface black-sky albedo",

	"net radiation"

};

static char *combo_units[] = {

	"watt/m^2",

	"m/s",

	"unitless",

	"unitless",

	"watt/m^2"

};

static float combo_FillValues[] = {

	1.e+36,

	1.e+36,

	1.e+36,

	1.e+36,

	1.e+36

};

static int nmcombo_vars = sizeof(combo_vars)/sizeof(*combo_vars);

static char *varname_total_prefix = "TOTAL_", *longname_total_prefix = "total ";

static char *time_name = "time", *lon_name = "lon", *lat_name = "lat";

static char *long_name_name = "long_name", *units_name = "units",

	*cell_method_name = "cell_method", *FillValue_name = "_FillValue",

	*missing_value_name = "missing_value";

static char *cell_method = "time: mean";

void wrap_nc(int status)

{

	if (status != NC_NOERR) {

		fprintf(stderr, "netCDF error: %s\n", nc_strerror(status));

		exit(1);

	}

	return;

}

void read_total_timeslice(int ncid, int varid, size_t nlevels, size_t d1,

	size_t d2, char FillFlag, float FillValue, size_t tslice, float *var,

	double *tvar)

{

	int i, j;

	size_t start[NC_MAX_VAR_DIMS], count[NC_MAX_VAR_DIMS];

	/* Initialize tvar to FillValue */

	for (i = 0; i < (d1 * d2); i++)

		tvar[i] = (double)FillValue;

	start[0] = tslice;

	count[0] = 1;

	start[2] = 0;

	count[2] = d1;

	start[3] = 0;

	count[3] = d2;

	/* March through the levels, totalling as we go */

	for (j = 0; j < nlevels; j++) {

		start[1] = j;

		count[1] = 1;

		wrap_nc(nc_get_vara_float(ncid, varid, start, count, var));

		for (i = 0; i < (d1 * d2); i++) {

			if (tvar[i] == (double)FillValue)

				tvar[i] = (double)var[i];

			else

				tvar[i] += (double)var[i];

		}

	}

	return;

}

void alloc_svars(size_t nsvars, size_t d1, size_t d2, float ***svarsp)

{

	int i;

	float **svars;

	if (!(*svarsp = (float **)malloc(nsvars * sizeof(float **)))) {

		perror("alloc_svars:svarsp");

		exit(2);

	}

	svars = *svarsp;

	for (i = 0; i < nsvars; i++) {

		if (!(svars[i] = (float *)malloc(d1 * d2 * sizeof(float)))) {

			perror("alloc_svars:svars[i]");

			exit(2);

		}

	}

	return;

}

void free_svars(size_t nsvars, float **svars)

{

	int i;

	for (i = 0; i < nsvars; i++)

		free(svars[i]);

	free(svars);

	return;

}

void read_timeslice(int ncid, size_t d1, size_t d2, size_t tslice, char *name,

	float *var)

{

	int varid;

	size_t start[NC_MAX_VAR_DIMS], count[NC_MAX_VAR_DIMS];

	start[0] = tslice;

	count[0] = 1;

	start[1] = 0;

	count[1] = d1;

	start[2] = 0;

	count[2] = d2;

	/* Get variable id */

	wrap_nc(nc_inq_varid(ncid, name, &varid));

	/* Assume the correct size and read the variable */

	wrap_nc(nc_get_vara_float(ncid, varid, start, count, var));

	return;

}

void read_compute_timeslice(int ncid, size_t d1, size_t d2, size_t tslice,

	int num, double *tvar)

{

	int i;

	size_t nsvars;

	float **svars; /* source variables */

	double denom;

	/* Initialize tvar to FillValue */

	for (i = 0; i < (d1 * d2); i++)

		tvar[i] = (double)combo_FillValues[num];

	switch (num) {

		case 0: /* LATENT */

			nsvars = 3;

			alloc_svars(nsvars, d1, d2, &svars);

			/* read timeslices */

			read_timeslice(ncid, d1, d2, tslice, "FCTR", svars[0]);

			read_timeslice(ncid, d1, d2, tslice, "FCEV", svars[1]);

			read_timeslice(ncid, d1, d2, tslice, "FGEV", svars[2]);

			/* compute new variable */

			for (i = 0; i < (d1 * d2); i++) {

				if (svars[0][i] != combo_FillValues[num] &&

					svars[1][i] != combo_FillValues[num] &&

					svars[2][i] != combo_FillValues[num]) {

					tvar[i] = (double)svars[0][i]

						+ (double)svars[1][i]

						+ (double)svars[2][i];

				}

			}

			free_svars(nsvars, svars);

			break;

		case 1: /* ET */

			nsvars = 3;

			alloc_svars(nsvars, d1, d2, &svars);

			/* read timeslices */

			read_timeslice(ncid, d1, d2, tslice, "QVEGE", svars[0]);

			read_timeslice(ncid, d1, d2, tslice, "QVEGT", svars[1]);

			read_timeslice(ncid, d1, d2, tslice, "QSOIL", svars[2]);

			/* compute new variable */

			for (i = 0; i < (d1 * d2); i++) {

				if (svars[0][i] != combo_FillValues[num] &&

					svars[1][i] != combo_FillValues[num] &&

					svars[2][i] != combo_FillValues[num]) {

					tvar[i] = (double)svars[0][i]

						+ (double)svars[1][i]

						+ (double)svars[2][i];

				}

			}

			free_svars(nsvars, svars);

			break;

		case 2: /* ALL_SKY_ALBEDO */

			nsvars = 2;

			alloc_svars(nsvars, d1, d2, &svars);

			/* read timeslices */

			read_timeslice(ncid, d1, d2, tslice, "FSR", svars[0]);

			read_timeslice(ncid, d1, d2, tslice, "FSDS", svars[1]);

			/* compute new variable */

			for (i = 0; i < (d1 * d2); i++) {

				if (svars[0][i] != combo_FillValues[num] &&

					svars[1][i] != combo_FillValues[num] &&

					(svars[1][i] > 1.e-35 || svars[1][i] < -1.e-35)) {

					tvar[i] = (double)svars[0][i]

						/ (double)svars[1][i];

				}

			}

			free_svars(nsvars, svars);

			break;

		case 3: /* BLACK_SKY_ALBEDO */

			nsvars = 4;

			alloc_svars(nsvars, d1, d2, &svars);

			/* read timeslices */

			read_timeslice(ncid, d1, d2, tslice, "FSRNDLN", svars[0]);

			read_timeslice(ncid, d1, d2, tslice, "FSRVDLN", svars[1]);

			read_timeslice(ncid, d1, d2, tslice, "FSDSNDLN", svars[2]);

			read_timeslice(ncid, d1, d2, tslice, "FSDSVDLN", svars[3]);

			/* compute new variable */

			for (i = 0; i < (d1 * d2); i++) {

				if (svars[0][i] != combo_FillValues[num] &&

					svars[1][i] != combo_FillValues[num] &&

					svars[2][i] != combo_FillValues[num] &&

					svars[3][i] != combo_FillValues[num]) {

					denom = (double)svars[2][i] + (double)svars[3][i];

					if (denom > 1.e-35 || denom < -1.e-35) {

						tvar[i] = ((double)svars[0][i]

							+ (double)svars[1][i]) /

							((double)svars[2][i]

							+ (double)svars[3][i]);

					}

				}

			}

			free_svars(nsvars, svars);

			break;

		case 4: /* NETRAD */

			nsvars = 2;

			alloc_svars(nsvars, d1, d2, &svars);

			/* read timeslices */

			read_timeslice(ncid, d1, d2, tslice, "FSA", svars[0]);

			read_timeslice(ncid, d1, d2, tslice, "FIRA", svars[1]);

			/* compute new variable */

			for (i = 0; i < (d1 * d2); i++) {

				if (svars[0][i] != combo_FillValues[num] &&

					svars[1][i] != combo_FillValues[num]) {

					tvar[i] = (double)svars[0][i]

						- (double)svars[1][i];

				}

			}

			free_svars(nsvars, svars);

			break;

		default:

			fprintf(stderr, "Unknown combination variable %d!\n", num);

			exit(3);

	}

	return;

}

void write_timeslice(int ncid, int varid, size_t d1, size_t d2,

	size_t tslice, float *var)

{

	size_t start[NC_MAX_VAR_DIMS], count[NC_MAX_VAR_DIMS];

	start[0] = tslice;

	count[0] = 1;

	start[1] = 0;

	count[1] = d1;

	start[2] = 0;

	count[2] = d2;

	wrap_nc(nc_put_vara_float(ncid, varid, start, count, var));

	return;

}

void usage(char *program)

{

	fprintf(stderr, "Usage: %s hist_file1.nc [ hist_file2.nc ... ]\n",

		program);

	fprintf(stderr, "Requires at least one history output file name\n");

	exit(5);

}

int main(int argc, char **argv)

{

	char **ifnames;

	int i, j, k, nifnames;

	int ncid, ngdims, nvars, ngatts, unlimdimid, varid, ndims, natts,

		dimids[NC_MAX_VAR_DIMS], new_ndims, new_dimids[NC_MAX_VAR_DIMS],

		new_varid, lat_dimid, lon_dimid;

	nc_type xtype, atype;

	char name[NC_MAX_NAME+1], new_name[NC_MAX_NAME+1],

		att_name[NC_MAX_NAME+1], *longname, *new_longname;

	size_t tlen, ts, alen, nlevels, d1, d2;

	float FillValue;

	char FillFlag;

	float *var;

	double *tvar;

#ifdef DEBUG

	printf("Number of total variables (nmtotal_vars) = %d\n", nmtotal_vars);

	printf("Number of combination variables (nmcombo_vars) = %d\n", nmcombo_vars);

#endif /* DEBUG */

	/* Require at least one parameter (the file on which to work) */

	if (argc < 2)

		usage(argv[0]);

	/* Put the list of filenames in an array of strings */

	if (!(ifnames = (char **)malloc((argc - 1) * sizeof(char *)))) {

		perror("ifnames");

		exit(2);

	}

	nifnames = 0;

	for (i = 1; i < argc; i++) {

		if (!(ifnames[nifnames++] = strdup(argv[i]))) {

			perror("ifnames[nifnames++]");

			exit(2);

		}

	}

	/*

	 * March through all input files adding new variables for totals

	 * across levels.

	 */

	for (i = 0; i < nifnames; i++) {

		printf("Processing %s\n", ifnames[i]);

		/* Open file */

		wrap_nc(nc_open(ifnames[i], NC_WRITE, &ncid));

		/*

		 * Inquire about number of dimensions, variables, global

		 * attributes and the ID of the unlimited dimension

		 */

		wrap_nc(nc_inq(ncid, &ngdims, &nvars, &ngatts, &unlimdimid));

		wrap_nc(nc_inq_dim(ncid, unlimdimid, name, &tlen));

		if (strcmp(name, time_name)) {

			fprintf(stderr, "%s is no the unlimited dimension for file %s!\n", time_name, ifnames[i]);

			exit(4);

		}

		/* Retrieve lat and lon dimension IDs */

		wrap_nc(nc_inq_dimid(ncid, lat_name, &lat_dimid));

		wrap_nc(nc_inq_dimid(ncid, lon_name, &lon_dimid));

		/* Put file into define mode */

		wrap_nc(nc_redef(ncid));

		/*

		 * Define all of the new variables first to improve

		 * performance.  This means some calls are made twice.

		 */

		/* First, add the new total variables */

		for (j = 0; j < nmtotal_vars; j++) {

			/* Figure out source variable information */

			wrap_nc(nc_inq_varid(ncid, total_vars[j], &varid));

			wrap_nc(nc_inq_var(ncid, varid, name, &xtype, &ndims, dimids, &natts));

			/* Make sure it has the correct number of dimensions */

			if (ndims != 4) {

				fprintf(stderr, "Variable %s has %d dimensions!\n", name, ndims);

				exit(4);

			}

			/* Make sure time is the first dimension */

			if (dimids[0] != unlimdimid) {

				fprintf(stderr, "First dimension of variable %s is not %s!\n", name, time_name);

				exit(4);

			}

			/* Make sure it is a float type */

			if (xtype != NC_FLOAT) {

				fprintf(stderr, "Only NC_FLOAT type is supported presently!\n");

				exit(4);

			}

			/* Set dimensions for new variable */

			new_ndims = 3;

			new_dimids[0] = unlimdimid;

			new_dimids[1] = dimids[2];

			new_dimids[2] = dimids[3];

			/* Define new variable */

			sprintf(new_name, "%s%s", varname_total_prefix, name);

			printf("\tAdding variable %s\n", new_name);

			wrap_nc(nc_def_var(ncid, new_name, xtype, new_ndims,

				new_dimids, &new_varid));

			/* Copy attributes from source variable to new one */

			for (k = 0; k < natts; k++) {

				wrap_nc(nc_inq_attname(ncid, varid, k, att_name));

				if (!strcmp(att_name, "long_name")) {

					wrap_nc(nc_inq_att(ncid, varid, att_name, &atype, &alen));

					if (!(longname = (char *)malloc((alen+1)*sizeof(char)))) {

						perror("longname");

						exit(2);

					}

					wrap_nc(nc_get_att_text(ncid, varid, att_name, longname));

					longname[alen] = '\0';

					if (!(new_longname = (char *)malloc((strlen(longname_total_prefix)+alen+1)*sizeof(char)))) {

						perror("new_longname");

						exit(2);

					}

					sprintf(new_longname, "%s%s", longname_total_prefix, longname);

					wrap_nc(nc_put_att_text(ncid, new_varid, att_name, strlen(new_longname), new_longname));

					free(new_longname);

					free(longname);

				}

				else

					wrap_nc(nc_copy_att(ncid, varid, att_name, ncid, new_varid));

			}

		}

		/* Second, add the new total variables */

		for (j = 0; j < nmcombo_vars; j++) {

			/* Set dimensions for new variable */

			new_ndims = 3;

			new_dimids[0] = unlimdimid;

			new_dimids[1] = lat_dimid;

			new_dimids[2] = lon_dimid;

			/* Define new variable */

			printf("\tAdding variable %s\n", combo_vars[j]);

			wrap_nc(nc_def_var(ncid, combo_vars[j], NC_FLOAT,

				new_ndims, new_dimids, &new_varid));

			/* Add attributes */

			wrap_nc(nc_put_att_text(ncid, new_varid, long_name_name,

				strlen(combo_long_names[j]),

				combo_long_names[j]));

			wrap_nc(nc_put_att_text(ncid, new_varid, units_name,

				strlen(combo_units[j]), combo_units[j]));

			wrap_nc(nc_put_att_text(ncid, new_varid,

				cell_method_name, strlen(cell_method),

				cell_method));

			wrap_nc(nc_put_att_float(ncid, new_varid,

				FillValue_name, NC_FLOAT, 1,

				&combo_FillValues[j]));

			wrap_nc(nc_put_att_float(ncid, new_varid,

				missing_value_name, NC_FLOAT, 1,

				&combo_FillValues[j]));

		}

		/* Leave define mode */

		wrap_nc(nc_enddef(ncid));

		/* Now actually compute and write out the new total variables */

		for (j = 0; j < nmtotal_vars; j++) {

			/* Figure out source variable information */

			wrap_nc(nc_inq_varid(ncid, total_vars[j], &varid));

			wrap_nc(nc_inq_var(ncid, varid, name, &xtype, &ndims, dimids, &natts));

			/* Check for _FillValue */

			FillFlag = 0;

			FillValue = 0.;

			if (nc_inq_att(ncid, varid, FillValue_name, &atype, &alen) == NC_NOERR) {

				if (atype == NC_FLOAT && alen) {

					wrap_nc(nc_get_att_float(ncid, varid,

						FillValue_name, &FillValue));

					FillFlag = 1;

				}

			}

			/* Set dimensions for new variable */

			new_ndims = 3;

			new_dimids[0] = unlimdimid;

			new_dimids[1] = dimids[2];

			new_dimids[2] = dimids[3];

			sprintf(new_name, "%s%s", varname_total_prefix, name);

			printf("\tComputing and writing %s\n", new_name);

			/* Retrieve the new varid again */

			wrap_nc(nc_inq_varid(ncid, new_name, &new_varid));

			/*

			 * Figure out dimensions and total size

			 * for a single time slice

			 */

			wrap_nc(nc_inq_dimlen(ncid, dimids[1], &nlevels));

			wrap_nc(nc_inq_dimlen(ncid, dimids[2], &d1));

			wrap_nc(nc_inq_dimlen(ncid, dimids[3], &d2));

			/*

			 * Allocate space for reading in time slice and for

			 * the sum

			 */

			if (!(var = (float *)malloc((d1*d2) * sizeof(float)))) {

				perror("var");

				exit(2);

			}

			if (!(tvar = (double *)malloc((d1*d2) * sizeof(double)))) {

				perror("tvar");

				exit(2);

			}

			/* Read timeslices, total, and write out new timeslices */

			for (ts = 0; ts < tlen; ts++) {

				read_total_timeslice(ncid, varid,

					nlevels, d1, d2, FillFlag, FillValue,

					ts, var, tvar);

				for (k = 0; k < (d1*d2); k++)

					var[k] = (float)tvar[k];

				write_timeslice(ncid, new_varid, d1, d2,

					ts, var);

			}

			free(var);

			free(tvar);

		}

		/* Now actually compute and write out the new combo variables */

		for (j = 0; j < nmcombo_vars; j++) {

			/* Set dimensions for new variable */

			new_ndims = 3;

			new_dimids[0] = unlimdimid;

			new_dimids[1] = lat_dimid;

			new_dimids[2] = lon_dimid;

			printf("\tComputing and writing %s\n", combo_vars[j]);

			/* Retrieve the new varid again */

			wrap_nc(nc_inq_varid(ncid, combo_vars[j], &new_varid));

			/*

			 * Figure out dimensions and total size

			 * for a single time slice

			 */

			wrap_nc(nc_inq_dimlen(ncid, new_dimids[1], &d1));

			wrap_nc(nc_inq_dimlen(ncid, new_dimids[2], &d2));

			/*

			 * Allocate space for reading in time slice and for

			 * the sum

			 */

			if (!(var = (float *)malloc((d1*d2) * sizeof(float)))) {

				perror("var");

				exit(2);

			}

			if (!(tvar = (double *)malloc((d1*d2) * sizeof(double)))) {

				perror("tvar");

				exit(2);

			}

			/* Read timeslices, compute new variable, and write */

			for (ts = 0; ts < tlen; ts++) {

				read_compute_timeslice(ncid, d1, d2, ts,

					j, tvar);

				for (k = 0; k < (d1*d2); k++)

					var[k] = (float)tvar[k];

				write_timeslice(ncid, new_varid, d1, d2,

					ts, var);

			}

			free(var);

			free(tvar);

		}

		/* Close file */

		wrap_nc(nc_close(ncid));

	}

	return 0;

}