/*
 * Copyright 2011-2016 Blender Foundation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "render/mesh.h"
#include "render/mesh_volume.h"
#include "render/attribute.h"
#ifdef WITH_OPENVDB
#include "render/openvdb.h"
#endif
#include "render/scene.h"

#include "util/util_foreach.h"
#include "util/util_logging.h"
#include "util/util_progress.h"
#include "util/util_sparse_grid.h"

CCL_NAMESPACE_BEGIN

enum {
	QUAD_X_MIN = 0,
	QUAD_X_MAX = 1,
	QUAD_Y_MIN = 2,
	QUAD_Y_MAX = 3,
	QUAD_Z_MIN = 4,
	QUAD_Z_MAX = 5,
};

const int quads_indices[6][4] = {
	/* QUAD_X_MIN */
	{ 4, 0, 3, 7 },
	/* QUAD_X_MAX */
	{ 1, 5, 6, 2 },
	/* QUAD_Y_MIN */
	{ 4, 5, 1, 0 },
	/* QUAD_Y_MAX */
	{ 3, 2, 6, 7 },
	/* QUAD_Z_MIN */
	{ 0, 1, 2, 3 },
	/* QUAD_Z_MAX */
	{ 5, 4, 7, 6 },
};

const float3 quads_normals[6] = {
	/* QUAD_X_MIN */
	make_float3(-1.0f, 0.0f, 0.0f),
	/* QUAD_X_MAX */
	make_float3(1.0f, 0.0f, 0.0f),
	/* QUAD_Y_MIN */
	make_float3(0.0f, -1.0f, 0.0f),
	/* QUAD_Y_MAX */
	make_float3(0.0f, 1.0f, 0.0f),
	/* QUAD_Z_MIN */
	make_float3(0.0f, 0.0f, -1.0f),
	/* QUAD_Z_MAX */
	make_float3(0.0f, 0.0f, 1.0f),
};

static void create_quad(int3 corners[8], vector<int3> &vertices, vector<QuadData> &quads, int face_index)
{
	size_t vertex_offset = vertices.size();

	QuadData quad;
	quad.v0 = vertex_offset + 0;
	quad.v1 = vertex_offset + 1;
	quad.v2 = vertex_offset + 2;
	quad.v3 = vertex_offset + 3;
	quad.normal = quads_normals[face_index];

	quads.push_back(quad);

	vertices.push_back(corners[quads_indices[face_index][0]]);
	vertices.push_back(corners[quads_indices[face_index][1]]);
	vertices.push_back(corners[quads_indices[face_index][2]]);
	vertices.push_back(corners[quads_indices[face_index][3]]);
}

static const int CUBE_SIZE = 8;

VolumeMeshBuilder::VolumeMeshBuilder(VolumeParams *volume_params)
{
	params = volume_params;
	number_of_nodes = 0;

	const size_t x = divide_up(params->resolution.x, CUBE_SIZE);
	const size_t y = divide_up(params->resolution.y, CUBE_SIZE);
	const size_t z = divide_up(params->resolution.z, CUBE_SIZE);

	/* Adding 2*pad_size since we pad in both positive and negative directions
	 * along the axis. */
	const size_t px = divide_up(params->resolution.x + 2*params->pad_size, CUBE_SIZE);
	const size_t py = divide_up(params->resolution.y + 2*params->pad_size, CUBE_SIZE);
	const size_t pz = divide_up(params->resolution.z + 2*params->pad_size, CUBE_SIZE);

	res = make_int3(px, py, pz);
	pad_offset = make_int3(px - x, py - y, pz - z);

	grid.resize(px*py*pz, 0);
}

void VolumeMeshBuilder::add_node(int x, int y, int z)
{
	/* Map coordinates to index space. */
	const int index_x = (x/CUBE_SIZE) + pad_offset.x;
	const int index_y = (y/CUBE_SIZE) + pad_offset.y;
	const int index_z = (z/CUBE_SIZE) + pad_offset.z;

	assert((index_x >= 0) && (index_y >= 0) && (index_z >= 0));

	const size_t index = compute_index(index_x, index_y, index_z, res);

	/* We already have a node here. */
	if(grid[index] == 1) {
		return;
	}

	++number_of_nodes;

	grid[index] = 1;
}

