#include "hdf5.h"
#include "api.h"
#include <stdio.h>

#define HDF_TYPE	'H5'

#define BUFF_SIZE	32768

#pragma warning(disable: 4244)


class zsRegPrimitive
{
public:
	zsRegPrimitive(const char* name, int type, void* func) { api_add_primitive(name, type, func); }
};

class zsHDF
{
public:
	hid_t	fid, gid, did;
};

/////////////////////////////////////////////////////////////////////

void input_error(void* caller) { api_runtime_error(caller, "bad arguments");  }


void hdf_destroy(void *ptr)
{
	zsHDF *o = (zsHDF*)ptr;
	H5Dclose(o->did);
	H5Gclose(o->gid);
	H5Fclose(o->fid);
	delete o;
}

herr_t hdf_error(void *data) { return 1; }

/////////////////////////////////////////////////////////////////////

void* hdf_open(void *caller, int nargs, void** args)
{
	if (nargs < 1) input_error(caller);
	const char *fname = api_get_string(caller, args[0]);
	zsHDF *o = new zsHDF;
	o->did  = -1;
	H5Eset_auto(hdf_error, o);
	if (nargs > 1) {
		const char *s = api_get_string(caller, args[1]);
		if      (s[0] == 'w' || s[0] == 'W') o->fid = H5Fopen(fname, H5F_ACC_RDWR, H5P_DEFAULT);
		else if (s[0] == 'c' || s[0] == 'C') o->fid = H5Fcreate(fname, H5F_ACC_EXCL, H5P_DEFAULT, H5P_DEFAULT);
		else                                 o->fid = H5Fopen(fname, H5F_ACC_RDONLY, H5P_DEFAULT);
	}
	else {
		o->fid = H5Fopen(fname, H5F_ACC_RDONLY, H5P_DEFAULT);
	}
	if (o->fid < 0) return 0; //api_create_null();
	o->gid = H5Gopen(o->fid, "/");
	if (o->gid < 0) o->gid = H5Gcreate(o->fid, "/", 1);
	if (o->gid < 0) api_runtime_error(caller, "failed to open or create group");
	return api_create_user(caller, o, 0, hdf_destroy, HDF_TYPE);
}
static zsRegPrimitive hdf0("hdf", 0, hdf_open);

/*
void* hdf_close(void *caller, int nargs, void** args)
{
	if (nargs < 1) input_error(caller);
	zsHDF *o = hdf_check_type(caller, args[0]);
	H5Dclose(o->did);
	o->did = -1;
	H5Gclose(o->gid);
	o->gid = -1;
	H5Fclose(o->fid);
	o->fid = -1;
	return 0;
}
static zsRegPrimitive hdf0a("close", HDF_TYPE, hdf_close);
*/

void* hdf_version(void *caller, int nargs, void** args)
{
	unsigned v1, v2, v3;
	H5get_libversion(&v1, &v2, &v3);
	static char buf[64];
	sprintf(buf, "version-%d.%d.%d", v1, v2, v3);
	return api_create_string(caller, buf);
}
static zsRegPrimitive hdf1("version", HDF_TYPE, hdf_version);


void* hdf_group(void *caller, int nargs, void** args)
{
	if (nargs < 2) input_error(caller);
	zsHDF *o = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	const char *name = api_get_string(caller, args[1]);
	if (o->gid >= 0) H5Gclose(o->gid);
	o->gid = H5Gopen(o->fid, name);
	if (o->gid < 0) o->gid = H5Gcreate(o->fid, name, strlen(name));
	if (o->gid < 0) api_runtime_error(caller, "failed to open/create group");
	return api_create_integer(caller, o->gid);
}
static zsRegPrimitive hdf2("group", HDF_TYPE, hdf_group);


