/*!\file gr3_mc.c
 *
 * Optimized OpenMP implementation of the marching cubes algorithm.
 *
 * This code is based on Paul Bourke's Marching Cubes implementation
 * (http://paulbourke.net/geometry/polygonise/)
 *
 * Creates an indexed mesh to reduce the number of vertices to calculate.
 * This is done by caching generated vertices depending on their location.
 * Multiple threads create independent meshes.
 * Caches values between adjacent cubes.
 *
 * Fabian Beule
 * 2014-02-10
 */

#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "gr3.h"
#include "gr3_mc_data.h"
#ifdef _OPENMP
#include <omp.h>
#endif

#define ABS(x) ((x) < 0 ? -(x) : (x))
#define INDEX(x,y,z) ((x)*mcdata.stride[0]+(y)*mcdata.stride[1]+(z)*mcdata.stride[2])
#define IDX2D(y,z) ((y)*mcdata.dim[2] + (z))

/* speedup does not grow much with a high number of threads */
#define THREADLIMIT 16

/* for smaller function headers */
typedef struct {
    const GR3_MC_DTYPE *data;
    GR3_MC_DTYPE isolevel;
    int dim[3];
    int stride[3];
    double step[3];
    double offset[3];
} mcdata_t;

/* calculate the gradient via difference qoutient */
static gr3_coord_t getgrad(mcdata_t mcdata, int x, int y, int z)
{
    int v[3];
    int neigh[3][2];
    int i;
    gr3_coord_t n;

	v[0] = x;
	v[1] = y;
	v[2] = z;

    for (i = 0; i < 3; i++) {
        if (v[i] > 0)
            neigh[i][0] = v[i] - 1;
        else
            neigh[i][0] = v[i];
        if (v[i] < mcdata.dim[i] - 1)
            neigh[i][1] = v[i] + 1;
        else
            neigh[i][1] = v[i];
    }
    n.x = (float) (mcdata.data[INDEX(neigh[0][1], y, z)]
           - mcdata.data[INDEX(neigh[0][0], y, z)])
          / (neigh[0][1] - neigh[0][0]) / mcdata.step[0];
    n.y = (float) (mcdata.data[INDEX(x, neigh[1][1], z)]
           - mcdata.data[INDEX(x, neigh[1][0], z)])
          / (neigh[1][1] - neigh[1][0]) / mcdata.step[1];
    n.z = (float) (mcdata.data[INDEX(x, y, neigh[2][1])]
           - mcdata.data[INDEX(x, y, neigh[2][0])])
          / (neigh[2][1] - neigh[2][0]) / mcdata.step[2];

    return n;
}

/* interpolate points and calulate normals */
static void interpolate(mcdata_t mcdata, int px, int py, int pz, 
                        GR3_MC_DTYPE v1,
                        int qx,  int qy,  int qz,
                        GR3_MC_DTYPE v2,
                        gr3_coord_t *p, gr3_coord_t *n)
{
    double mu;
    gr3_coord_t n1, n2;
    double norm;

    if (ABS(mcdata.isolevel - v1) < 0.00001)
        mu = 0.0;
    else if (ABS(mcdata.isolevel - v2) < 0.00001)
        mu = 1.0;
    else if (ABS(v1 - v2) < 0.00001)
        mu = 0.5;
    else
        mu = 1.0 * (mcdata.isolevel - v1) / (v2 - v1);

    p->x = (px + mu * (qx - px)) * mcdata.step[0] + mcdata.offset[0];
    p->y = (py + mu * (qy - py)) * mcdata.step[1] + mcdata.offset[1];
    p->z = (pz + mu * (qz - pz)) * mcdata.step[2] + mcdata.offset[2];

    n1 = getgrad(mcdata, px, py, pz);
    n2 = getgrad(mcdata, qx, qy, qz);
    n->x = -(n1.x + mu * (n2.x - n1.x));
    n->y = -(n1.y + mu * (n2.y - n1.y));
    n->z = -(n1.z + mu * (n2.z - n1.z));

    norm = sqrt(n->x * n->x + n->y * n->y + n->z * n->z);
    if (norm > 0.0) {
        n->x /= norm;
        n->y /= norm;
        n->z /= norm;
    }
}

