#include "VolumeCache.h"

#include <boost/thread/locks.hpp>
#include <boost/thread/mutex.hpp>

#include <openvdb/io/File.h>
#include <openvdb/math/Vec3.h>
#include <openvdb/tools/RayIntersector.h>

#include <nanovdb/NanoVDB.h>
#include <nanovdb/util/IO.h>
#include <nanovdb/util/Ray.h>
#include <nanovdb/util/HDDA.h>
#include <nanovdb/util/SampleFromVoxels.h>


CacheInstance::CacheInstance()
	: m_Valid(false)
{}

CacheInstance::~CacheInstance() {}

//OpenVDBCacheInstance::FloatGridSamplerPtr OpenVDBCacheInstance::_gridSampler;

OpenVDBCacheInstance::OpenVDBCacheInstance(VolumeCache::Handle h, openvdb::GridBase::Ptr grid)
	: m_Grid(grid), m_Handle(h)
{
	m_Valid = true;
	m_CacheType = kOpenVDB;
}
template <typename GridType>
auto OpenVDBCacheInstance::sample(int interpolation, const Imath::V3d& wsP) const -> typename GridType::ValueType
{
	using namespace openvdb::tools;
	using Result = GridType::ValueType;
	GridType::Ptr grid = openvdb::gridConstPtrCast<GridType>(m_Grid);
	auto accessor = grid->getAccessor();
	auto transform = m_Grid->transform();
	
	openvdb::Vec3d sample_point(wsP.x, wsP.y, wsP.z);
	auto index_point = m_Grid->worldToIndex(sample_point);
	Result result = openvdb::zeroVal<Result>();
	switch (interpolation)
	{
	case Interpolation::kClosest:
		PointSampler::sample(accessor, index_point, result);
		break;
	case Interpolation::kTriliear:
		BoxSampler::sample(accessor, index_point, result);
		break;
	case Interpolation::kTricubic:
		QuadraticSampler::sample(accessor, index_point, result);
		break;
	}
	return result
}

template<typename Intersector>
void OpenVDBCacheInstance::march_extents(const AtVolumeIntersectionInfo * info, const Ray & ray, const Intersector * master_intersector)
{
	Intersector intersector(*master_intersector);

	if (intersector.setWorldRay(ray))
	{
		Intersector::RealType t_min, t_max;
		while (intersector.march(t_min, t_max))
		{
			const auto near_pos = intersector.getWorldPos(t_min);
			const auto far_pos = intersector.getWorldPos(t_max);

			t_min = (near_pos - ray.eye()).length();
			t_max = (far_pos - ray.eye()).length();

			AiVolumeAddIntersection(info, t_min, t_max);
		}
	}
}

void OpenVDBCacheInstance::bounds(Imath::V3d& bbmin, Imath::V3d& bbmax) const
{
	openvdb::CoordBBox bbox = m_Grid->evalActiveVoxelBoundingBox();
	openvdb::Vec3d mn = bbox.min().asVec3d();
	openvdb::Vec3d mx = bbox.max().asVec3d();
	mn = m_Grid->indexToWorld(mn);
	mx = m_Grid->indexToWorld(mx);
	bbmin = Imath::V3d(mn.x(), mn.y(), mn.z());
	bbmax = Imath::V3d(mx.x(), mx.y(), mx.z());
}

size_t OpenVDBCacheInstance::memUsage() const
{
	return m_Grid->memUsage();
}

Imath::V3d OpenVDBCacheInstance::voxelSize() const
{
	openvdb::Vec3d v = m_Grid->voxelSize();
	return Imath::V3d(v.x(), v.y(), v.z());
}

VolumeCache::StaticInit VolumeCache::m_StaticInit;
VolumeCache* VolumeCache::m_Instance(0);
VolumeCache::StaticInit::StaticInit()
{
	openvdb::initialize();
}