void* hdf_dataset(void *caller, int nargs, void** args)
{
	if (nargs < 2) input_error(caller);
	zsHDF *o = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	const char *name = api_get_string(caller, args[1]);
	if (nargs == 2) {
		if (o->did >= 0) H5Dclose(o->did);
		o->did = H5Dopen(o->gid, name);
		if (o->did < 0) return 0; // api_create_null();
		return api_create_integer(caller, o->did);
	}
	if (nargs < 4) input_error(caller);
	const char *type = api_get_string(caller, args[2]);
	hid_t htype;
	if      (!strcmp(type, "char"))   htype = H5Tcopy(H5T_NATIVE_CHAR);
	else if (!strcmp(type, "uchar"))  htype = H5Tcopy(H5T_NATIVE_UCHAR);
	else if (!strcmp(type, "short"))  htype = H5Tcopy(H5T_NATIVE_SHORT);
	else if (!strcmp(type, "ushort")) htype = H5Tcopy(H5T_NATIVE_USHORT);
	else if (!strcmp(type, "int"))    htype = H5Tcopy(H5T_NATIVE_INT);
	else if (!strcmp(type, "uint"))   htype = H5Tcopy(H5T_NATIVE_UINT);
	else if (!strcmp(type, "float"))  htype = H5Tcopy(H5T_NATIVE_FLOAT);
	else if (!strcmp(type, "double")) htype = H5Tcopy(H5T_NATIVE_DOUBLE);
	else 	input_error(caller);
	if (o->did >= 0) H5Dclose(o->did);
	hsize_t	dim[H5S_MAX_RANK];
	int i, k;
	for (i = 3, k = 0; i < nargs && k < H5S_MAX_RANK; i++, k++) dim[k] = api_get_integer(caller, args[i]);
	hid_t dspace = H5Screate_simple(nargs-3, dim, 0);
	if (dspace < 0) api_runtime_error(caller, "failed to create data space");
	o->did = H5Dcreate(o->gid, name, htype, dspace, H5P_DEFAULT);
	H5Sclose(dspace);
	H5Tclose(htype);
	if (o->did < 0) api_runtime_error(caller, "failed to create dataset");
	return api_create_integer(caller, o->did);
}
static zsRegPrimitive hdf3("dataset", HDF_TYPE, hdf_dataset);


int get_dataset_size(void *caller, hid_t id, int &ndim, hsize_t dim[], hid_t &type, int &tsize)
{
	if (id < 0) api_runtime_error(caller, "dataset not opened yet");
	hid_t dspace = H5Dget_space(id);
	if (dspace < 0) api_runtime_error(caller, "failed to get dataset space");
	ndim = H5Sget_simple_extent_dims(dspace, dim, 0);
	H5Sclose(dspace);
	type = H5Dget_type(id);
	if      (H5Tequal(type, H5T_NATIVE_CHAR) > 0)   tsize = 1;
	else if (H5Tequal(type, H5T_NATIVE_UCHAR) > 0)  tsize = 1;
	else if (H5Tequal(type, H5T_NATIVE_SHORT) > 0)  tsize = sizeof(short);
	else if (H5Tequal(type, H5T_NATIVE_USHORT) > 0) tsize = sizeof(short);
	else if (H5Tequal(type, H5T_NATIVE_INT) > 0)    tsize = sizeof(int);
	else if (H5Tequal(type, H5T_NATIVE_UINT) > 0)   tsize = sizeof(int);
	else if (H5Tequal(type, H5T_NATIVE_FLOAT) > 0)  tsize = sizeof(float);
	else if (H5Tequal(type, H5T_NATIVE_DOUBLE) > 0) tsize = sizeof(double);
	else api_runtime_error(caller, "un-supported data type");
	int i, size = 1;
	for (i = 0; i < ndim; i++) size *= dim[i];
	return size;
}