/*!
 * marching cubes algorithm for one x-layer.
 * created vertices are cached between calls using vindex.
 * vindex associates the intersected edge with the vertex index.
 * the edge is identified by its location (low, high), direction (x, y, z)
 * and coordinates (py, pz) of its starting point.
 * direction and location are the first index:
 * (x, y_low, z_low, y_high, z_high) (see mc_edgeprop) 
 * second index is py * mcdata.dim[1] + pz.
 * py and pz are the coordinates of the lower one of both edge vertices
 */
static void layer(mcdata_t mcdata, int x, int **vindex,
                  unsigned int *num_vertices, gr3_coord_t **vertices,
                  gr3_coord_t **normals, unsigned int *vertcapacity,
                  unsigned int *num_faces, unsigned int **indices,
                  unsigned int *facecapacity)
{
    int i, j;
    int y, z;
    int cubeindex;
    GR3_MC_DTYPE cubeval[8]; /* also cache between adjacent cubes */

    for (y = 0; y < mcdata.dim[1] - 1; y++) {
        /* init z-cache */
        for (i = 0; i < 4; i++) {
            int zi = mc_zvertices[0][i];

            cubeval[mc_zvertices[1][i]] = mcdata.data[INDEX(
                x + mc_cubeverts[zi][0],
                y + mc_cubeverts[zi][1],
                0 + mc_cubeverts[zi][2])];
        }
        for (z = 0; z < mcdata.dim[2] - 1; z++) {
            cubeindex = 0;
            /* shift old values (z-cache) */
            for (i = 0; i < 4; i++) {
                int zi = mc_zvertices[0][i];

                cubeval[zi] = cubeval[mc_zvertices[1][i]];
                if (cubeval[zi] < mcdata.isolevel) {
                    cubeindex |= 1 << zi;
                }
            }
            /* read new cube values */
            for (i = 0; i < 4; i++) {
                int zi = mc_zvertices[1][i];

                cubeval[zi] = mcdata.data[INDEX(x + mc_cubeverts[zi][0],
                    y + mc_cubeverts[zi][1], z + mc_cubeverts[zi][2])];
                if (cubeval[zi] < mcdata.isolevel) {
                    cubeindex |= 1 << zi;
                }
            }
            if (cubeindex != 0 && cubeindex != 255) {
                /* create triangles */
                for (i = 0; i < mc_tricount[cubeindex]; i++) {
                    if (*facecapacity <= *num_faces) {
                        (*facecapacity) = (unsigned int) 
                                          (*num_faces * 1.5) + 50;
                        *indices = realloc(*indices, (*facecapacity)
                                           * 3 * sizeof(int));
                    }
                    /* create triangle vertices */
                    for (j = 0; j < 3; j++) {
                        int trival = mc_tritable[cubeindex][i*3+j];
                        const int *edge = mc_cubeedges[trival];
                        int dir = mc_edgeprop[trival];
                        int px = x + mc_cubeverts[edge[0]][0];
                        int py = y + mc_cubeverts[edge[0]][1];
                        int pz = z + mc_cubeverts[edge[0]][2];
                        /* lookup if vertex already exists */
                        int node = vindex[dir][IDX2D(py, pz)];
                        if (node < 0) {
                            /* it does not, create it */
                            GR3_MC_DTYPE v1 = cubeval[edge[0]];
                            GR3_MC_DTYPE v2 = cubeval[edge[1]];
                            if (*vertcapacity <= *num_vertices) {
                                (*vertcapacity) = (unsigned int)
                                                  (*num_vertices * 1.5)
                                                  + 50;
                                *vertices =
                                    realloc(*vertices, (*vertcapacity)
                                            * sizeof(gr3_coord_t));
                                *normals =
                                    realloc(*normals, (*vertcapacity)
                                            * sizeof(gr3_coord_t));
                            }
                            node = *num_vertices;
                            interpolate(mcdata, px, py, pz, v1,
                                        x + mc_cubeverts[edge[1]][0],
                                        y + mc_cubeverts[edge[1]][1],
                                        z + mc_cubeverts[edge[1]][2], v2,
                                        *vertices + node,
                                        *normals + node);
                            vindex[dir][IDX2D(py, pz)] = node;
                            (*num_vertices)++;
                        }
                        /* add vertex index to the element array */
                        (*indices)[*num_faces * 3 + j] = node;
                    }
                    (*num_faces)++;
                }
            }
        }
    }
}

/*!
 * handle consecutive calls to layer
 */