VolumeCache* VolumeCache::instance()
{
	if (!m_Instance)
	{
		boost::mutex::scoped_lock m;
		if (!m_Instance)
		{
			m_Instance = new VolumeCache;
		}
	}
	return m_Instance;
}

CacheInstance::Ptr VolumeCache::getCacheInstance(Handle h)
{
	CacheInstance::Ptr c;
	if (h >= 0)
	{
		CacheInstanceMap::iterator it = m_CacheInstanceMap.find(h);
		if (it != m_CacheInstanceMap.end())
		{
			c = it->second;
		}
	}

	return c;
}

VolumeCache::Handle VolumeCache::readField(const std::string& filename, const std::string& fieldName, bool openToNanoVDB = false)
{
	Handle h;

	boost::mutex::scoped_lock lock;

	// check in our cache entries to see if we've loaded this cache already
	std::string fileHandle = filename + std::string("/") + fieldName;
	FileHandleMap::const_iterator it = m_FileHandleMap.find(fileHandle);
	if (it != m_FileHandleMap.end())
	{
		// just return the existing handle
		h = it->second;
	}
	else
	{
		// try and load it
		// first figure out what type of file we have
		// std::string ext = cacheName.substr(cacheName.length() - 3, std::string::npos);
		std::string ext = filename.substr(filename.find_last_of(".") + 1);

		if (ext == "vdb")
		{
			h = readOpenVDBGrid(filename, fieldName, false);
			m_FileHandleMap[fileHandle] = h;
		}
		else if (ext == "nvdb")
		{
			h = readNanoVDBGrid(filename, fieldName, false);
			m_FileHandleMap[fileHandle] = h;
		}
		else if (ext == "f3d")
		{

		}
		else
		{
			std::cerr << "[VolumeCache] Unsupported file type with extension '" << ext << "'" << std::endl;
		}
	}

	return h;
}

void VolumeCache::evict(Handle h)
{
	CacheInstanceMap::iterator it = m_CacheInstanceMap.find(h);
	if (it != m_CacheInstanceMap.end()) m_CacheInstanceMap.erase(it);
}

size_t VolumeCache::memUsage() const
{
	CacheInstanceMap::const_iterator it = m_CacheInstanceMap.begin();
	size_t mem = 0;
	for (; it != m_CacheInstanceMap.end(); ++it)
	{
		mem += it->second->memUsage();
	}
	return mem;
}

VolumeCache::Handle VolumeCache::readOpenVDBGrid(const std::string& filename, const std::string& gridName, bool opentoNanoVDB = false)
{
	VolumeCache::Handle h;
	try
	{
		openvdb::io::File file(filename);
		file.open();
		openvdb::GridBase::Ptr gridBase = file.readGrid(gridName);
		file.close();
		if (gridBase)
		{
			h = m_NumEntries++;
			OpenVDBCacheInstance::Ptr oc(new OpenVDBCacheInstance(h, gridBase));
			m_CacheInstanceMap[h] = oc;
		}
		else
		{
			h = -1;
		}
	}
	catch (openvdb::Exception& e)
	{
		std::cerr << "[VolumeCache] " << e.what() << std::endl;
		h = kInvalidFile;
	}

	return h;
}

VolumeCache::Handle VolumeCache::readNanoVDBGrid(const std::string & filename, const std::string & gridName, bool openToOpenVDB)
{
	VolumeCache::Handle h;
	
	auto gridHandle = nanovdb::io::readGrid<>(filename);
	
	if (gridHandle)
	{
		h = m_NumEntries++;
		NanoVDBCacheInstance::Ptr nc(new NanoVDBCacheInstance(h, std::move(gridHandle)));
		m_CacheInstanceMap[h] = nc;
	}
	else
	{
		h = -1;
	}
	
	return h;
}

NanoVDBCacheInstance::NanoVDBCacheInstance(VolumeCache::Handle h, nanovdb::GridHandle<> gridHandle)
	:m_Handle(h), m_GridHandle(&gridHandle)
{
}