void* hdf_size(void *caller, int nargs, void** args)
{
	if (nargs < 1) input_error(caller);
	zsHDF *o = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	hsize_t dim[H5S_MAX_RANK];
	hid_t type;
	int ndim, tsize;
	int size = get_dataset_size(caller, o->did, ndim, dim, type, tsize);
	H5Tclose(type);
	void *arr = api_create_array(caller, ndim+2);
	api_set_array_object(caller, arr, "0", api_create_integer(caller,  size));
	api_set_array_object(caller, arr, "1", api_create_integer(caller, tsize));
	for (int i = 0; i < ndim; i++) {
		api_set_array_object2(caller, arr, i+2, api_create_integer(caller, dim[i]));
	}
	return arr;
}
static zsRegPrimitive hdf4("size", HDF_TYPE, hdf_size);


void* hdf_getdat(void *caller, int nargs, void** args)
{
	if (nargs < 4) input_error(caller);
	zsHDF *o  = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	void *ptr = api_get_ptr    (caller, args[1]);
	int n     = api_get_integer(caller, args[2]);
	int esize = api_get_integer(caller, args[3]);

	hsize_t dim[H5S_MAX_RANK];
	hid_t type, status;
	int ndim, tsize;
	int size = get_dataset_size(caller, o->did, ndim, dim, type, tsize);
	if (tsize != esize) api_runtime_error(caller, "incompatible type");

	if (nargs == 4) {
		if (size != n) api_runtime_error(caller, "incompatible size");
		status = H5Dread(o->did, type, H5S_ALL, H5S_ALL, H5P_DEFAULT, ptr);
	}
	else {
		hsize_t  mstart[1], dstart[H5S_MAX_RANK], dcount[H5S_MAX_RANK], mcount[1];
		mstart[0] = 0;
		mcount[0] = 1;
		int i, k;
		for (i = 5, k = 0; i < nargs && k < ndim; i += 2, k++) {
			dstart[k] = api_get_integer(caller, args[i-1]);
			dcount[k] = api_get_integer(caller, args[i  ]);
			if (dstart[k] < 0 ||
				dcount[k] < 1 ||
				dstart[k] + dcount[k] > dim[k]) api_runtime_error(caller, "bad dimemsion spec");
			mcount[0] *= dcount[k];
		}
		hid_t dspace  = H5Dget_space(o->did);
		hid_t status1 = H5Sselect_hyperslab(dspace, H5S_SELECT_SET, dstart, NULL, dcount, NULL);
		hid_t mspace  = H5Screate_simple(1, mcount, NULL);
		hid_t status2 = H5Sselect_hyperslab(mspace, H5S_SELECT_SET, mstart, NULL, mcount, NULL);
		if (status1 < 0 || status2 < 0) api_runtime_error(caller, "failed to create space");
		status = H5Dread(o->did, type, mspace, dspace, H5P_DEFAULT, ptr);
		H5Sclose(mspace);
		H5Sclose(dspace);
	}
	H5Tclose(type);
	if (status < 0) api_runtime_error(caller, "failed to get dataset");
	return 0;
}
static zsRegPrimitive hdf5("getdat", HDF_TYPE, hdf_getdat);