static void layerblock(mcdata_t mcdata, int from, int to,
                       unsigned int *num_vertices, gr3_coord_t **vertices,
                       gr3_coord_t **normals, unsigned int *num_faces,
                       unsigned int **faces)
{
    int x;
    int y;
    int z;
    unsigned int vertcapacity;
    unsigned int facecapacity;
    /* cache for the vertex indices of the x-layer
     * [x, y_bot, z_bot, y_top, z_top] */
    int *vindex[5], *ntmp;

    *num_vertices = 0;
    vertcapacity = 0;
    *vertices = NULL;
    *normals = NULL;
    *num_faces = 0;
    facecapacity = 0;
    *faces = NULL;

    vindex[0] = malloc(5 * mcdata.dim[1] * mcdata.dim[2] * sizeof(int));
    vindex[1] = vindex[0] + mcdata.dim[1] * mcdata.dim[2];
    vindex[2] = vindex[0] + 2 * mcdata.dim[1] * mcdata.dim[2];
    vindex[3] = vindex[0] + 3 * mcdata.dim[1] * mcdata.dim[2];
    vindex[4] = vindex[0] + 4 * mcdata.dim[1] * mcdata.dim[2];

    for (y = 0; y < mcdata.dim[1]; y++) {
        for (z = 0; z < mcdata.dim[2]; z++) {
            vindex[0][IDX2D(y, z)] = -1;
            vindex[3][IDX2D(y, z)] = -1;
            vindex[4][IDX2D(y, z)] = -1;
        }
    }
    /*
     * iterate layer-by-layer through the data
     * create an indexed mesh
     * indices are cached in vindex[direction of the edge][y, z]
     * the top cache becomes the bottom in the next iterarion
     */
    for (x = from; x < to; x++) {
        ntmp = vindex[1];
        vindex[1] = vindex[3];
        vindex[3] = ntmp;
        ntmp = vindex[2];
        vindex[2] = vindex[4];
        vindex[4] = ntmp;
        for (y = 0; y < mcdata.dim[1]; y++) {
            for (z = 0; z < mcdata.dim[2]; z++) {
                vindex[0][IDX2D(y, z)] = -1;
                vindex[3][IDX2D(y, z)] = -1;
                vindex[4][IDX2D(y, z)] = -1;
            }
        }
        layer(mcdata, x, vindex,
              num_vertices, vertices, normals, &vertcapacity,
              num_faces, faces, &facecapacity);
    }
    free(vindex[0]);
}

/*!
 * Create an isosurface (as indexed mesh) from voxel data
 * with the marching cubes algorithm.
 * This function manages the parallelization:
 * Divide the data into blocks along the x-axis. Allocate memory,
 * call layerblock and merge the individual meshes into a single one.
 *
 * \param [in]  data          the volume (voxel) data
 * \param [in]  isolevel      value where the isosurface will be extracted
 * \param [in]  dim_x         number of elements in x-direction
 * \param [in]  dim_y         number of elements in y-direction
 * \param [in]  dim_z         number of elements in z-direction
 * \param [in]  stride_x      number of elements to step when traversing
 *                            the data in x-direction
 * \param [in]  stride_y      number of elements to step when traversing
 *                            the data in y-direction
 * \param [in]  stride_z      number of elements to step when traversing
 *                            the data in z-direction
 * \param [in]  step_x        distance between the voxels in x-direction
 * \param [in]  step_y        distance between the voxels in y-direction
 * \param [in]  step_z        distance between the voxels in z-direction
 * \param [in]  offset_x      coordinate origin
 * \param [in]  offset_y      coordinate origin
 * \param [in]  offset_z      coordinate origin
 * \param [out] num_vertices  number of vertices created
 * \param [out] vertices      array of vertex coordinates
 * \param [out] normals       array of vertex normal vectors
 * \param [out] num_indices   number of indices created
 *                            (3 times the number of triangles)
 * \param [out] indices       array of vertex indices that make the triangles
 */