void VolumeMeshBuilder::add_node_with_padding(int x, int y, int z)
{
	for(int px = x - params->pad_size; px < x + params->pad_size; ++px) {
		for(int py = y - params->pad_size; py < y + params->pad_size; ++py) {
			for(int pz = z - params->pad_size; pz < z + params->pad_size; ++pz) {
				add_node(px, py, pz);
			}
		}
	}
}

void VolumeMeshBuilder::create_mesh(vector<float3> &vertices,
                                    vector<int> &indices,
                                    vector<float3> &face_normals)
{
	/* We create vertices in index space (is), and only convert them to object
	 * space when done. */
	vector<int3> vertices_is;
	vector<QuadData> quads;

	generate_vertices_and_quads(vertices_is, quads);

	deduplicate_vertices(vertices_is, quads);

	convert_object_space(vertices_is, vertices);

	convert_quads_to_tris(quads, indices, face_normals);
}

void VolumeMeshBuilder::generate_vertices_and_quads(
		vector<ccl::int3> &vertices_is,
		vector<QuadData> &quads)
{
	/* Overallocation, we could count the number of quads and vertices to create
	 * in a pre-pass if memory becomes an issue. */
	vertices_is.reserve(number_of_nodes*8);
	quads.reserve(number_of_nodes*6);

	for(int z = 0; z < res.z; ++z) {
		for(int y = 0; y < res.y; ++y) {
			for(int x = 0; x < res.x; ++x) {
				size_t voxel_index = compute_index(x, y, z, res);
				if(grid[voxel_index] == 0) {
					continue;
				}

				/* Compute min and max coords of the node in index space. */
				int3 min = make_int3((x - pad_offset.x)*CUBE_SIZE,
				                     (y - pad_offset.y)*CUBE_SIZE,
				                     (z - pad_offset.z)*CUBE_SIZE);

				/* Maximum is just CUBE_SIZE voxels away from minimum on each axis. */
				int3 max = make_int3(min.x + CUBE_SIZE, min.y + CUBE_SIZE, min.z + CUBE_SIZE);

				int3 corners[8] = {
					make_int3(min[0], min[1], min[2]),
					make_int3(max[0], min[1], min[2]),
					make_int3(max[0], max[1], min[2]),
					make_int3(min[0], max[1], min[2]),
					make_int3(min[0], min[1], max[2]),
					make_int3(max[0], min[1], max[2]),
					make_int3(max[0], max[1], max[2]),
					make_int3(min[0], max[1], max[2]),
				};

				/* Only create a quad if on the border between an active and
				 * an inactive node.
				 */

				voxel_index = compute_index(x - 1, y, z, res);
				if(voxel_index == -1 || grid[voxel_index] == 0) {
					create_quad(corners, vertices_is, quads, QUAD_X_MIN);
				}

				voxel_index = compute_index(x + 1, y, z, res);
				if(voxel_index == -1 || grid[voxel_index] == 0) {
					create_quad(corners, vertices_is, quads, QUAD_X_MAX);
				}

				voxel_index = compute_index(x, y - 1, z, res);
				if(voxel_index == -1 || grid[voxel_index] == 0) {
					create_quad(corners, vertices_is, quads, QUAD_Y_MIN);
				}

				voxel_index = compute_index(x, y + 1, z, res);
				if(voxel_index == -1 || grid[voxel_index] == 0) {
					create_quad(corners, vertices_is, quads, QUAD_Y_MAX);
				}

				voxel_index = compute_index(x, y, z - 1, res);
				if(voxel_index == -1 || grid[voxel_index] == 0) {
					create_quad(corners, vertices_is, quads, QUAD_Z_MIN);
				}

				voxel_index = compute_index(x, y, z + 1, res);
				if(voxel_index == -1 || grid[voxel_index] == 0) {
					create_quad(corners, vertices_is, quads, QUAD_Z_MAX);
				}
			}
		}
	}
}

