#include "common.h"
#include <matio.h>

#include "e2e.h"
#include "export.h"

static size_t dims_scalar = 1;

#define COUNT(x) (sizeof(x) / sizeof((x)[0]))

matvar_t *create_int(const int *value)
{
	return Mat_VarCreate(NULL, MAT_C_INT32, MAT_T_INT32, 1, &dims_scalar,
	                     (int*)value, 0);
}

matvar_t *create_struct(const char **fields, int nfields)
{
	return Mat_VarCreateStruct(NULL, 1, &dims_scalar, fields, nfields);
}


static int export_tomogram(const struct e2e_slice *slices, matvar_t **tomogram,
                           matvar_t **slice_ids)
{
	size_t i, dims[3], zstride;
	int *ids;
	const struct e2e_slice *slice;
	float *image3d;

	dims[0] = slices->image->width;
	dims[1] = slices->image->height;
	dims[2] = HASH_COUNT(slices);

	ids = calloc(dims[2], sizeof(int));
	if (!ids)
		return ENOMEM;

	i = 0;
	HASH_FOR(slice, slices, hh)
		ids[i++] = slice->id;

	*slice_ids = Mat_VarCreate(NULL, MAT_C_INT32, MAT_T_INT32, 1, dims + 2,
	                           ids, 0);
	free(ids);
	if (!*slice_ids)
		return ENOMEM;

	// FIXME: check overflow
	zstride = dims[0] * dims[1];
	image3d = calloc(zstride * dims[2], sizeof(float));
	if (!image3d)
		goto error_image3d;

	i = 0;
	HASH_FOR(slice, slices, hh)
		memcpy(image3d + i * zstride, slice->image->tomogram,
	               zstride * sizeof(float));

	*tomogram = Mat_VarCreate(NULL, MAT_C_SINGLE, MAT_T_SINGLE, 3, dims,
	                          image3d, 0);
	free(image3d);
	if (!*tomogram)
		goto error_image3d;

	return 0;

error_image3d:
	Mat_VarFree(*slice_ids);
	return ENOMEM;
}

static int export_series(const struct e2e_series *series, matvar_t **out)
{
	int rv;
	size_t i, dims;
	const struct e2e_series *series1;

	dims = HASH_COUNT(series);
	*out = Mat_VarCreate(NULL, MAT_C_CELL, MAT_T_CELL, 1, &dims, NULL, 0);
	if (!*out)
		return ENOMEM;

	i = 0;
	HASH_FOR(series1, series, hh) {
		const char *fields[] = {
			"series_id",
			"tomogram",
			"tomogram_slice_ids",
			"fundi",
		};
		matvar_t *st, *id, *tomogram, *slice_ids;

		st = create_struct(fields, COUNT(fields));
		if (!st) {
			rv = ENOMEM;
			goto error;
		}

		id = create_int(&series1->id);
		if (!id) {
			rv = ENOMEM;
			goto error_id;
		}

		Mat_VarSetStructFieldByName(st, "series_id", 0, id);

		if (HASH_COUNT(series1->slices)) {
			rv = export_tomogram(series1->slices, &tomogram, &slice_ids);
			if (rv)
				goto error_tomogram;

			Mat_VarSetStructFieldByName(st, "tomogram", 0, tomogram);
			Mat_VarSetStructFieldByName(st, "tomogram_slice_ids", 0,
			                            slice_ids);
		}
		// TODO: export fundus images
		
		Mat_VarSetCell(*out, i++, st);
		continue;

	error_tomogram:
		Mat_VarFree(id);
	error_id:
		Mat_VarFree(st);
		goto error;
	}

	return 0;

error:
	Mat_VarFree(*out);
	return rv;
}

static int export_studies(const struct e2e_study *studies, matvar_t **out)
{
	int rv;
	size_t i, dims;
	const struct e2e_study *study;

	dims = HASH_COUNT(studies);
	*out = Mat_VarCreate(NULL, MAT_C_CELL, MAT_T_CELL, 1, &dims, NULL, 0);
	if (!*out)
		return ENOMEM;

	i = 0;
	HASH_FOR(study, studies, hh) {
		const char *fields[] = {"study_id", "series"};
		matvar_t *st, *id, *series;

		st = create_struct(fields, COUNT(fields));
		if (!st) {
			rv = ENOMEM;
			goto error;
		}

		id = create_int(&study->id);
		if (!id) {
			rv = ENOMEM;
			goto error_id;
		}

		rv = export_series(studies->series, &series);
		if (rv)
			goto error_series;

		Mat_VarSetStructFieldByName(st, "study_id", 0, id);
		Mat_VarSetStructFieldByName(st, "series", 0, series);
		Mat_VarSetCell(*out, i++, st);
		continue;

	error_series:
		Mat_VarFree(id);
	error_id:
		Mat_VarFree(st);
		goto error;
	}

	return 0;

error:
	Mat_VarFree(*out);
	return rv;
}

static int export_patients(const struct e2e_patient *patients, matvar_t **out)
{
	int rv;
	size_t i, dims;
	const struct e2e_patient *patient;

	dims = HASH_COUNT(patients);
	*out = Mat_VarCreate("patients", MAT_C_CELL, MAT_T_CELL, 1, &dims,
	                     NULL, 0);
	if (!*out)
		return ENOMEM;

	i = 0;
	HASH_FOR(patient, patients, hh) {
		const char *fields[] = {"patient_id", "studies"};
		matvar_t *st, *id, *studies;

		st = create_struct(fields, COUNT(fields));
		if (!st) {
			rv = ENOMEM;
			goto error;
		}

		id = create_int(&patient->id);
		if (!id) {
			rv = ENOMEM;
			goto error_id;
		}

		rv = export_studies(patients->studies, &studies);
		if (rv)
			goto error_studies;

		Mat_VarSetStructFieldByName(st, "patient_id", 0, id);
		Mat_VarSetStructFieldByName(st, "studies", 0, studies);
		Mat_VarSetCell(*out, i++, st);
		continue;

	error_studies:
		Mat_VarFree(id);
	error_id:
		Mat_VarFree(st);
		goto error;
	}

	return 0;

error:
	Mat_VarFree(*out);
	return rv;
}

int export(const struct e2e_data *data, const char *path)
{
	int rv = 0;
	mat_t *mat;
	matvar_t *patients;

	mat = Mat_Create(path, NULL);
	if (!mat) {
		errorf("error: couldn't open '%s' for writing\n", path);
		return EIO;
	}

	rv = export_patients(data->patients, &patients);
	if (!rv) {
		Mat_VarWrite(mat, patients, MAT_COMPRESSION_ZLIB);
		Mat_VarFree(patients);
	}

	Mat_Close(mat);
	return rv;
}