GR3API void gr3_triangulateindexed(const GR3_MC_DTYPE *data,
                       GR3_MC_DTYPE isolevel,
                       unsigned int dim_x, unsigned int dim_y,
                       unsigned int dim_z, unsigned int stride_x,
                       unsigned int stride_y, unsigned int stride_z,
                       double step_x, double step_y, double step_z,
                       double offset_x, double offset_y, double offset_z,
                       unsigned int *num_vertices, gr3_coord_t **vertices,
                       gr3_coord_t **normals, unsigned int *num_indices,
                       unsigned int **indices)
{
    int num_threads;
    unsigned int num_faces;
    unsigned int *num_t_vertices, *num_t_faces, **t_faces;
    gr3_coord_t **t_vertices, **t_normals;
    unsigned int *vertblock, *faceblock;
    mcdata_t mcdata;
#if defined(_OPENMP) && defined(THREADLIMIT)
    int max_threads;

    max_threads = omp_get_max_threads();
    if (max_threads > THREADLIMIT)
        omp_set_num_threads(THREADLIMIT);
#endif

    if (stride_x == 0)
        stride_x = dim_z * dim_y;
    if (stride_y == 0)
        stride_y = dim_z;
    if (stride_z == 0)
        stride_z = 1;

    mcdata.data = data;
    mcdata.isolevel = isolevel;
    mcdata.dim[0] = dim_x;
    mcdata.dim[1] = dim_y;
    mcdata.dim[2] = dim_z;
    mcdata.stride[0] = stride_x;
    mcdata.stride[1] = stride_y;
    mcdata.stride[2] = stride_z;
    mcdata.step[0] = step_x;
    mcdata.step[1] = step_y;
    mcdata.step[2] = step_z;
    mcdata.offset[0] = offset_x;
    mcdata.offset[1] = offset_y;
    mcdata.offset[2] = offset_z;

    *num_vertices = 0;
    *vertices = NULL;
    *normals = NULL;
    *num_indices = 0;
    *indices = NULL;

#ifdef _OPENMP
#pragma omp parallel default(none) shared(num_threads, num_t_vertices, \
t_vertices, t_normals, num_t_faces, t_faces, mcdata, \
vertblock, faceblock, num_vertices, num_faces, vertices, normals, indices)
#endif
    {
        int thread_id;
        unsigned int from, to;
        unsigned int i;
#ifdef _OPENMP
#pragma omp single
#endif
        {
            /* allocate temporary memory for each thread */
#ifdef _OPENMP
            num_threads = omp_get_num_threads();
#else
            num_threads = 1;
#endif
            num_t_vertices = malloc(num_threads * sizeof(unsigned int));
            t_vertices = malloc(num_threads * sizeof(gr3_coord_t *));
            t_normals = malloc(num_threads * sizeof(gr3_coord_t *));
            num_t_faces = malloc(num_threads * sizeof(unsigned int));
            t_faces = malloc(num_threads * sizeof(unsigned int *));
        }
        /* create a mesh per thread */
#ifdef _OPENMP
        thread_id = omp_get_thread_num();
#else
        thread_id = 0;
#endif
        from = thread_id * (mcdata.dim[0] - 1) / num_threads;
        to = (thread_id + 1) * (mcdata.dim[0] - 1) / num_threads;
        num_t_vertices[thread_id] = 0;
        t_vertices[thread_id] = NULL;
        t_normals[thread_id] = NULL;
        num_t_faces[thread_id] = 0;
        t_faces[thread_id] = NULL;
        layerblock(mcdata, from, to,
                   num_t_vertices + thread_id,
                   t_vertices + thread_id, t_normals + thread_id,
                   num_t_faces + thread_id, t_faces + thread_id);
#ifdef _OPENMP
#pragma omp barrier
#pragma omp single
#endif
        {
            /* calculate beginning indices of thread blocks */
            vertblock = malloc((num_threads + 1) * sizeof(unsigned int));
            vertblock[0] = 0;
            faceblock = malloc((num_threads + 1) * sizeof(unsigned int));
            faceblock[0] = 0;
            for (i = 0; i < (unsigned int) num_threads; i++) {
                vertblock[i + 1] = vertblock[i] + num_t_vertices[i];
                faceblock[i + 1] = faceblock[i] + num_t_faces[i];
            }
            *num_vertices = vertblock[num_threads];
            num_faces = faceblock[num_threads];
            *vertices =
                realloc(*vertices, *num_vertices * sizeof(gr3_coord_t));
            *normals =
                realloc(*normals, *num_vertices * sizeof(gr3_coord_t));
            *indices =
                realloc(*indices, num_faces * 3 * sizeof(unsigned int));
        }
        /* copy thread meshes into the arrays */
        memcpy(*vertices + vertblock[thread_id],
               t_vertices[thread_id],
               num_t_vertices[thread_id] * sizeof(gr3_coord_t));
        memcpy(*normals + vertblock[thread_id],
               t_normals[thread_id],
               num_t_vertices[thread_id] * sizeof(gr3_coord_t));
        /* translate thread indices to global indices */
        for (i = 0; i < num_t_faces[thread_id]; i++) {
            (*indices)[(faceblock[thread_id] + i) * 3 + 0]
                = t_faces[thread_id][i * 3 + 0] + vertblock[thread_id];
            (*indices)[(faceblock[thread_id] + i) * 3 + 1]
                = t_faces[thread_id][i * 3 + 1] + vertblock[thread_id];
            (*indices)[(faceblock[thread_id] + i) * 3 + 2]
                = t_faces[thread_id][i * 3 + 2] + vertblock[thread_id];
        }
        free(t_vertices[thread_id]);
        free(t_normals[thread_id]);
        free(t_faces[thread_id]);
    }
    free(faceblock);
    free(vertblock);
    free(t_faces);
    free(num_t_faces);
    free(t_normals);
    free(t_vertices);
    free(num_t_vertices);
    *num_indices = num_faces * 3;