void VolumeMeshBuilder::deduplicate_vertices(vector<int3> &vertices,
                                             vector<QuadData> &quads)
{
	vector<int3> sorted_vertices = vertices;
	std::sort(sorted_vertices.begin(), sorted_vertices.end());
	vector<int3>::iterator it = std::unique(sorted_vertices.begin(), sorted_vertices.end());
	sorted_vertices.resize(std::distance(sorted_vertices.begin(), it));

	vector<QuadData> new_quads = quads;

	for(size_t i = 0; i < vertices.size(); ++i) {
		for(size_t j = 0; j < sorted_vertices.size(); ++j) {
			if(vertices[i] != sorted_vertices[j]) {
				continue;
			}

			for(int k = 0; k < quads.size(); ++k) {
				if(quads[k].v0 == i) {
					new_quads[k].v0 = j;
				}
				else if(quads[k].v1 == i) {
					new_quads[k].v1 = j;
				}
				else if(quads[k].v2 == i) {
					new_quads[k].v2 = j;
				}
				else if(quads[k].v3 == i) {
					new_quads[k].v3 = j;
				}
			}

			break;
		}
	}

	vertices = sorted_vertices;
	quads = new_quads;
}

void VolumeMeshBuilder::convert_object_space(const vector<int3> &vertices,
	                                         vector<float3> &out_vertices)
{
	out_vertices.reserve(vertices.size());

	for(size_t i = 0; i < vertices.size(); ++i) {
		float3 vertex = make_float3(vertices[i].x, vertices[i].y, vertices[i].z);
		vertex *= params->cell_size;
		vertex += params->start_point;

		out_vertices.push_back(vertex);
	}
}

void VolumeMeshBuilder::convert_quads_to_tris(const vector<QuadData> &quads,
                                              vector<int> &tris,
                                              vector<float3> &face_normals)
{
	int index_offset = 0;
	tris.resize(quads.size()*6);
	face_normals.reserve(quads.size()*2);

	for(size_t i = 0; i < quads.size(); ++i) {
		tris[index_offset++] = quads[i].v0;
		tris[index_offset++] = quads[i].v2;
		tris[index_offset++] = quads[i].v1;

		face_normals.push_back(quads[i].normal);

		tris[index_offset++] = quads[i].v0;
		tris[index_offset++] = quads[i].v3;
		tris[index_offset++] = quads[i].v2;

		face_normals.push_back(quads[i].normal);
	}
}

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

struct VoxelAttributeGrid {
	void *data;
	int *grid_info;
	int channels;
};