void* hdf_putdat(void *caller, int nargs, void** args)
{
	if (nargs < 2) input_error(caller);
	zsHDF *o = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	void *ptr = api_get_ptr   (caller, args[1]);
	
	hsize_t dim[H5S_MAX_RANK];
	hid_t type, status;
	int ndim, tsize;
	get_dataset_size(caller, o->did, ndim, dim, type, tsize);

	if (nargs == 2) {
		status = H5Dwrite(o->did, type, H5S_ALL, H5S_ALL, H5P_DEFAULT, ptr);
	}
	else {
		hsize_t  dstart[H5S_MAX_RANK], dcount[H5S_MAX_RANK], mstart[1], mcount[1];
		mstart[0] = 0;
		mcount[0] = 1;
		int i, k;
		for (i = 3, k = 0; i < nargs && k < ndim; i += 2, k++) {
			dstart[k] = api_get_integer(caller, args[i-1]);
			dcount[k] = api_get_integer(caller, args[i  ]);
			if (dstart[k] < 0 ||
				dcount[k] < 1 ||
				dstart[k] + dcount[k] > dim[k]) api_runtime_error(caller, "bad dimemsion spec");
			mcount[0] *= dcount[k];
		}
		hid_t dspace  = H5Dget_space(o->did);
		hid_t status1 = H5Sselect_hyperslab(dspace, H5S_SELECT_SET, dstart, NULL, dcount, NULL);
		hid_t mspace  = H5Screate_simple(1, mcount, NULL);
		hid_t status2 = H5Sselect_hyperslab(mspace, H5S_SELECT_SET, mstart, NULL, mcount, NULL);
		if (status1 < 0 || status2 < 0) api_runtime_error(caller, "failed to create space");
		status = H5Dwrite(o->did, type, mspace, dspace, H5P_DEFAULT, ptr);
		H5Sclose(mspace);
		H5Sclose(dspace);
	}
	H5Tclose(type);
	if (status < 0) api_runtime_error(caller, "failed to put dataset");

	return 0;
}
static zsRegPrimitive hdf6("putdat", HDF_TYPE, hdf_putdat);


void* hdf_getatt(void *caller, int nargs, void** args)
{
	if (nargs < 3) input_error(caller);
	zsHDF *o         = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	hid_t id         = api_get_integer(caller, args[1]);
	const char *name = api_get_string(caller,  args[2]);

	hid_t att_id = H5Aopen_name(id, name);
	if (att_id < 0) return 0;

	hid_t h = H5Aget_space(att_id);
	hsize_t dim[1];
	if (H5Sget_simple_extent_dims(h, dim, NULL) != 1) api_runtime_error(caller, "un-supported attribute type");
	if (dim[0]*sizeof(double) > BUFF_SIZE) api_runtime_error(caller, "attribute too long");
	H5Sclose(h);

	char buf[BUFF_SIZE];

	h = H5Aget_type(att_id);
	hid_t type = H5Tget_native_type(h, H5T_DIR_ASCEND);
	H5Tclose(h);
	if (H5Aread(att_id, type, (void*)buf) < 0) api_runtime_error(caller, "H5Aread() error");
	H5Aclose(att_id);

	if (H5Tequal(type, H5T_NATIVE_CHAR) > 0) {
		buf[dim[0]] = 0;
		H5Tclose(type);
		return api_create_string(caller, buf);
	}

	void *arr = 0;
	if (dim[0] > 1) arr = api_create_array(caller, dim[0]);

	if (H5Tequal(type, H5T_NATIVE_UCHAR) > 0) {
		if (dim[0] == 1) return api_create_integer(caller, ((unsigned char*)buf)[0]);
		for (int i = 0; i < dim[0]; i++) {
			api_set_array_object2(caller, arr, i, api_create_integer(caller, ((unsigned char*)buf)[i]));
		}
		return arr;
	}

	if (H5Tequal(type, H5T_NATIVE_SHORT) > 0) {
		if (dim[0] == 1) return api_create_integer(caller, ((short*)buf)[0]);
		for (int i = 0; i < dim[0]; i++) {
			api_set_array_object2(caller, arr, i, api_create_integer(caller, ((short*)buf)[i]));
		}
		return arr;
	}

	if (H5Tequal(type, H5T_NATIVE_USHORT) > 0) {
		if (dim[0] == 1) return api_create_integer(caller, ((unsigned short*)buf)[0]);
		for (int i = 0; i < dim[0]; i++) {
			api_set_array_object2(caller, arr, i, api_create_integer(caller, ((unsigned short*)buf)[i]));
		}
		return arr;
	}

	if (H5Tequal(type, H5T_NATIVE_INT) > 0) {
		if (dim[0] == 1) return api_create_integer(caller, ((int*)buf)[0]);
		for (int i = 0; i < dim[0]; i++) {
			api_set_array_object2(caller, arr, i, api_create_integer(caller, ((int*)buf)[i]));
		}
		return arr;
	}

	if (H5Tequal(type, H5T_NATIVE_UINT) > 0) {
		if (dim[0] == 1) return api_create_integer(caller, ((unsigned int*)buf)[0]);
		for (int i = 0; i < dim[0]; i++) {
			api_set_array_object2(caller, arr, i, api_create_integer(caller, ((unsigned int*)buf)[i]));
		}
		return arr;
	}

	if (H5Tequal(type, H5T_NATIVE_FLOAT) > 0) {
		if (dim[0] == 1) return api_create_real(caller, ((float*)buf)[0]);
		for (int i = 0; i < dim[0]; i++) {
			api_set_array_object2(caller, arr, i, api_create_real(caller, ((float*)buf)[i]));
		}
		return arr;
	}
	
	if (H5Tequal(type, H5T_NATIVE_DOUBLE) > 0) {
		if (dim[0] == 1) return api_create_real(caller, ((double*)buf)[0]);
		for (int i = 0; i < dim[0]; i++) {
			api_set_array_object2(caller, arr, i, api_create_real(caller, ((double*)buf)[i]));
		}
		return arr;
	}

	api_runtime_error(caller, "un-supported attribute type");

	return 0;
}
static zsRegPrimitive hdf7("getatt", HDF_TYPE, hdf_getatt);