bool NanoVDBCacheInstance::march_extents(const AtVolumeIntersectionInfo* info, const Imath::V3f &ray_origin, const Imath::V3f &ray_direction, float t0, float t1)
{
	nanovdb::Vec3f ray_eye(ray_origin.x, ray_origin.y, ray_origin.z);
	nanovdb::Vec3f ray_dir(ray_direction.x, ray_direction.y, ray_direction.z);
	const nanovdb::Ray<float> ray(ray_eye, ray_dir, t0, t1);

	const auto grid = m_GridHandle->grid<float>();
	const auto& tree = grid->tree();
	auto grid_accessor = tree.getAccessor();
	auto index_ray = ray.worldToIndexF(*grid);

	nanovdb::TreeMarcher<nanovdb::NanoTree<float>::LeafNodeType, nanovdb::Ray<float>, decltype(grid_accessor)> marcher(grid_accessor);

	if (!marcher.init(index_ray))
		return false;

	const typename nanovdb::NanoTree<float>::LeafNodeType* node = nullptr;
	float t_min, t_max;
	while (marcher.step(&node, t_min, t_max))
	{
		const auto near_pos = index_ray.indexToWorldF(*grid)(t_min);
		const auto far_pos = index_ray.indexToWorldF(*grid)(t_max);
		t_min = (near_pos - ray.eye()).length();
		t_max = (far_pos - ray.eye()).length();
		AiVolumeAddIntersection(info, t_min, t_max);
	}
}

void NanoVDBCacheInstance::bounds(Imath::V3d & bbmin, Imath::V3d & bbmax) const
{
	auto worldBBox = m_GridHandle->gridMetaData()->worldBBox();
	bbmin.x = worldBBox.min()[0];
	bbmin.y = worldBBox.min()[1];
	bbmin.z = worldBBox.min()[2];
	bbmax.x = worldBBox.max()[0];
	bbmax.y = worldBBox.max()[1];
	bbmax.z = worldBBox.max()[2];
}

size_t NanoVDBCacheInstance::memUsage() const
{
	size_t memUsage;
	
	switch (m_GridHandle->gridMetaData()->gridType())
	{
	case nanovdb::GridType::Float:
		memUsage = m_GridHandle->grid<float>()->memUsage();
		break;
	case nanovdb::GridType::Vec3f:
		memUsage = m_GridHandle->grid<nanovdb::Vec3f>()->memUsage();
		break;
	}
	return memUsage;
}

Imath::V3d NanoVDBCacheInstance::voxelSize() const
{
	auto voxelSize = m_GridHandle->gridMetaData()->voxelSize();
	return Imath::V3d(voxelSize[0], voxelSize[1], voxelSize[2]);
}

template <typename GridType>
auto NanoVDBCacheInstance::sample(int interpolation, const Imath::V3d& wsP) const -> typename GridType
{
	auto grid = m_GridHandle->grid<GridType>();
	auto& grid_tree = grid->tree();
	auto grid_accessor = grid_tree.getAccessor();

	nanovdb::Vec3f sample_point(wsP.x, wsP.y, wsP.z);
	// Transform sample point from object space to world space
	auto world_point = grid->map().applyInverseMap(sample_point);

	auto index_point = grid->worldToIndexF(world_point);

	GridType result;

	switch (interpolation)
	{
	case Interpolation::kClosest:
	{
		auto sampler = nanovdb::createSampler<0, decltype(grid_accessor), false>(grid_accessor);
		result = sampler(index_point);
		break;
	}
	case Interpolation::kTriliear:
	{
		auto sampler = nanovdb::createSampler<1, decltype(grid_accessor), false>(grid_accessor);
		result = sampler(index_point);
		break;
	}
	case Interpolation::kTricubic:
		
		auto sampler = nanovdb::createSampler<2, decltype(grid_accessor), false>(grid_accessor);
		result = sampler(index_point);
		break;
	}
	return result;
}