void MeshManager::create_volume_mesh(Scene *scene,
                                     Mesh *mesh,
                                     Progress& progress)
{
	string msg = string_printf("Computing Volume Mesh %s", mesh->name.c_str());
	progress.set_status("Updating Mesh", msg);

	vector<VoxelAttributeGrid> voxel_grids;

	/* Compute volume parameters. */
	VolumeParams volume_params;
	volume_params.resolution = make_int3(0, 0, 0);
	bool is_openvdb = false;
	int grid_mem = -1;

	foreach(Attribute& attr, mesh->attributes.attributes) {
		if(attr.element != ATTR_ELEMENT_VOXEL) {
			continue;
		}

		VoxelAttribute *voxel = attr.data_voxel();
		device_memory *image_memory = scene->image_manager->image_memory(voxel->slot);
		device_memory *grid_info = image_memory->grid_info;
		int3 resolution = make_int3(image_memory->real_width,
		                            image_memory->real_height,
		                            image_memory->real_depth);

		if(volume_params.resolution == make_int3(0, 0, 0)) {
			/* First volume grid. */
			volume_params.resolution = resolution;
			grid_mem = image_memory->memory_size();
		}
		else if(volume_params.resolution != resolution) {
			VLOG(1) << "Can't create volume mesh, all voxel grid resolutions must be equal\n";
			return;
		}

		is_openvdb = (image_memory->grid_type == IMAGE_GRID_TYPE_OPENVDB);

		VoxelAttributeGrid voxel_grid;
		voxel_grid.data = image_memory->host_pointer;
		voxel_grid.channels = image_memory->data_elements;
		voxel_grid.grid_info = grid_info ? static_cast<int*>(grid_info->host_pointer) : NULL;
		voxel_grids.push_back(voxel_grid);
	}

	if(voxel_grids.empty()) {
		return;
	}

	/* Compute padding. */
	Shader *volume_shader = NULL;
	int pad_size = 0;

	foreach(Shader *shader, mesh->used_shaders) {
		if(!shader->has_volume) {
			continue;
		}

		volume_shader = shader;

		if(shader->volume_interpolation_method == VOLUME_INTERPOLATION_LINEAR) {
			pad_size = max(1, pad_size);
		}
		else if(shader->volume_interpolation_method == VOLUME_INTERPOLATION_CUBIC) {
			pad_size = max(2, pad_size);
		}

		break;
	}

	if(!volume_shader) {
		return;
	}

	/* Compute start point and cell size from transform. */
	Attribute *attr = mesh->attributes.find(ATTR_STD_GENERATED_TRANSFORM);
	const int3 resolution = volume_params.resolution;
	float3 start_point = make_float3(0.0f, 0.0f, 0.0f);
	float3 cell_size = make_float3(1.0f/resolution.x,
	                               1.0f/resolution.y,
	                               1.0f/resolution.z);
	const int3 tiled_res = make_int3(get_tile_res(resolution.x),
	                                 get_tile_res(resolution.y),
	                                 get_tile_res(resolution.z));
	const bool using_cuda = (scene->device->info.type == DEVICE_CUDA);

	if(attr) {
		const Transform *tfm = attr->data_transform();
		const Transform itfm = transform_inverse(*tfm);
		start_point = transform_point(&itfm, start_point);
		cell_size = transform_direction(&itfm, cell_size);
	}

	volume_params.start_point = start_point;
	volume_params.cell_size = cell_size;
	volume_params.pad_size = pad_size;

	/* Build bounding mesh around non-empty volume cells. */
	VolumeMeshBuilder builder(&volume_params);
	const float isovalue = mesh->volume_isovalue;

	for(size_t i = 0; i < voxel_grids.size(); ++i) {
		const VoxelAttributeGrid &voxel_grid = voxel_grids[i];
		const int channels = voxel_grid.channels;

		if(is_openvdb) {
#ifdef WITH_OPENVDB
			openvdb_build_mesh(&builder, voxel_grid.data, resolution, isovalue, channels > 1);
#else
			assert(0);
#endif
		}
		else {
			const int *grid_info = voxel_grid.grid_info;
			float *data = static_cast<float*>(voxel_grid.data);

			for(int z = 0; z < resolution.z; ++z) {
				for(int y = 0; y < resolution.y; ++y) {
					for(int x = 0; x < resolution.x; ++x) {
						int voxel_index;

						if(grid_info) {
							voxel_index = using_cuda ?
							    compute_index_cuda(grid_info, x, y, z,
							                       resolution.x, resolution.y, resolution.z,
							                       tiled_res.x, tiled_res.y, tiled_res.z) :
							    compute_index(grid_info, x, y, z,
							                  resolution.x, resolution.y, resolution.z);

							if(voxel_index < 0) {
								continue;
							}
						}
						else {
							voxel_index = compute_index(x, y, z, resolution);
						}

						voxel_index *= channels;
						for(int c = 0; c < channels; c++) {
							if(data[voxel_index + c] >= isovalue) {
								builder.add_node_with_padding(x, y, z);
								break;
							}
						}
					}
				}
			}
		}
	}

	/* Create mesh. */
	vector<float3> vertices;
	vector<int> indices;
	vector<float3> face_normals;
	builder.create_mesh(vertices, indices, face_normals);

	mesh->clear(true);
	mesh->reserve_mesh(vertices.size(), indices.size()/3);
	mesh->used_shaders.push_back(volume_shader);

	for(size_t i = 0; i < vertices.size(); ++i) {
		mesh->add_vertex(vertices[i]);
	}

	for(size_t i = 0; i < indices.size(); i += 3) {
		mesh->add_triangle(indices[i], indices[i + 1], indices[i + 2], 0, false);
	}

	Attribute *attr_fN = mesh->attributes.add(ATTR_STD_FACE_NORMAL);
	float3 *fN = attr_fN->data_float3();

	for(size_t i = 0; i < face_normals.size(); ++i) {
		fN[i] = face_normals[i];
	}

	/* Print stats. */
	VLOG(1) << "Memory usage volume mesh: "
	        << ((vertices.size() + face_normals.size())*sizeof(float3) + indices.size()*sizeof(int))/(1024.0*1024.0)
	        << "Mb.";

	VLOG(1) << "Memory usage volume grid: " << grid_mem / (1024.0 * 1024.0) << "Mb.";
}

CCL_NAMESPACE_END