void* hdf_putatt(void *caller, int nargs, void** args)
{
	if (nargs < 4) input_error(caller);
	zsHDF *o         = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	hid_t id         = api_get_integer(caller, args[1]);
	const char *name = api_get_string(caller,  args[2]);

	hid_t att_id, status, space;
	hsize_t i, dim[1] = { nargs - 3 };

	H5Adelete(id, name);

	if (api_is_real(args[3])) {
		if (dim[0] > 256) api_runtime_error(caller, "too many attributes");
		space = H5Screate_simple(1, dim, NULL);
		att_id = H5Acreate(id, name, H5T_NATIVE_DOUBLE, space, 0);
		if (att_id < 0) api_runtime_error(caller, "failed to create attribute");
		double data[256];
		for (i = 0; i < dim[0]; i++) data[i] = api_get_real(caller, args[i+3]);
		status = H5Awrite(att_id, H5T_NATIVE_DOUBLE, data);
	}
	else if (api_is_integer(args[3])) {
		if (dim[0] > 256) api_runtime_error(caller, "too many attributes");
		space = H5Screate_simple(1, dim, NULL);
		att_id = H5Acreate(id, name, H5T_NATIVE_INT, space, 0);
		if (att_id < 0) api_runtime_error(caller, "failed to create attribute.");
		int data[256];
		for (i = 0; i < dim[0]; i++) data[i] = api_get_integer(caller, args[i+3]);
		status = H5Awrite(att_id, H5T_NATIVE_INT, data);
	}
	else if (api_is_string(args[3])) {
		const char *ptr = api_get_string(caller, args[3]);
		dim[0] = strlen(ptr);
		space = H5Screate_simple(1, dim, NULL);
		att_id = H5Acreate(id, name, H5T_NATIVE_CHAR, space, 0);
		if (att_id < 0) api_runtime_error(caller, "failed to create attribute");
		status = H5Awrite(att_id, H5T_NATIVE_CHAR, ptr);
	}

	H5Sclose(space);
	H5Aclose(att_id);

	if (status < 0) api_runtime_error(caller, "H5Twrite() error");

	return 0;
}
static zsRegPrimitive hdf8("putatt", HDF_TYPE, hdf_putatt);


void* hdf_getcmm(void *caller, int nargs, void** args)
{
	if (nargs < 3) input_error(caller);
	zsHDF *o         = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	hid_t id		 = api_get_integer(caller,  args[1]);
	const char *name = api_get_string (caller,  args[2]);
	char buf[BUFF_SIZE];
	if (H5Gget_comment(id, name, BUFF_SIZE-2, buf) < 0) api_runtime_error(caller, "failed to get comment");
	return api_create_string(caller, buf);
}
static zsRegPrimitive hdf9("getcmm", HDF_TYPE, hdf_getcmm);