#if defined(_OPENMP) && defined(THREADLIMIT)
    omp_set_num_threads(max_threads);
#endif
}

/*!
 * Create an isosurface (as mesh) from voxel data
 * with the marching cubes algorithm.
 * This function calls gr3_triangulateindexed and copies the values.
 *
 * \param [in]  data          the volume (voxel) data
 * \param [in]  isolevel      value where the isosurface will be extracted
 * \param [in]  dim_x         number of elements in x-direction
 * \param [in]  dim_y         number of elements in y-direction
 * \param [in]  dim_z         number of elements in z-direction
 * \param [in]  stride_x      number of elements to step when traversing
 *                           the data in x-direction
 * \param [in]  stride_y      number of elements to step when traversing
 *                           the data in y-direction
 * \param [in]  stride_z      number of elements to step when traversing
 *                           the data in z-direction
 * \param [in]  step_x        distance between the voxels in x-direction
 * \param [in]  step_y        distance between the voxels in y-direction
 * \param [in]  step_z        distance between the voxels in z-direction
 * \param [in]  offset_x      coordinate origin
 * \param [in]  offset_y      coordinate origin
 * \param [in]  offset_z      coordinate origin
 * \param [out] triangles_p   array of triangle data
 *
 * \returns the number of triangles created
 */
GR3API unsigned int gr3_triangulate(const GR3_MC_DTYPE *data,
                GR3_MC_DTYPE isolevel,
                unsigned int dim_x, unsigned int dim_y, unsigned int dim_z,
                unsigned int stride_x, unsigned int stride_y,
                unsigned int stride_z, double step_x, double step_y,
                double step_z, double offset_x, double offset_y,
                double offset_z, gr3_triangle_t **triangles_p)
{
    unsigned int num_vertices;
    gr3_coord_t *vertices, *normals;
    unsigned int num_indices;
    unsigned int *indices;
    unsigned int i, j;
#if defined(_OPENMP) && defined(THREADLIMIT)
    int max_threads;

    max_threads = omp_get_max_threads();
    if (max_threads > THREADLIMIT)
        omp_set_num_threads(THREADLIMIT);
#endif

    gr3_triangulateindexed(data, isolevel,
                           dim_x, dim_y, dim_z,
                           stride_x, stride_y, stride_z,
                           step_x, step_y, step_z,
                           offset_x, offset_y, offset_z,
                           &num_vertices, &vertices, &normals,
                           &num_indices, &indices);

    *triangles_p = malloc(num_indices / 3 * sizeof(gr3_triangle_t));
#ifdef _OPENMP
#pragma omp parallel for default(none) private(j) \
shared(num_indices, triangles_p, indices, vertices, normals)
#endif
    for (i = 0; i < num_indices / 3; i++) {
        for (j = 0; j < 3; j++) {
            (*triangles_p)[i].vertex[j] = vertices[indices[i * 3 + j]];
            (*triangles_p)[i].normal[j] = normals[indices[i * 3 + j]];
        }
    }
    free(vertices);
    free(normals);
    free(indices);

#if defined(_OPENMP) && defined(THREADLIMIT)
    omp_set_num_threads(max_threads);
#endif
    return num_indices / 3;
}