void* hdf_putcmm(void *caller, int nargs, void** args)
{
	if (nargs < 4) input_error(caller);
	zsHDF *o         = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	hid_t id		 = api_get_integer(caller,  args[1]);
	const char *name = api_get_string (caller,  args[2]);
	const char *comm = api_get_string (caller,  args[3]);
	if (H5Gset_comment(id, name, comm) < 0) api_runtime_error(caller, "failed to put comment");
	return 0;
}
static zsRegPrimitive hdf10("putcmm", HDF_TYPE, hdf_putcmm);


void* hdf_list(void *caller, int nargs, void** args)
{
	if (nargs < 2) input_error(caller);
	zsHDF *o = (zsHDF*)api_get_user(caller, args[0], HDF_TYPE);
	hid_t id = api_get_integer(caller, args[1]);

	hsize_t k, n;
	char buf[BUFF_SIZE];

	if (id == o->gid) {
		H5Gget_num_objs(id, &n);
		printf("------ group ------\n");
		for (k = 0; k < n; k++) {
			H5Gget_objname_by_idx(id, k, buf, BUFF_SIZE);
			hid_t type = H5Gget_objtype_by_idx(id, k);
			switch(type) {
			case H5G_LINK:
				printf("link");
				break;
			case H5G_GROUP:
				printf("group");
				break;
			case H5G_DATASET:
				printf("dataset");
				break;
			case H5G_TYPE:
				printf("type");
				break;
			default:
				printf("?");

			}
			printf(": %s\n", buf);
			H5Tclose(type);
		}
	}
	else if (id == o->did) {
		hid_t h = H5Dget_type(id);
		printf("------ dataset ------\ntype: ");
		if      (H5Tequal(h, H5T_NATIVE_CHAR) > 0)   printf("H5T_NATIVE_CHAR\n");
		else if (H5Tequal(h, H5T_NATIVE_UCHAR) > 0)  printf("H5T_NATIVE_UCHAR\n");
		else if (H5Tequal(h, H5T_NATIVE_SHORT) > 0)  printf("H5T_NATIVE_SHORT\n");
		else if (H5Tequal(h, H5T_NATIVE_USHORT) > 0) printf("H5T_NATIVE_USHORT\n");
		else if (H5Tequal(h, H5T_NATIVE_INT) > 0)    printf("H5T_NATIVE_INT\n");
		else if (H5Tequal(h, H5T_NATIVE_UINT) > 0)   printf("H5T_NATIVE_UINT\n");
		else if (H5Tequal(h, H5T_NATIVE_FLOAT) > 0)  printf("H5T_NATIVE_FLOAT\n");
		else if (H5Tequal(h, H5T_NATIVE_DOUBLE) > 0) printf("H5T_NATIVE_DOUBLE\n");
		else    printf("?\n");
		H5Tclose(h);
		h = H5Dget_space(id);
		hsize_t dim[H5S_MAX_RANK];
		n = H5Sget_simple_extent_dims(h, dim, 0);
		H5Sclose(h);
		printf("dimensions:");
		for (k = 0; k < n; k++) printf(" %d", dim[k]);
		printf("\n");
	}
	else {
		input_error(caller);
	}

	n = H5Aget_num_attrs(id);
	printf("attibute:");
	for (k = 0; k < n; k++) {
		hid_t aid = H5Aopen_idx(id, (unsigned int)k);
		H5Aget_name(aid, BUFF_SIZE, buf);
		H5Aclose(aid);
		printf(" %s", buf);
	}
	printf("\n");

	fflush(stdout);
	return 0;
}
static zsRegPrimitive hdf11("list", HDF_TYPE, hdf_list);