/* Copyright (c) 2023-2025, Qualcomm Innovation Center, Inc. All rights reserved.
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 "mobile_nerf.h"
#include "filesystem/legacy.h"
#include "glm/gtx/matrix_decompose.hpp"
#include "gltf_loader.h"
#include "platform/platform.h"
#include "rendering/subpasses/forward_subpass.h"
#include "scene_graph/components/material.h"
#include "scene_graph/components/mesh.h"
#include "scene_graph/components/perspective_camera.h"

namespace
{
constexpr uint32_t MIN_THREAD_COUNT = 1;

template <typename T>
struct CopyBuffer
{
	std::vector<T> operator()(std::unordered_map<std::string, vkb::core::BufferC> &buffers, const char *buffer_name)
	{
		auto iter = buffers.find(buffer_name);
		if (iter == buffers.cend())
		{
			return {};
		}
		auto &buffer = iter->second;

		std::vector<T> out;

		const size_t sz = buffer.get_size();
		out.resize(sz / sizeof(T));
		const bool already_mapped = buffer.get_data() != nullptr;
		if (!already_mapped)
		{
			buffer.map();
		}
		memcpy(&out[0], buffer.get_data(), sz);
		if (!already_mapped)
		{
			buffer.unmap();
		}
		return out;
	}
};

void camera_set_look_at(vkb::Camera &camera, const glm::vec3 look, const glm::vec3 up)
{
	auto view_matrix = glm::lookAt(camera.position, look, up);

	glm::vec3 scale;
	glm::quat orientation;
	glm::vec3 translation;
	glm::vec3 skew;
	glm::vec4 perspective;
	glm::decompose(view_matrix, scale, orientation, translation, skew, perspective);

	camera.set_rotation(glm::eulerAngles(orientation) * glm::pi<float>() / 180.f);
	camera.set_position(translation);
}

}        // namespace

MobileNerf::MobileNerf()
{
	title = "Mobile NeRF";
	// SPIRV 1.4 requires Vulkan 1.1
	set_api_version(VK_API_VERSION_1_1);
	add_device_extension(VK_KHR_SPIRV_1_4_EXTENSION_NAME);
	// Required by VK_KHR_spirv_1_4
	add_device_extension(VK_KHR_SHADER_FLOAT_CONTROLS_EXTENSION_NAME);
	// For choosing different sets of weights
	add_device_extension(VK_EXT_DESCRIPTOR_INDEXING_EXTENSION_NAME);
}

MobileNerf::~MobileNerf()
{
	if (has_device())
	{
		if (render_pass_nerf)
		{
			vkDestroyRenderPass(get_device().get_handle(), render_pass_nerf, nullptr);
		}

		for (uint32_t i = 0; i < nerf_framebuffers.size(); i++)
		{
			if (nerf_framebuffers[i])
			{
				vkDestroyFramebuffer(get_device().get_handle(), nerf_framebuffers[i], nullptr);
			}
		}

		auto device_ptr = get_device().get_handle();

		for (auto &model : models)
		{
			model.vertex_buffer.reset();
			model.index_buffer.reset();

			vkDestroySampler(get_device().get_handle(), model.texture_input_0.sampler, nullptr);
			vkDestroySampler(get_device().get_handle(), model.texture_input_1.sampler, nullptr);

			vkDestroyPipeline(device_ptr, model.pipeline_first_pass, nullptr);
		}

		for (auto &weights_buffer : weights_buffers)
		{
			weights_buffer.reset();
		}

		for (auto &uniform_buffer : uniform_buffers)
		{
			uniform_buffer.reset();
		}

		vkDestroyPipelineLayout(device_ptr, pipeline_first_pass_layout, nullptr);
		vkDestroyDescriptorSetLayout(device_ptr, descriptor_set_first_pass_layout, nullptr);

		if (pipeline_baseline)
		{
			vkDestroyPipeline(get_device().get_handle(), pipeline_baseline, nullptr);
			vkDestroyPipelineLayout(get_device().get_handle(), pipeline_layout_baseline, nullptr);
			vkDestroyDescriptorSetLayout(get_device().get_handle(), descriptor_set_layout_baseline, nullptr);
		}

		for (auto &attachment : frameAttachments)
		{
			attachment.feature_0.destroy();
			attachment.feature_1.destroy();
			attachment.feature_2.destroy();
			attachment.weights_idx.destroy();
		}
	}
}

void MobileNerf::read_json_map()
{
	std::string assetBase = vkb::fs::path::get(vkb::fs::path::Type::Assets);
	LOGI("Base assets path: {}", assetBase);

#if defined(NERF_JSON_FILE)
	const std::string nerf_obj_map = assetBase + "scenes/mobile_nerf_models.json";

	std::ifstream f(nerf_obj_map);

	if (!f)
	{
		LOGE("Failed to open nerf obj map data");
		assert(0);
	}

	LOGI("Parsing nerf obj map data {}", nerf_obj_map);

	json raw_asset_map = json::parse(f);
#else

	const std::string nerf_obj_json =
	    R"V0G0N(
        {
            "width": 0,

            "height": 0,

            "texture_type": "8bit",

            "target_model": "lego_combo",

            "deferred": false,

            "rotation": true,

            "lego_ball":{
                "path": "scenes/morpheus_team/lego_ball_phone/",
                "num_sub_model": 1,
                "original": false,
                "camera": [-1, 1, 1],
                "instancing":{
                    "dim": [1, 1, 1],
                    "interval": [2.0, 2.0, 2.0]
                }
            },

            "lego_boba_fett":{
                "path": "scenes/morpheus_team/lego_boba_fett_phone/",
                "num_sub_model": 1,
                "original": false,
                "camera": [-1, 1, 1],
                "instancing":{
                    "dim": [1, 1, 1],
                    "interval": [2.0, 2.0, 2.0]
                }
            },

            "lego_monster_truck":{
                "path": "scenes/morpheus_team/lego_monster_truck_phone/",
                "num_sub_model": 1,
                "original": false,
                "camera": [-1, 1, 1],
                "instancing":{
                    "dim": [1, 1, 1],
                    "interval": [2.0, 2.0, 2.0]
                }
            },

            "lego_tractor":{
                "path": "scenes/morpheus_team/lego_tractor_phone/",
                "num_sub_model": 1,
                "original": false,
                "camera": [-1, 1, 1],
                "instancing":{
                    "dim": [1, 1, 1],
                    "interval": [2.0, 2.0, 2.0]
                }
            },

            "lego_combo":{
                "combo": true,
                "models": ["scenes/morpheus_team/lego_ball_phone/", "scenes/morpheus_team/lego_boba_fett_phone/",
                            "scenes/morpheus_team/lego_monster_truck_phone/", "scenes/morpheus_team/lego_tractor_phone/"],
                "original": [false, false, false, false],
                "camera": [-0.0381453, 1.84186, -1.51744],
                "instancing":{
                    "dim": [2, 2, 2],
                    "interval": [1.5, 1.5, 1.5]
                }
            }
        }
        )V0G0N";

	json raw_asset_map = json::parse(nerf_obj_json);

#endif

	std::string target_model = raw_asset_map["target_model"].get<std::string>();
	asset_map                = raw_asset_map[target_model];

	// Load combo models or a single model
	if (!asset_map["combo"].is_null())
	{
		combo_mode = asset_map["combo"].get<bool>();
	}
	else
	{
		combo_mode = false;
	}

	if (combo_mode)
	{
		model_path.resize(asset_map["models"].size());
		using_original_nerf_models.resize(asset_map["models"].size());

		for (int i = 0; i < model_path.size(); i++)
		{
			model_path[i]                 = asset_map["models"][i].get<std::string>();
			using_original_nerf_models[i] = asset_map["original"][i].get<bool>();
			LOGI("Target model: {}, asset path: {}", target_model, model_path[i]);
		}
	}
	else
	{
		model_path.resize(1);
		model_path[0] = asset_map["path"].get<std::string>();
		using_original_nerf_models.resize(1);
		using_original_nerf_models[0] = asset_map["original"].get<bool>();
		LOGI("Target model: {}, asset path: {}", target_model, model_path[0]);
	}

	std::string textureType = raw_asset_map["texture_type"].get<std::string>();

	if (textureType == "8bit")
	{
		LOGI("Using VK_FORMAT_R8G8B8A8_UNORM for feature texture");
		feature_map_format = VK_FORMAT_R8G8B8A8_UNORM;
	}
	else if (textureType == "16bit")
	{
		LOGI("Using VK_FORMAT_R16G16B16A16_SFLOAT for feature texture");
		feature_map_format = VK_FORMAT_R16G16B16A16_SFLOAT;
	}
	else if (textureType == "32bit")
	{
		LOGI("Using VK_FORMAT_R32G32B32A32_SFLOAT for feature texture");
		feature_map_format = VK_FORMAT_R32G32B32A32_SFLOAT;
	}
	else if (textureType == "8bit")
	{
		LOGI("Using VK_FORMAT_R8G8B8A8_UNORM for feature texture");
		feature_map_format = VK_FORMAT_R8G8B8A8_UNORM;
	}
	else
	{
		LOGW("Unrecognized feature texture type, using VK_FORMAT_R32G32B32A32_SFLOAT");
		feature_map_format = VK_FORMAT_R32G32B32A32_SFLOAT;
	}
	use_deferred = raw_asset_map["deferred"].get<bool>();
	do_rotation  = raw_asset_map["rotation"].get<bool>();

	view_port_width  = raw_asset_map["width"].get<int>();
	view_port_height = raw_asset_map["height"].get<int>();

	if (asset_map["camera"].is_array() && asset_map["camera"].size() == 3)
	{
		camera_pos = glm::vec3(asset_map["camera"][0].get<float>(), asset_map["camera"][1].get<float>(), asset_map["camera"][2].get<float>());
	}
	else
	{
		LOGW("Fail to read camera position. Use defualt value.");
	}

	json instacing_map = asset_map["instancing"];
	if (instacing_map["dim"].is_array() && instacing_map["dim"].size() == 3)
	{
		instancing_info.dim = glm::vec3(instacing_map["dim"][0].get<int>(), instacing_map["dim"][1].get<int>(), instacing_map["dim"][2].get<int>());
	}
	else
	{
		LOGE("Wrong instancing dimension. Terminating...");
		exit(1);
	}

	if (instacing_map["interval"].is_array() && instacing_map["interval"].size() == 3)
	{
		instancing_info.interval = glm::vec3(instacing_map["interval"][0].get<float>(), instacing_map["interval"][1].get<float>(), instacing_map["interval"][2].get<float>());
	}
	else
	{
		LOGE("Wrong instancing interval. Terminating...");
		exit(1);
	}

	if (instancing_info.dim.x <= 0 || instancing_info.dim.y <= 0 || instancing_info.dim.z <= 0 || instancing_info.interval.x <= 0.f || instancing_info.interval.y <= 0.f || instancing_info.interval.z <= 0.f)
	{
		LOGE("Instancing settings must be positive. Terminating...");
		exit(1);
	}
}

void MobileNerf::load_shaders()
{
	// Loading first pass shaders
	if (use_deferred)
	{
		// Loading first pass shaders
		shader_stages_first_pass[0] = load_shader("mobile_nerf/raster.vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
		shader_stages_first_pass[1] = load_shader(
		    combo_mode ?
		        (using_original_nerf_models[0] ? "mobile_nerf/raster_combo.frag.spv" : "mobile_nerf/raster_morpheus_combo.frag.spv") :
		        (using_original_nerf_models[0] ? "mobile_nerf/raster.frag.spv" : "mobile_nerf/raster_morpheus.frag.spv"),
		    VK_SHADER_STAGE_FRAGMENT_BIT);

		// Loading second pass shaders
		shader_stages_second_pass[0] = load_shader("mobile_nerf/quad.vert", VK_SHADER_STAGE_VERTEX_BIT);
		shader_stages_second_pass[1] = load_shader(
		    combo_mode ?
		        (using_original_nerf_models[0] ? "mobile_nerf/mlp_combo.frag.spv" : "mobile_nerf/mlp_morpheus_combo.frag.spv") :
		        (using_original_nerf_models[0] ? "mobile_nerf/mlp.frag.spv" : "mobile_nerf/mlp_morpheus.frag.spv"),
		    VK_SHADER_STAGE_FRAGMENT_BIT);
	}
	else
	{
		// Loading one pass shaders
		shader_stages_first_pass[0] = load_shader("mobile_nerf/raster.vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
		shader_stages_first_pass[1] = load_shader(
		    using_original_nerf_models[0] ? "mobile_nerf/merged.frag.spv" : "mobile_nerf/merged_morpheus.frag.spv",
		    VK_SHADER_STAGE_FRAGMENT_BIT);
	}
}

bool MobileNerf::prepare(const vkb::ApplicationOptions &options)
{
	read_json_map();

	// Load the mlp for each model
	mlp_weight_vector.resize(model_path.size());

	for (int i = 0; i < model_path.size(); i++)
	{
		initialize_mlp_uniform_buffers(i);
	}

	if (!ApiVulkanSample::prepare(options))
	{
		return false;
	}

	if (view_port_width == 0 || view_port_height == 0)
	{
		view_port_width        = width;
		view_port_height       = height;
		use_native_screen_size = true;
	}

	load_shaders();

	if (use_deferred)
	{
		update_render_pass_nerf_baseline();
	}
	else
	{
		update_render_pass_nerf_forward();
	}

	setup_nerf_framebuffer_baseline();
	// Because we have our own customized render pass, the UI render pass need to be updated with load on load so it won't
	// clear out the written color attachment
	update_render_pass_flags(RenderPassCreateFlags::ColorAttachmentLoad);

	camera.type  = vkb::CameraType::LookAt;
	camera_pos.y = -camera_pos.y;        // flip y to keep consistency of the init pos between rayquery and rasterization
	camera.set_position(camera_pos);
	camera_set_look_at(camera, glm::vec3(0.0f), glm::vec3(0.0f, 1.0f, 0.0f));

	camera.set_perspective(60.0f, static_cast<float>(width) / static_cast<float>(height), 0.01f, 256.0f);

	int models_entry = 0;

	for (int model_index = 0; model_index < model_path.size(); model_index++)
	{
		int num_sub_model = models[models_entry].sub_model_num;

		for (int sub_model_index = 0; sub_model_index < num_sub_model; sub_model_index++)
		{
			load_scene(model_index, sub_model_index, models_entry);
			create_texture(model_index, sub_model_index, models_entry);
			create_static_object_buffers(model_index, sub_model_index, models_entry);
			models_entry++;
		}
	}
	create_uniforms();
	prepare_instance_data();
	create_pipeline_layout_fist_pass();

	if (use_deferred)
	{
		create_pipeline_layout_baseline();
	}
	create_descriptor_pool();

	for (auto &model : models)
	{
		create_descriptor_sets_first_pass(model);
	}

	if (use_deferred)
	{
		create_descriptor_sets_baseline();
	}
	prepare_pipelines();
	build_command_buffers();

	prepared = true;
	LOGI("Prepare Done!");
	return true;
}

bool MobileNerf::resize(const uint32_t width, const uint32_t height)
{
	ApiVulkanSample::resize(width, height);
	rebuild_command_buffers();
	return true;
}

void MobileNerf::request_gpu_features(vkb::PhysicalDevice &gpu)
{
	REQUEST_REQUIRED_FEATURE(gpu, VkPhysicalDeviceDescriptorIndexingFeaturesEXT, VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES_EXT, shaderUniformBufferArrayNonUniformIndexing);
	REQUEST_REQUIRED_FEATURE(gpu, VkPhysicalDeviceDescriptorIndexingFeaturesEXT, VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES_EXT, runtimeDescriptorArray);
	REQUEST_REQUIRED_FEATURE(gpu, VkPhysicalDeviceDescriptorIndexingFeaturesEXT, VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES_EXT, descriptorBindingVariableDescriptorCount);
}

void MobileNerf::render(float delta_time)
{
	if (!prepared)
	{
		return;
	}
	draw();
	update_uniform_buffers();
}

void MobileNerf::FrameBufferAttachment::destroy()
{
	if (!image)
	{
		return;
	}

	auto &device = image->get_device();
	vkDestroyImageView(device.get_handle(), view, nullptr);
	image.reset();
}

void MobileNerf::setup_attachment(VkFormat format, VkImageUsageFlags usage, FrameBufferAttachment &attachment)
{
	if (attachment)
	{
		attachment.destroy();
	}

	VkImageAspectFlags aspectMask  = VK_IMAGE_ASPECT_COLOR_BIT;
	VkImageLayout      imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	if (usage & VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT)
	{
		aspectMask  = VK_IMAGE_ASPECT_DEPTH_BIT;
		imageLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
	}

	auto surfaceExtent = get_render_context().get_surface_extent();
	attachment.image   = std::make_unique<vkb::core::Image>(
        get_device(),
        VkExtent3D{surfaceExtent.width, surfaceExtent.height, 1},
        format,
        usage,
        VMA_MEMORY_USAGE_GPU_ONLY);

	with_command_buffer([&](VkCommandBuffer command_buffer) {
		vkb::image_layout_transition(command_buffer, attachment.image->get_handle(),
		                             VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
		                             VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
		                             {},
		                             {},
		                             VK_IMAGE_LAYOUT_UNDEFINED,
		                             VK_IMAGE_LAYOUT_GENERAL,
		                             {aspectMask, 0, 1, 0, 1});
	});

	VkImageViewCreateInfo color_image_view           = vkb::initializers::image_view_create_info();
	color_image_view.viewType                        = VK_IMAGE_VIEW_TYPE_2D;
	color_image_view.format                          = format;
	color_image_view.subresourceRange.aspectMask     = aspectMask;
	color_image_view.subresourceRange.baseMipLevel   = 0;
	color_image_view.subresourceRange.levelCount     = 1;
	color_image_view.subresourceRange.baseArrayLayer = 0;
	color_image_view.subresourceRange.layerCount     = 1;
	color_image_view.image                           = attachment.image->get_handle();
	VK_CHECK(vkCreateImageView(get_device().get_handle(), &color_image_view, nullptr, &attachment.view));
}

void MobileNerf::setup_nerf_framebuffer_baseline()
{
	if (use_deferred)
	{
		frameAttachments.resize(get_render_context().get_render_frames().size());

		for (auto i = 0; i < frameAttachments.size(); i++)
		{
			setup_attachment(feature_map_format, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].feature_0);
			setup_attachment(feature_map_format, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].feature_1);
			setup_attachment(VK_FORMAT_R16G16B16A16_SFLOAT, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].feature_2);
			if (combo_mode)
				setup_attachment(VK_FORMAT_R8_UINT, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_INPUT_ATTACHMENT_BIT, frameAttachments[i].weights_idx);
		}
	}

	// Delete existing frame buffers
	if (nerf_framebuffers.size() > 0)
	{
		for (uint32_t i = 0; i < nerf_framebuffers.size(); i++)
		{
			if (nerf_framebuffers[i] != VK_NULL_HANDLE)
			{
				vkDestroyFramebuffer(get_device().get_handle(), nerf_framebuffers[i], nullptr);
			}
		}
	}

	std::vector<VkImageView> views;

	if (use_deferred)
	{
		views.resize(combo_mode ? 6 : 5);
		views[depth_attach_idx] = depth_stencil.view;
	}
	else
	{
		views.resize(2);
		views[0] = depth_stencil.view;
	}

	// Depth/Stencil attachment is the same for all frame buffers

	VkFramebufferCreateInfo framebuffer_create_info = {};
	framebuffer_create_info.sType                   = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
	framebuffer_create_info.pNext                   = NULL;
	framebuffer_create_info.renderPass              = render_pass_nerf;
	framebuffer_create_info.attachmentCount         = static_cast<uint32_t>(views.size());
	framebuffer_create_info.pAttachments            = views.data();
	framebuffer_create_info.width                   = get_render_context().get_surface_extent().width;
	framebuffer_create_info.height                  = get_render_context().get_surface_extent().height;
	framebuffer_create_info.layers                  = 1;

	nerf_framebuffers.resize(swapchain_buffers.size());

	for (uint32_t i = 0; i < nerf_framebuffers.size(); i++)
	{
		if (use_deferred)
		{
			views[color_attach_0_idx] = frameAttachments[i].feature_0.view;
			views[color_attach_1_idx] = frameAttachments[i].feature_1.view;
			views[color_attach_2_idx] = frameAttachments[i].feature_2.view;
			if (combo_mode)
				views[color_attach_3_idx] = frameAttachments[i].weights_idx.view;
			views[swapchain_attach_idx] = swapchain_buffers[i].view;
		}
		else
		{
			views[1] = swapchain_buffers[i].view;
		}

		VK_CHECK(vkCreateFramebuffer(get_device().get_handle(), &framebuffer_create_info, nullptr, &nerf_framebuffers[i]));
	}
}

void MobileNerf::update_descriptor_sets_baseline()
{
	for (int i = 0; i < nerf_framebuffers.size(); i++)
	{
		std::vector<VkDescriptorImageInfo> attachment_input_descriptors;
		attachment_input_descriptors.resize(combo_mode ? 4 : 3);

		attachment_input_descriptors[0].sampler     = VK_NULL_HANDLE;
		attachment_input_descriptors[0].imageView   = frameAttachments[i].feature_0.view;
		attachment_input_descriptors[0].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		attachment_input_descriptors[1].sampler     = VK_NULL_HANDLE;
		attachment_input_descriptors[1].imageView   = frameAttachments[i].feature_1.view;
		attachment_input_descriptors[1].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		attachment_input_descriptors[2].sampler     = VK_NULL_HANDLE;
		attachment_input_descriptors[2].imageView   = frameAttachments[i].feature_2.view;
		attachment_input_descriptors[2].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		if (combo_mode)
		{
			attachment_input_descriptors[3].sampler     = VK_NULL_HANDLE;
			attachment_input_descriptors[3].imageView   = frameAttachments[i].weights_idx.view;
			attachment_input_descriptors[3].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
		}

		VkWriteDescriptorSet texture_input_write_0 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 0, &attachment_input_descriptors[0]);
		VkWriteDescriptorSet texture_input_write_1 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1, &attachment_input_descriptors[1]);
		VkWriteDescriptorSet texture_input_write_2 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 2, &attachment_input_descriptors[2]);

		if (combo_mode)
		{
			VkWriteDescriptorSet texture_input_write_3 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 3, &attachment_input_descriptors[3]);

			std::vector<VkWriteDescriptorSet> write_descriptor_sets = {
			    texture_input_write_0,
			    texture_input_write_1,
			    texture_input_write_2,
			    texture_input_write_3};

			vkUpdateDescriptorSets(get_device().get_handle(), static_cast<uint32_t>(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE);
		}
		else
		{
			std::vector<VkWriteDescriptorSet> write_descriptor_sets = {
			    texture_input_write_0,
			    texture_input_write_1,
			    texture_input_write_2};

			vkUpdateDescriptorSets(get_device().get_handle(), static_cast<uint32_t>(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE);
		}
	}
}

void MobileNerf::build_command_buffers()
{
	if (use_native_screen_size)
	{
		view_port_height = height;
		view_port_width  = width;
	}
	build_command_buffers_baseline();
}

void MobileNerf::build_command_buffers_baseline()
{
	// In case the screen is resized, need to update the storage image size and descriptor set
	// Note that the texture_rendered image has already been recreated at this point
	if (!prepared)
	{
		setup_nerf_framebuffer_baseline();

		if (use_deferred)
		{
			update_descriptor_sets_baseline();
		}
	}

	VkCommandBufferBeginInfo  command_buffer_begin_info = vkb::initializers::command_buffer_begin_info();
	std::vector<VkClearValue> clear_values;

	if (use_deferred)
	{
		if (combo_mode)
		{
			clear_values.resize(6);
			clear_values[0].color        = {{0.025f, 0.025f, 0.025f, 0.5f}};        // default_clear_color;
			clear_values[1].color        = {{0.025f, 0.025f, 0.025f, 0.5f}};        // default_clear_color;
			clear_values[2].color        = {{0.025f, 0.025f, 0.025f, 0.5f}};        // default_clear_color;
			clear_values[3].color        = {{0, 0, 0, 0}};                          // default clear index for weights;
			clear_values[4].depthStencil = {1.0f, 0};
			clear_values[5].color        = {{1.0f, 1.0f, 1.0f, 0.5f}};        // default_clear_color;
		}
		else
		{
			clear_values.resize(5);
			clear_values[0].color        = {{0.025f, 0.025f, 0.025f, 0.5f}};        // default_clear_color;
			clear_values[1].color        = {{0.025f, 0.025f, 0.025f, 0.5f}};        // default_clear_color;
			clear_values[2].color        = {{0.025f, 0.025f, 0.025f, 0.5f}};        // default_clear_color;
			clear_values[3].depthStencil = {1.0f, 0};
			clear_values[4].color        = {{1.0f, 1.0f, 1.0f, 0.5f}};        // default_clear_color;
		}
	}
	else
	{
		clear_values.resize(2);
		clear_values[0].depthStencil = {1.0f, 0};
		clear_values[1].color        = {{0.0f, 0.0f, 0.0f, 1.0f}};        // let's use this to distinguish forward rendering and deferred renderding
	}

	VkRenderPassBeginInfo render_pass_begin_info    = vkb::initializers::render_pass_begin_info();
	render_pass_begin_info.renderPass               = render_pass_nerf;
	render_pass_begin_info.renderArea.offset.x      = 0;
	render_pass_begin_info.renderArea.offset.y      = 0;
	render_pass_begin_info.renderArea.extent.width  = width;
	render_pass_begin_info.renderArea.extent.height = height;
	render_pass_begin_info.clearValueCount          = static_cast<uint32_t>(clear_values.size());
	render_pass_begin_info.pClearValues             = clear_values.data();

	VkClearValue clear_values_UI[2];
	clear_values_UI[0].color        = default_clear_color;
	clear_values_UI[1].depthStencil = {1.0f, 0};

	VkRenderPassBeginInfo render_pass_begin_info_UI    = vkb::initializers::render_pass_begin_info();
	render_pass_begin_info_UI.renderPass               = render_pass;
	render_pass_begin_info_UI.renderArea.offset.x      = 0;
	render_pass_begin_info_UI.renderArea.offset.y      = 0;
	render_pass_begin_info_UI.renderArea.extent.width  = width;
	render_pass_begin_info_UI.renderArea.extent.height = height;
	render_pass_begin_info_UI.clearValueCount          = 2;
	render_pass_begin_info_UI.pClearValues             = clear_values_UI;

	VkImageSubresourceRange subresource_range = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};

	for (size_t i = 0; i < draw_cmd_buffers.size(); ++i)
	{
		render_pass_begin_info.framebuffer = nerf_framebuffers[i];

		VK_CHECK(vkBeginCommandBuffer(draw_cmd_buffers[i], &command_buffer_begin_info));

		vkCmdBeginRenderPass(draw_cmd_buffers[i], &render_pass_begin_info, VK_SUBPASS_CONTENTS_INLINE);

		// First sub pass
		// Fills the attachments

		VkViewport viewport = vkb::initializers::viewport(static_cast<float>(width), static_cast<float>(height), 0.0f, 1.0f);
		const auto scissor  = vkb::initializers::rect2D(static_cast<int32_t>(width), static_cast<int32_t>(height), 0, 0);
		vkCmdSetViewport(draw_cmd_buffers[i], 0, 1, &viewport);
		vkCmdSetScissor(draw_cmd_buffers[i], 0, 1, &scissor);

		auto &ii = instancing_info;
		for (auto &model : models)
		{
			vkCmdBindPipeline(draw_cmd_buffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, model.pipeline_first_pass);
			// If deferred, only use the first descriptor bounded with the model
			// If forward, each model has the swapchan number of descriptor
			int descriptorIndex = use_deferred ? 0 : static_cast<uint32_t>(i);
			vkCmdBindDescriptorSets(draw_cmd_buffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_first_pass_layout,
			                        0, 1, &model.descriptor_set_first_pass[descriptorIndex], 0, nullptr);
			VkDeviceSize offsets[1] = {0};
			vkCmdBindVertexBuffers(draw_cmd_buffers[i], 0, 1, model.vertex_buffer->get(), offsets);
			vkCmdBindVertexBuffers(draw_cmd_buffers[i], 1, 1, instance_buffer->get(), offsets);
			vkCmdBindIndexBuffer(draw_cmd_buffers[i], model.index_buffer->get_handle(), 0, VK_INDEX_TYPE_UINT32);
			if (use_deferred && combo_mode)
			{
				PushConstants constants = {static_cast<unsigned int>(model.model_index)};
				vkCmdPushConstants(
				    draw_cmd_buffers[i],
				    pipeline_first_pass_layout,
				    VK_SHADER_STAGE_FRAGMENT_BIT,
				    0,
				    sizeof(PushConstants),
				    &constants);
			}
			vkCmdDrawIndexed(draw_cmd_buffers[i], static_cast<uint32_t>(model.indices.size()) * 3, ii.dim.x * ii.dim.y * ii.dim.z, 0, 0, 0);
		}

		if (use_deferred)
		{
			// Second sub pass
			// Render a full screen quad, reading from the previously written attachments via input attachments

			vkCmdNextSubpass(draw_cmd_buffers[i], VK_SUBPASS_CONTENTS_INLINE);
			vkCmdBindPipeline(draw_cmd_buffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_baseline);
			vkCmdBindDescriptorSets(draw_cmd_buffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_layout_baseline, 0, 1, &descriptor_set_baseline[i], 0, NULL);
			vkCmdDraw(draw_cmd_buffers[i], 3, 1, 0, 0);

			vkCmdEndRenderPass(draw_cmd_buffers[i]);
		}
		else
		{
			vkCmdEndRenderPass(draw_cmd_buffers[i]);
		}

		// Render UI

		render_pass_begin_info_UI.framebuffer = framebuffers[i];

		vkCmdBeginRenderPass(draw_cmd_buffers[i], &render_pass_begin_info_UI, VK_SUBPASS_CONTENTS_INLINE);
		draw_ui(draw_cmd_buffers[i]);
		vkCmdEndRenderPass(draw_cmd_buffers[i]);

		VK_CHECK(vkEndCommandBuffer(draw_cmd_buffers[i]));
	}
}

void MobileNerf::load_scene(int model_index, int sub_model_index, int models_entry)
{
	Model &model = models[models_entry];

	vkb::GLTFLoader loader{get_device()};
	int             total_sub_sub_model = using_original_nerf_models[model_index] ? 8 : 1;

	for (int sub_model = 0; sub_model < total_sub_sub_model; sub_model++)
	{
		std::string inputfile(model_path[model_index] + "shape" + std::to_string(sub_model_index));

		if (total_sub_sub_model > 1)
		{
			inputfile += ("_" + std::to_string(sub_model) + ".gltf");
		}
		else
		{
			inputfile += (".gltf");
		}

		LOGI("Parsing nerf obj {}", inputfile);

		auto scene = loader.read_scene_from_file(inputfile);

		for (auto &&mesh : scene->get_components<vkb::sg::Mesh>())
		{
			for (auto &&sub_mesh : mesh->get_submeshes())
			{
				auto       pts_               = CopyBuffer<glm::vec3>{}(sub_mesh->vertex_buffers, "position");
				const auto texcoord_          = CopyBuffer<glm::vec2>{}(sub_mesh->vertex_buffers, "texcoord_0");
				const auto vertex_start_index = static_cast<uint32_t>(model.vertices.size());

				// Copy vertex data
				{
					model.vertices.resize(vertex_start_index + pts_.size());
					for (size_t i = 0; i < pts_.size(); ++i)
					{
						model.vertices[vertex_start_index + i].position  = pts_[i];
						model.vertices[vertex_start_index + i].tex_coord = glm::vec2(texcoord_[i].x, 1.0f - texcoord_[i].y);
					}
				}

				// Copy index data
				{
					auto index_buffer_ = sub_mesh->index_buffer.get();
					if (index_buffer_)
					{
						assert(sub_mesh->index_type == VkIndexType::VK_INDEX_TYPE_UINT32);
						const size_t sz                   = index_buffer_->get_size();
						const size_t nTriangles           = sz / sizeof(uint32_t) / 3;
						const auto   triangle_start_index = static_cast<uint32_t>(model.indices.size());
						model.indices.resize(triangle_start_index + nTriangles);
						auto ptr = index_buffer_->get_data();
						assert(!!ptr);
						std::vector<uint32_t> tempBuffer(nTriangles * 3);
						memcpy(&tempBuffer[0], ptr, sz);
						for (size_t i = 0; i < nTriangles; ++i)
						{
							model.indices[triangle_start_index + i] = {vertex_start_index + static_cast<uint32_t>(tempBuffer[3 * i]),
							                                           vertex_start_index + static_cast<uint32_t>(tempBuffer[3 * i + 1]),
							                                           vertex_start_index + static_cast<uint32_t>(tempBuffer[3 * i + 2])};
						}
					}
				}
			}
		}
	}
}

void MobileNerf::create_descriptor_pool()
{
	if (use_deferred)
	{
		std::vector<VkDescriptorPoolSize> pool_sizes = {
		    // First Pass
		    {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 2 * static_cast<uint32_t>(models.size())},
		    {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast<uint32_t>(models.size())},
		};
		// Second Pass
		if (combo_mode)
		{
			pool_sizes.push_back({VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 4 * static_cast<uint32_t>(framebuffers.size())});
			pool_sizes.push_back({VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast<uint32_t>(framebuffers.size()) * static_cast<uint32_t>(model_path.size())});
		}
		else
		{
			pool_sizes.push_back({VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 3 * static_cast<uint32_t>(framebuffers.size())});
			pool_sizes.push_back({VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast<uint32_t>(framebuffers.size())});
		}

		VkDescriptorPoolCreateInfo descriptor_pool_create_info = vkb::initializers::descriptor_pool_create_info(pool_sizes, static_cast<uint32_t>(models.size() + framebuffers.size()));
		VK_CHECK(vkCreateDescriptorPool(get_device().get_handle(), &descriptor_pool_create_info, nullptr, &descriptor_pool));
	}
	else
	{
		std::vector<VkDescriptorPoolSize> pool_sizes = {
		    // First Pass
		    {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 2 * static_cast<uint32_t>(models.size()) * static_cast<uint32_t>(framebuffers.size())},
		    {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast<uint32_t>(models.size()) * static_cast<uint32_t>(framebuffers.size())},
		    {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 * static_cast<uint32_t>(models.size()) * static_cast<uint32_t>(framebuffers.size())}};

		VkDescriptorPoolCreateInfo descriptor_pool_create_info = vkb::initializers::descriptor_pool_create_info(pool_sizes, static_cast<uint32_t>(models.size() * framebuffers.size()));
		VK_CHECK(vkCreateDescriptorPool(get_device().get_handle(), &descriptor_pool_create_info, nullptr, &descriptor_pool));
	}
}

void MobileNerf::create_pipeline_layout_fist_pass()
{
	// First Pass Descriptor set and layout

	std::vector<VkDescriptorSetLayoutBinding> set_layout_bindings = {
	    vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_FRAGMENT_BIT, 0),
	    vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_FRAGMENT_BIT, 1),
	    vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_VERTEX_BIT, 2)};

	// If use forward, add uniform buffer descriptor for the weights
	if (!use_deferred)
	{
		set_layout_bindings.push_back(vkb::initializers::descriptor_set_layout_binding(
		    VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_FRAGMENT_BIT, 3));
	}

	VkDescriptorSetLayoutCreateInfo descriptor_layout = vkb::initializers::descriptor_set_layout_create_info(set_layout_bindings.data(), static_cast<uint32_t>(set_layout_bindings.size()));
	VK_CHECK(vkCreateDescriptorSetLayout(get_device().get_handle(), &descriptor_layout, nullptr, &descriptor_set_first_pass_layout));

	VkPipelineLayoutCreateInfo pipeline_layout_create_info =
	    vkb::initializers::pipeline_layout_create_info(
	        &descriptor_set_first_pass_layout,
	        1);

	if (use_deferred && combo_mode)
	{
		VkPushConstantRange pushConstantRange = vkb::initializers::push_constant_range(VK_SHADER_STAGE_FRAGMENT_BIT, sizeof(PushConstants), 0);

		pipeline_layout_create_info.pushConstantRangeCount = 1;
		pipeline_layout_create_info.pPushConstantRanges    = &pushConstantRange;
		VK_CHECK(vkCreatePipelineLayout(get_device().get_handle(), &pipeline_layout_create_info, nullptr, &pipeline_first_pass_layout));
	}
	else
	{
		VK_CHECK(vkCreatePipelineLayout(get_device().get_handle(), &pipeline_layout_create_info, nullptr, &pipeline_first_pass_layout));
	}
}

void MobileNerf::create_pipeline_layout_baseline()
{
	// Second Pass Descriptor set and layout

	std::vector<VkDescriptorSetLayoutBinding> set_layout_bindings = {
	    // Two output color from the first pass and ray direction
	    vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 0),
	    vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 1),
	    vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 2),
	};
	if (combo_mode)
	{
		set_layout_bindings.push_back(vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, VK_SHADER_STAGE_FRAGMENT_BIT, 3));
		// MLP weights array, using descriptor indexing
		set_layout_bindings.push_back(vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_FRAGMENT_BIT, 4, static_cast<uint32_t>(model_path.size())));
	}
	else
	{
		// MLP weights
		set_layout_bindings.push_back(vkb::initializers::descriptor_set_layout_binding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_FRAGMENT_BIT, 3));
	}

	VkDescriptorSetLayoutCreateInfo descriptor_layout = vkb::initializers::descriptor_set_layout_create_info(set_layout_bindings.data(), static_cast<uint32_t>(set_layout_bindings.size()));
	if (combo_mode)
	{
		VkDescriptorBindingFlagsEXT                    flags[5] = {0, 0, 0, 0, VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT_EXT};
		VkDescriptorSetLayoutBindingFlagsCreateInfoEXT setLayoutBindingFlags{};
		setLayoutBindingFlags.sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO_EXT;
		setLayoutBindingFlags.bindingCount  = 5;
		setLayoutBindingFlags.pBindingFlags = flags;
		descriptor_layout.pNext             = &setLayoutBindingFlags;
		VK_CHECK(vkCreateDescriptorSetLayout(get_device().get_handle(), &descriptor_layout, nullptr, &descriptor_set_layout_baseline));
	}
	else
	{
		VK_CHECK(vkCreateDescriptorSetLayout(get_device().get_handle(), &descriptor_layout, nullptr, &descriptor_set_layout_baseline));
	}

	VkPipelineLayoutCreateInfo pipeline_layout_create_info =
	    vkb::initializers::pipeline_layout_create_info(
	        &descriptor_set_layout_baseline,
	        1);

	VK_CHECK(vkCreatePipelineLayout(get_device().get_handle(), &pipeline_layout_create_info, nullptr, &pipeline_layout_baseline));
}

void MobileNerf::create_descriptor_sets_first_pass(Model &model)
{
	int numDescriptorPerModel = use_deferred ? 1 : static_cast<int>(nerf_framebuffers.size());
	model.descriptor_set_first_pass.resize(numDescriptorPerModel);

	for (int i = 0; i < numDescriptorPerModel; i++)
	{
		VkDescriptorSetAllocateInfo descriptor_set_allocate_info =
		    vkb::initializers::descriptor_set_allocate_info(descriptor_pool, &descriptor_set_first_pass_layout, 1);
		VK_CHECK(vkAllocateDescriptorSets(get_device().get_handle(), &descriptor_set_allocate_info, &model.descriptor_set_first_pass[i]));

		std::array<VkDescriptorImageInfo, 2> texture_input_descriptors;

		texture_input_descriptors[0].sampler     = model.texture_input_0.sampler;
		texture_input_descriptors[0].imageView   = model.texture_input_0.image->get_vk_image_view().get_handle();
		texture_input_descriptors[0].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		texture_input_descriptors[1].sampler     = model.texture_input_1.sampler;
		texture_input_descriptors[1].imageView   = model.texture_input_1.image->get_vk_image_view().get_handle();
		texture_input_descriptors[1].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		VkDescriptorBufferInfo buffer_descriptor = create_descriptor(*uniform_buffers[model.model_index]);

		VkWriteDescriptorSet texture_input_write_0 = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i],
		                                                                                     VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 0, &texture_input_descriptors[0]);
		VkWriteDescriptorSet texture_input_write_1 = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i],
		                                                                                     VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1, &texture_input_descriptors[1]);
		VkWriteDescriptorSet uniform_buffer_write  = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i],
		                                                                                     VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 2, &buffer_descriptor);

		std::vector<VkWriteDescriptorSet> write_descriptor_sets = {
		    texture_input_write_0,
		    texture_input_write_1,
		    uniform_buffer_write};

		VkDescriptorBufferInfo weights_buffer_descriptor;

		if (!use_deferred)
		{
			// Add in descriptor sets for MLP weights
			weights_buffer_descriptor = create_descriptor(*weights_buffers[model.model_index]);
			// Add in descriptor sets for MLP weights
			VkWriteDescriptorSet weights_buffer_write = vkb::initializers::write_descriptor_set(model.descriptor_set_first_pass[i],
			                                                                                    VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 3, &weights_buffer_descriptor);
			write_descriptor_sets.push_back(weights_buffer_write);
		}

		vkUpdateDescriptorSets(get_device().get_handle(), static_cast<uint32_t>(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE);
	}
}

void MobileNerf::create_descriptor_sets_baseline()
{
	descriptor_set_baseline.resize(nerf_framebuffers.size());

	for (int i = 0; i < nerf_framebuffers.size(); i++)
	{
		VkDescriptorSetAllocateInfo descriptor_set_allocate_info = vkb::initializers::descriptor_set_allocate_info(descriptor_pool, &descriptor_set_layout_baseline, 1);
		if (combo_mode)
		{
			uint32_t counts[1];
			counts[0] = static_cast<uint32_t>(model_path.size());

			VkDescriptorSetVariableDescriptorCountAllocateInfo set_counts = {};
			set_counts.sType                                              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_ALLOCATE_INFO;
			set_counts.descriptorSetCount                                 = 1;
			set_counts.pDescriptorCounts                                  = counts;

			descriptor_set_allocate_info.pNext = &set_counts;

			VK_CHECK(vkAllocateDescriptorSets(get_device().get_handle(), &descriptor_set_allocate_info, &descriptor_set_baseline[i]));
		}
		else
		{
			VK_CHECK(vkAllocateDescriptorSets(get_device().get_handle(), &descriptor_set_allocate_info, &descriptor_set_baseline[i]));
		}

		std::vector<VkDescriptorImageInfo> attachment_input_descriptors;
		attachment_input_descriptors.resize(combo_mode ? 4 : 3);

		attachment_input_descriptors[0].sampler     = VK_NULL_HANDLE;
		attachment_input_descriptors[0].imageView   = frameAttachments[i].feature_0.view;
		attachment_input_descriptors[0].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		attachment_input_descriptors[1].sampler     = VK_NULL_HANDLE;
		attachment_input_descriptors[1].imageView   = frameAttachments[i].feature_1.view;
		attachment_input_descriptors[1].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		attachment_input_descriptors[2].sampler     = VK_NULL_HANDLE;
		attachment_input_descriptors[2].imageView   = frameAttachments[i].feature_2.view;
		attachment_input_descriptors[2].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

		VkWriteDescriptorSet texture_input_write_0 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 0, &attachment_input_descriptors[0]);
		VkWriteDescriptorSet texture_input_write_1 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1, &attachment_input_descriptors[1]);
		VkWriteDescriptorSet texture_input_write_2 = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 2, &attachment_input_descriptors[2]);

		if (combo_mode)
		{
			attachment_input_descriptors[3].sampler     = VK_NULL_HANDLE;
			attachment_input_descriptors[3].imageView   = frameAttachments[i].weights_idx.view;
			attachment_input_descriptors[3].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
			VkWriteDescriptorSet texture_input_write_3  = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 3, &attachment_input_descriptors[3]);

			std::vector<VkDescriptorBufferInfo> weights_buffer_descriptors;
			weights_buffer_descriptors.reserve(mlp_weight_vector.size());
			for (auto &weight_buffer : weights_buffers)
			{
				weights_buffer_descriptors.emplace_back(create_descriptor(*weight_buffer));
			}
			VkWriteDescriptorSet weights_buffer_write = vkb::initializers::write_descriptor_set(
			    descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 4, weights_buffer_descriptors.data(), static_cast<uint32_t>(weights_buffer_descriptors.size()));

			std::vector<VkWriteDescriptorSet> write_descriptor_sets = {
			    texture_input_write_0,
			    texture_input_write_1,
			    texture_input_write_2,
			    texture_input_write_3,
			    weights_buffer_write};

			vkUpdateDescriptorSets(get_device().get_handle(), static_cast<uint32_t>(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE);
		}
		else
		{
			VkDescriptorBufferInfo weights_buffer_descriptor = create_descriptor(*weights_buffers[models[0].model_index]);
			VkWriteDescriptorSet   weights_buffer_write      = vkb::initializers::write_descriptor_set(descriptor_set_baseline[i], VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 3, &weights_buffer_descriptor);        // UBO

			std::vector<VkWriteDescriptorSet> write_descriptor_sets = {
			    texture_input_write_0,
			    texture_input_write_1,
			    texture_input_write_2,
			    weights_buffer_write};

			vkUpdateDescriptorSets(get_device().get_handle(), static_cast<uint32_t>(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE);
		}
	}
}

void MobileNerf::prepare_pipelines()
{
	VkPipelineInputAssemblyStateCreateInfo input_assembly_state = vkb::initializers::pipeline_input_assembly_state_create_info(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, 0, VK_FALSE);

	VkPipelineRasterizationStateCreateInfo rasterization_state = vkb::initializers::pipeline_rasterization_state_create_info(VK_POLYGON_MODE_FILL, /*VK_CULL_MODE_BACK_BIT*/ VK_CULL_MODE_NONE, VK_FRONT_FACE_COUNTER_CLOCKWISE /*VK_FRONT_FACE_CLOCKWISE*/, 0);

	std::vector<VkPipelineColorBlendAttachmentState> blend_attachment_states;
	blend_attachment_states.push_back(vkb::initializers::pipeline_color_blend_attachment_state(0xf, VK_FALSE));

	if (use_deferred)
	{
		blend_attachment_states.push_back(vkb::initializers::pipeline_color_blend_attachment_state(0xf, VK_FALSE));
		blend_attachment_states.push_back(vkb::initializers::pipeline_color_blend_attachment_state(0xf, VK_FALSE));
		if (combo_mode)
			blend_attachment_states.push_back(vkb::initializers::pipeline_color_blend_attachment_state(0xf, VK_FALSE));
	}

	VkPipelineColorBlendStateCreateInfo color_blend_state = vkb::initializers::pipeline_color_blend_state_create_info(static_cast<uint32_t>(blend_attachment_states.size()), blend_attachment_states.data());

	VkPipelineDepthStencilStateCreateInfo depth_stencil_state = vkb::initializers::pipeline_depth_stencil_state_create_info(VK_TRUE, VK_TRUE, VK_COMPARE_OP_LESS);
	depth_stencil_state.depthBoundsTestEnable                 = VK_FALSE;
	depth_stencil_state.minDepthBounds                        = 0.f;
	depth_stencil_state.maxDepthBounds                        = 1.f;

	VkPipelineViewportStateCreateInfo viewport_state = vkb::initializers::pipeline_viewport_state_create_info(1, 1, 0);

	std::vector<VkDynamicState> dynamic_state_enables = {
	    VK_DYNAMIC_STATE_VIEWPORT,
	    VK_DYNAMIC_STATE_SCISSOR};

	VkPipelineDynamicStateCreateInfo dynamic_state =
	    vkb::initializers::pipeline_dynamic_state_create_info(
	        dynamic_state_enables.data(),
	        static_cast<uint32_t>(dynamic_state_enables.size()),
	        0);

	VkPipelineMultisampleStateCreateInfo multisample_state = vkb::initializers::pipeline_multisample_state_create_info(VK_SAMPLE_COUNT_1_BIT, 0);

	// Vertex bindings and attributes
	const std::vector<VkVertexInputBindingDescription> vertex_input_bindings = {
	    vkb::initializers::vertex_input_binding_description(0, sizeof(Vertex), VK_VERTEX_INPUT_RATE_VERTEX),
	    vkb::initializers::vertex_input_binding_description(1, sizeof(InstanceData), VK_VERTEX_INPUT_RATE_INSTANCE),
	};
	const std::vector<VkVertexInputAttributeDescription> vertex_input_attributes = {
	    vkb::initializers::vertex_input_attribute_description(0, 0, VK_FORMAT_R32G32B32_SFLOAT, offsetof(Vertex, position)),
	    vkb::initializers::vertex_input_attribute_description(0, 1, VK_FORMAT_R32G32_SFLOAT, offsetof(Vertex, tex_coord)),
	    vkb::initializers::vertex_input_attribute_description(1, 2, VK_FORMAT_R32G32B32_SFLOAT, 0),
	};
	VkPipelineVertexInputStateCreateInfo vertex_input_state = vkb::initializers::pipeline_vertex_input_state_create_info();
	vertex_input_state.vertexBindingDescriptionCount        = static_cast<uint32_t>(vertex_input_bindings.size());
	vertex_input_state.pVertexBindingDescriptions           = vertex_input_bindings.data();
	vertex_input_state.vertexAttributeDescriptionCount      = static_cast<uint32_t>(vertex_input_attributes.size());
	vertex_input_state.pVertexAttributeDescriptions         = vertex_input_attributes.data();

	// First Pass

	VkGraphicsPipelineCreateInfo pipeline_create_info = vkb::initializers::pipeline_create_info(pipeline_first_pass_layout, render_pass_nerf, 0);
	pipeline_create_info.pVertexInputState            = &vertex_input_state;
	pipeline_create_info.pInputAssemblyState          = &input_assembly_state;
	pipeline_create_info.pRasterizationState          = &rasterization_state;
	pipeline_create_info.pColorBlendState             = &color_blend_state;
	pipeline_create_info.pMultisampleState            = &multisample_state;
	pipeline_create_info.pViewportState               = &viewport_state;
	pipeline_create_info.pDepthStencilState           = &depth_stencil_state;
	pipeline_create_info.pDynamicState                = &dynamic_state;
	pipeline_create_info.subpass                      = 0;
	pipeline_create_info.stageCount                   = static_cast<uint32_t>(shader_stages_first_pass.size());
	pipeline_create_info.pStages                      = shader_stages_first_pass.data();

	// Each model will have its own pipeline
	for (auto &model : models)
	{
		VK_CHECK(vkCreateGraphicsPipelines(get_device().get_handle(), pipeline_cache, 1, &pipeline_create_info, nullptr, &model.pipeline_first_pass));
	}

	if (use_deferred)
	{
		// Second Pass

		pipeline_create_info.layout  = pipeline_layout_baseline;
		pipeline_create_info.subpass = 1;

		VkPipelineVertexInputStateCreateInfo emptyInputStateCI{};
		emptyInputStateCI.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;

		pipeline_create_info.pVertexInputState = &emptyInputStateCI;
		color_blend_state.attachmentCount      = 1;
		rasterization_state.cullMode           = VK_CULL_MODE_NONE;
		depth_stencil_state.depthWriteEnable   = VK_FALSE;
		pipeline_create_info.stageCount        = static_cast<uint32_t>(shader_stages_second_pass.size());
		pipeline_create_info.pStages           = shader_stages_second_pass.data();

		VK_CHECK(vkCreateGraphicsPipelines(get_device().get_handle(), pipeline_cache, 1, &pipeline_create_info, nullptr, &pipeline_baseline));
	}
}

void MobileNerf::create_static_object_buffers(int model_index, int sub_model_index, int models_entry)
{
	LOGI("Creating static object buffers");
	Model &model              = models[models_entry];
	auto   vertex_buffer_size = model.vertices.size() * sizeof(Vertex);
	auto   index_buffer_size  = model.indices.size() * sizeof(model.indices[0]);

	// Create destination buffers
	model.vertex_buffer = std::make_unique<vkb::core::BufferC>(
	    get_device(),
	    vertex_buffer_size,
	    VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
	    VMA_MEMORY_USAGE_GPU_ONLY);
	model.vertex_buffer->set_debug_name(fmt::format("Model #{} Sub-Model #{} vertices", model_index, sub_model_index));
	model.index_buffer = std::make_unique<vkb::core::BufferC>(
	    get_device(),
	    index_buffer_size,
	    VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
	    VMA_MEMORY_USAGE_GPU_ONLY);
	model.index_buffer->set_debug_name(fmt::format("Model #{} Sub-Model #{} indices", model_index, sub_model_index));

	// Create staging buffers
	std::unique_ptr<vkb::core::BufferC> staging_vertex_buffer = std::make_unique<vkb::core::BufferC>(
	    get_device(),
	    vertex_buffer_size,
	    VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
	    VMA_MEMORY_USAGE_CPU_TO_GPU);
	staging_vertex_buffer->update(model.vertices);

	std::unique_ptr<vkb::core::BufferC> staging_index_buffer = std::make_unique<vkb::core::BufferC>(
	    get_device(),
	    index_buffer_size,
	    VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
	    VMA_MEMORY_USAGE_CPU_TO_GPU);
	staging_index_buffer->update(model.indices);

	// Copy over the data for each of the models
	with_vkb_command_buffer(
	    [&](vkb::core::CommandBufferC &cmd) {
		    cmd.copy_buffer(*staging_vertex_buffer, *model.vertex_buffer, staging_vertex_buffer->get_size());
		    cmd.copy_buffer(*staging_index_buffer, *model.index_buffer, staging_index_buffer->get_size());
	    });

	LOGI("Done Creating static object buffers");
}

void MobileNerf::create_uniforms()
{
	uniform_buffers.resize(model_path.size());
	weights_buffers.resize(model_path.size());

	for (int i = 0; i < model_path.size(); i++)
	{
		LOGI("Creating camera view uniform buffer for model {}", i);
		uniform_buffers[i] = std::make_unique<vkb::core::BufferC>(get_device(),
		                                                          sizeof(global_uniform),
		                                                          VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
		                                                          VMA_MEMORY_USAGE_CPU_TO_GPU);

		LOGI("Creating mlp weights uniform buffer for model {}", i);
		weights_buffers[i] = std::make_unique<vkb::core::BufferC>(get_device(),
		                                                          sizeof(MLP_Weights),
		                                                          VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
		                                                          VMA_MEMORY_USAGE_CPU_TO_GPU);
	}

	update_uniform_buffers();
	update_weights_buffers();
}

void MobileNerf::initialize_mlp_uniform_buffers(int model_index)
{
	std::string assetBase   = vkb::fs::path::get(vkb::fs::path::Type::Assets);
	std::string mlpJsonPath = assetBase + model_path[model_index] + "mlp.json";

	using json = nlohmann::json;

	std::ifstream f(mlpJsonPath);

	if (!f)
	{
		LOGE("Failed to open mlp data");
		assert(0);
	}

	LOGI("Parsing mlp data {}", mlpJsonPath);
	json data = json::parse(f);

	// Record a index of the first sub-model
	int first_sub_model = static_cast<int>(models.size());
	int obj_num         = data["obj_num"].get<int>();

	// Here we know the actual number of sub models
	int next_sub_model_index = static_cast<int>(models.size());
	models.resize(models.size() + obj_num);

	for (int i = next_sub_model_index; i < models.size(); i++)
	{
		models[i].model_index = model_index;
	}

	auto weights_0_array_raw = data["0_weights"].get<std::vector<std::vector<float>>>();

	std::vector<float> weights_0_array;

	for (auto ii = weights_0_array_raw.begin(); ii != weights_0_array_raw.end(); ii++)
	{
		weights_0_array.insert(weights_0_array.end(), (*ii).begin(), (*ii).end());
	}

	if (weights_0_array.size() != WEIGHTS_0_COUNT)
	{
		LOGE("MLP data layer 0 weights count is {}, rather than {}", weights_0_array.size(), WEIGHTS_0_COUNT);
	}

	auto bias_0_array = data["0_bias"].get<std::vector<float>>();

	if (bias_0_array.size() != BIAS_0_COUNT)
	{
		LOGE("MLP data layer 0 bias count is {}, rather than {}", bias_0_array.size(), BIAS_0_COUNT);
	}

	auto weights_1_array_raw = data["1_weights"].get<std::vector<std::vector<float>>>();

	std::vector<float> weights_1_array;

	for (auto ii = weights_1_array_raw.begin(); ii != weights_1_array_raw.end(); ii++)
	{
		weights_1_array.insert(weights_1_array.end(), (*ii).begin(), (*ii).end());
	}

	if (weights_1_array.size() != WEIGHTS_1_COUNT)
	{
		LOGE("MLP data layer 1 weights count is {}, rather than {}", weights_1_array.size(), WEIGHTS_1_COUNT);
	}

	auto bias_1_array = data["1_bias"].get<std::vector<float>>();

	if (bias_1_array.size() != BIAS_1_COUNT)
	{
		LOGE("MLP data layer 1 bias count is {}, rather than {}", bias_1_array.size(), BIAS_1_COUNT);
	}

	auto weights_2_array_raw = data["2_weights"].get<std::vector<std::vector<float>>>();

	std::vector<float> weights_2_array;

	for (auto ii = weights_2_array_raw.begin(); ii != weights_2_array_raw.end(); ii++)
	{
		weights_2_array.insert(weights_2_array.end(), (*ii).begin(), (*ii).end());
	}

	// We need to pad the layer 2's weights with 16 zeros
	if (weights_2_array.size() != WEIGHTS_2_COUNT - 16)
	{
		LOGE("MLP data layer 2 weights count is {}, rather than {}", weights_2_array.size(), WEIGHTS_2_COUNT);
	}

	auto bias_2_array = data["2_bias"].get<std::vector<float>>();

	if (bias_2_array.size() != BIAS_2_COUNT - 1)
	{
		LOGE("MLP data layer 2 bias count is {}, rather than {}", bias_2_array.size(), BIAS_2_COUNT);
	}

	// Each sub model will share the same mlp weights data
	MLP_Weights &model_mlp = mlp_weight_vector[model_index];

	for (int ii = 0; ii < WEIGHTS_0_COUNT; ii++)
	{
		model_mlp.data[ii] = weights_0_array[ii];
	}

	for (int ii = 0; ii < WEIGHTS_1_COUNT; ii++)
	{
		model_mlp.data[WEIGHTS_0_COUNT + ii] = weights_1_array[ii];
	}

	// We need to pad the layer 2's weights with zeros for every 3 weights to make it 16 bytes aligned
	int raw_weight_cnt = 0;
	for (int ii = 0; ii < WEIGHTS_2_COUNT; ii++)
	{
		if ((ii + 1) % 4 == 0)
		{
			model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + ii] = 0.0f;
		}
		else
		{
			model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + ii] = weights_2_array[raw_weight_cnt++];
		}
	}

	for (int ii = 0; ii < BIAS_0_COUNT; ii++)
	{
		model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT + ii] = bias_0_array[ii];
	}

	for (int ii = 0; ii < BIAS_1_COUNT; ii++)
	{
		model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT +
		               BIAS_0_COUNT + ii] = bias_1_array[ii];
	}

	// We need to pad the layer 2's bias with zeros for every 3 weights to make it 16 bytes aligned
	for (int ii = 0; ii < BIAS_2_COUNT; ii++)
	{
		if ((ii + 1) % 4 == 0)
		{
			model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT +
			               BIAS_0_COUNT + BIAS_1_COUNT + ii] = 0.0f;
		}
		else
		{
			model_mlp.data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT +
			               BIAS_0_COUNT + BIAS_1_COUNT + ii] = bias_2_array[ii];
		}
	}

	// Update all sub model with the same mlp weight
	for (int i = 0; i < obj_num; i++)
	{
		models[first_sub_model + i].sub_model_num = obj_num;
	}
}

void MobileNerf::update_uniform_buffers()
{
	assert(uniform_buffers[0]);

	const float tan_half_fov = tan(0.5 * fov / 180.0f * 3.141592653589793f);

	global_uniform.proj            = camera.matrices.perspective;
	global_uniform.view            = camera.matrices.view;
	global_uniform.camera_position = camera.position;
	global_uniform.camera_side     = glm::vec3(camera.matrices.view[0][0], camera.matrices.view[1][0], camera.matrices.view[2][0]);
	global_uniform.camera_up       = glm::vec3(camera.matrices.view[0][1], camera.matrices.view[1][1], camera.matrices.view[2][1]);
	global_uniform.camera_lookat   = -glm::vec3(camera.matrices.view[0][2], camera.matrices.view[1][2], camera.matrices.view[2][2]);
	global_uniform.img_dim         = glm::vec2(width, height);
	global_uniform.tan_half_fov    = tan_half_fov;

	for (int i = 0; i < model_path.size(); i++)
	{
		// Note that this is a hard-coded scene setting for the lego_combo
		global_uniform.model = combo_mode ? combo_model_transform[i] : glm::translate(glm::vec3(0.0f));
		uniform_buffers[i]->update(&global_uniform, sizeof(global_uniform));
	}
}

void MobileNerf::update_weights_buffers()
{
	// No need to be updated for every frames
	for (int i = 0; i < model_path.size(); i++)
	{
		weights_buffers[i]->update(&(mlp_weight_vector[i].data[0]), sizeof(MLP_Weights));
	}
}

void MobileNerf::prepare_instance_data()
{
	auto &ii = instancing_info;

	std::vector<InstanceData> instance_data;
	instance_data.resize(ii.dim.x * ii.dim.y * ii.dim.z);

	const glm::vec3 corner_pos = -ii.interval * 0.5f * (glm::vec3(ii.dim - 1));
	int             idx        = 0;
	glm::vec3       offset;
	for (int x = 0; x < ii.dim.x; ++x)
	{
		offset.x = corner_pos.x + ii.interval.x * x;
		for (int y = 0; y < ii.dim.y; ++y)
		{
			offset.y = corner_pos.y + ii.interval.y * y;
			for (int z = 0; z < ii.dim.z; ++z)
			{
				offset.z                        = corner_pos.z + ii.interval.z * z;
				instance_data[idx++].pos_offset = offset;
			}
		}
	}

	auto instance_buffer_size = instance_data.size() * sizeof(InstanceData);

	instance_buffer = std::make_unique<vkb::core::BufferC>(
	    get_device(),
	    instance_buffer_size,
	    VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
	    VMA_MEMORY_USAGE_GPU_ONLY);

	// Copy over the data for each of the models
	auto staging_instance_buffer = std::make_unique<vkb::core::BufferC>(
	    get_device(),
	    instance_buffer_size,
	    VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
	    VMA_MEMORY_USAGE_CPU_TO_GPU);
	staging_instance_buffer->update(instance_data);

	// now transfer over to the end buffer
	with_vkb_command_buffer([&](vkb::core::CommandBufferC &cmd) { cmd.copy_buffer(*staging_instance_buffer, *instance_buffer, staging_instance_buffer->get_size()); });
}

void MobileNerf::draw()
{
	ApiVulkanSample::prepare_frame();

	// Command buffer to be submitted to the queue
	submit_info.commandBufferCount = 1;
	submit_info.pCommandBuffers    = &draw_cmd_buffers[current_buffer];

	// Submit to queue
	VK_CHECK(vkQueueSubmit(queue, 1, &submit_info, VK_NULL_HANDLE));

	ApiVulkanSample::submit_frame();
}

void MobileNerf::create_texture(int model_index, int sub_model_index, int models_entry)
{
	// Set up the input texture image

	// TODO(tomatkinson): should load different scenes's feature map from command line
	std::string feature_0_path = model_path[model_index] + "shape" + std::to_string(sub_model_index) + ".pngfeat0.png";
	std::string feature_1_path = model_path[model_index] + "shape" + std::to_string(sub_model_index) + ".pngfeat1.png";

	LOGI("Creating feature texture 0");
	create_texture_helper(feature_0_path, models[models_entry].texture_input_0);
	LOGI("Done Creating feature texture 0");

	LOGI("Creating feature texture 1");
	create_texture_helper(feature_1_path, models[models_entry].texture_input_1);
	LOGI("Done Creating feature texture 0");
}

void MobileNerf::create_texture_helper(std::string const &texturePath, Texture &texture_input)
{
	// Feature textures are in linear space instead of sRGB space
	texture_input = load_texture(texturePath, vkb::sg::Image::Other);
	vkDestroySampler(get_device().get_handle(), texture_input.sampler, nullptr);

	// Calculate valid filter
	VkFilter filter = using_original_nerf_models[0] ? VK_FILTER_NEAREST : VK_FILTER_LINEAR;
	vkb::make_filters_valid(get_device().get_gpu().get_handle(), texture_input.image->get_format(), &filter);

	VkSamplerCreateInfo samplerCreateInfo     = {};
	samplerCreateInfo.sType                   = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
	samplerCreateInfo.magFilter               = filter;
	samplerCreateInfo.minFilter               = filter;
	samplerCreateInfo.mipmapMode              = VK_SAMPLER_MIPMAP_MODE_NEAREST;
	samplerCreateInfo.addressModeU            = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE;
	samplerCreateInfo.addressModeV            = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE;
	samplerCreateInfo.addressModeW            = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE;
	samplerCreateInfo.minLod                  = 0.0f;
	samplerCreateInfo.maxLod                  = 16.0f;
	samplerCreateInfo.unnormalizedCoordinates = VK_FALSE;
	VK_CHECK(vkCreateSampler(get_device().get_handle(), &samplerCreateInfo, 0, &texture_input.sampler));
}

void MobileNerf::update_render_pass_nerf_forward()
{
	// For merged shaders, we need 2 attachments (as opposed to 5)
	// 0: Depth attachment
	// 1: Swapchain attachment
	std::array<VkAttachmentDescription, 2> attachments = {};
	// Depth attachment
	attachments[0].format         = depth_format;
	attachments[0].samples        = VK_SAMPLE_COUNT_1_BIT;
	attachments[0].loadOp         = VK_ATTACHMENT_LOAD_OP_CLEAR;
	attachments[0].storeOp        = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	attachments[0].stencilLoadOp  = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	attachments[0].stencilStoreOp = VK_ATTACHMENT_STORE_OP_STORE;
	attachments[0].initialLayout  = VK_IMAGE_LAYOUT_UNDEFINED;
	attachments[0].finalLayout    = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
	// Swapchain  attachment
	attachments[1].format         = get_render_context().get_swapchain().get_format();
	attachments[1].samples        = VK_SAMPLE_COUNT_1_BIT;
	attachments[1].loadOp         = VK_ATTACHMENT_LOAD_OP_CLEAR;
	attachments[1].storeOp        = VK_ATTACHMENT_STORE_OP_STORE;
	attachments[1].stencilLoadOp  = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	attachments[1].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	attachments[1].initialLayout  = VK_IMAGE_LAYOUT_UNDEFINED;
	attachments[1].finalLayout    = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

	VkAttachmentReference depth_reference = {};
	depth_reference.attachment            = 0;
	depth_reference.layout                = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

	VkAttachmentReference swapchain_reference = {};
	swapchain_reference.attachment            = 1;
	swapchain_reference.layout                = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

	VkSubpassDescription subpass    = {};
	subpass.pipelineBindPoint       = VK_PIPELINE_BIND_POINT_GRAPHICS;
	subpass.colorAttachmentCount    = 1;
	subpass.pColorAttachments       = &swapchain_reference;
	subpass.pDepthStencilAttachment = &depth_reference;
	subpass.inputAttachmentCount    = 0;
	subpass.pInputAttachments       = nullptr;
	subpass.preserveAttachmentCount = 0;
	subpass.pPreserveAttachments    = nullptr;
	subpass.pResolveAttachments     = nullptr;

	VkRenderPassCreateInfo render_pass_create_info = {};
	render_pass_create_info.sType                  = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
	render_pass_create_info.attachmentCount        = static_cast<uint32_t>(attachments.size());
	render_pass_create_info.pAttachments           = attachments.data();
	render_pass_create_info.subpassCount           = 1;
	render_pass_create_info.pSubpasses             = &subpass;

	VK_CHECK(vkCreateRenderPass(get_device().get_handle(), &render_pass_create_info, nullptr, &render_pass_nerf));
}

void MobileNerf::update_render_pass_nerf_baseline()
{
	unsigned int attachment_idx = 0;

	// Color attachment 0 - feature 0 G-buffer
	color_attach_0_idx                          = attachment_idx++;
	VkAttachmentDescription color_description_0 = {};
	color_description_0.format                  = feature_map_format;
	color_description_0.samples                 = VK_SAMPLE_COUNT_1_BIT;
	color_description_0.loadOp                  = VK_ATTACHMENT_LOAD_OP_CLEAR;
	color_description_0.storeOp                 = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	color_description_0.stencilLoadOp           = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	color_description_0.stencilStoreOp          = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	color_description_0.initialLayout           = VK_IMAGE_LAYOUT_UNDEFINED;
	color_description_0.finalLayout             = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	// Color attachment 1 - feature 1 G-buffer
	color_attach_1_idx                          = attachment_idx++;
	VkAttachmentDescription color_description_1 = {};
	color_description_1.format                  = feature_map_format;
	color_description_1.samples                 = VK_SAMPLE_COUNT_1_BIT;
	color_description_1.loadOp                  = VK_ATTACHMENT_LOAD_OP_CLEAR;
	color_description_1.storeOp                 = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	color_description_1.stencilLoadOp           = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	color_description_1.stencilStoreOp          = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	color_description_1.initialLayout           = VK_IMAGE_LAYOUT_UNDEFINED;
	color_description_1.finalLayout             = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	// Color attachment 2 - ray direction G-buffer
	color_attach_2_idx                          = attachment_idx++;
	VkAttachmentDescription color_description_2 = {};
	color_description_2.format                  = VK_FORMAT_R16G16B16A16_SFLOAT;
	color_description_2.samples                 = VK_SAMPLE_COUNT_1_BIT;
	color_description_2.loadOp                  = VK_ATTACHMENT_LOAD_OP_CLEAR;
	color_description_2.storeOp                 = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	color_description_2.stencilLoadOp           = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	color_description_2.stencilStoreOp          = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	color_description_2.initialLayout           = VK_IMAGE_LAYOUT_UNDEFINED;
	color_description_2.finalLayout             = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	// Color attachment 3 - weight index G-buffer
	VkAttachmentDescription color_description_3 = {};
	color_attach_3_idx                          = 3;
	if (combo_mode)
	{
		color_attach_3_idx                 = attachment_idx++;
		color_description_3.format         = VK_FORMAT_R8_UINT;
		color_description_3.samples        = VK_SAMPLE_COUNT_1_BIT;
		color_description_3.loadOp         = VK_ATTACHMENT_LOAD_OP_CLEAR;
		color_description_3.storeOp        = VK_ATTACHMENT_STORE_OP_DONT_CARE;
		color_description_3.stencilLoadOp  = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
		color_description_3.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
		color_description_3.initialLayout  = VK_IMAGE_LAYOUT_UNDEFINED;
		color_description_3.finalLayout    = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	}
	// Depth attachment
	depth_attach_idx                          = attachment_idx++;
	VkAttachmentDescription depth_description = {};
	depth_description.format                  = depth_format;
	depth_description.samples                 = VK_SAMPLE_COUNT_1_BIT;
	depth_description.loadOp                  = VK_ATTACHMENT_LOAD_OP_CLEAR;
	depth_description.storeOp                 = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	depth_description.stencilLoadOp           = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	depth_description.stencilStoreOp          = VK_ATTACHMENT_STORE_OP_STORE;
	depth_description.initialLayout           = VK_IMAGE_LAYOUT_UNDEFINED;
	depth_description.finalLayout             = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
	// Swapchain  attachment
	swapchain_attach_idx                          = attachment_idx++;
	VkAttachmentDescription swapchain_description = {};
	swapchain_description.format                  = get_render_context().get_swapchain().get_format();
	swapchain_description.samples                 = VK_SAMPLE_COUNT_1_BIT;
	swapchain_description.loadOp                  = VK_ATTACHMENT_LOAD_OP_CLEAR;
	swapchain_description.storeOp                 = VK_ATTACHMENT_STORE_OP_STORE;
	swapchain_description.stencilLoadOp           = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
	swapchain_description.stencilStoreOp          = VK_ATTACHMENT_STORE_OP_DONT_CARE;
	swapchain_description.initialLayout           = VK_IMAGE_LAYOUT_UNDEFINED;
	swapchain_description.finalLayout             = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

	std::vector<VkAttachmentDescription> attachments;
	attachments.push_back(color_description_0);
	attachments.push_back(color_description_1);
	attachments.push_back(color_description_2);
	if (combo_mode)
		attachments.push_back(color_description_3);
	attachments.push_back(depth_description);
	attachments.push_back(swapchain_description);

	VkAttachmentReference color_reference_0 = {};
	color_reference_0.attachment            = color_attach_0_idx;
	color_reference_0.layout                = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

	VkAttachmentReference color_reference_1 = {};
	color_reference_1.attachment            = color_attach_1_idx;
	color_reference_1.layout                = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

	VkAttachmentReference color_reference_2 = {};
	color_reference_2.attachment            = color_attach_2_idx;
	color_reference_2.layout                = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

	VkAttachmentReference color_reference_3 = {};
	color_reference_3.attachment            = color_attach_3_idx;
	color_reference_3.layout                = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

	VkAttachmentReference depth_reference = {};
	depth_reference.attachment            = depth_attach_idx;
	depth_reference.layout                = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

	VkAttachmentReference swapchain_reference = {};
	swapchain_reference.attachment            = swapchain_attach_idx;
	swapchain_reference.layout                = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

	std::array<VkSubpassDescription, 2> subpassDescriptions{};

	std::vector<VkAttachmentReference> color_references_feature_maps = {color_reference_0, color_reference_1, color_reference_2};
	if (combo_mode)
		color_references_feature_maps.push_back(color_reference_3);

	subpassDescriptions[0].pipelineBindPoint       = VK_PIPELINE_BIND_POINT_GRAPHICS;
	subpassDescriptions[0].colorAttachmentCount    = static_cast<uint32_t>(color_references_feature_maps.size());
	subpassDescriptions[0].pColorAttachments       = color_references_feature_maps.data();
	subpassDescriptions[0].pDepthStencilAttachment = &depth_reference;
	subpassDescriptions[0].inputAttachmentCount    = 0;
	subpassDescriptions[0].pInputAttachments       = nullptr;
	subpassDescriptions[0].preserveAttachmentCount = 0;
	subpassDescriptions[0].pPreserveAttachments    = nullptr;
	subpassDescriptions[0].pResolveAttachments     = nullptr;

	// Color attachments written to in first sub pass will be used as input attachments to be read in the fragment shader
	std::vector<VkAttachmentReference> inputReferences = {
	    {color_attach_0_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL},        // Color attachment 0 - feature 0 G-buffer
	    {color_attach_1_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL},        // Color attachment 1 - feature 1 G-buffer
	    {color_attach_2_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}         // Color attachment 2 - ray direction G-buffer
	};
	if (combo_mode)
		inputReferences.push_back({color_attach_3_idx, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL});        // Color attachment 3 - weight index G-buffer

	subpassDescriptions[1].pipelineBindPoint       = VK_PIPELINE_BIND_POINT_GRAPHICS;
	subpassDescriptions[1].colorAttachmentCount    = 1;
	subpassDescriptions[1].pColorAttachments       = &swapchain_reference;
	subpassDescriptions[1].pDepthStencilAttachment = nullptr;
	subpassDescriptions[1].inputAttachmentCount    = static_cast<uint32_t>(inputReferences.size());
	subpassDescriptions[1].pInputAttachments       = inputReferences.data();
	subpassDescriptions[1].preserveAttachmentCount = 0;
	subpassDescriptions[1].pPreserveAttachments    = nullptr;
	subpassDescriptions[1].pResolveAttachments     = nullptr;

	// Subpass dependencies for layout transitions
	std::array<VkSubpassDependency, 3> dependencies{};

	dependencies[0].srcSubpass      = VK_SUBPASS_EXTERNAL;
	dependencies[0].dstSubpass      = 0;
	dependencies[0].srcStageMask    = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
	dependencies[0].dstStageMask    = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
	dependencies[0].srcAccessMask   = VK_ACCESS_NONE;
	dependencies[0].dstAccessMask   = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
	dependencies[0].dependencyFlags = VK_DEPENDENCY_BY_REGION_BIT;

	dependencies[1].srcSubpass      = 0;
	dependencies[1].dstSubpass      = 1;
	dependencies[1].srcStageMask    = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
	dependencies[1].dstStageMask    = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
	dependencies[1].srcAccessMask   = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
	dependencies[1].dstAccessMask   = VK_ACCESS_INPUT_ATTACHMENT_READ_BIT;
	dependencies[1].dependencyFlags = VK_DEPENDENCY_BY_REGION_BIT;

	dependencies[2].srcSubpass      = 1;
	dependencies[2].dstSubpass      = VK_SUBPASS_EXTERNAL;
	dependencies[2].srcStageMask    = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
	dependencies[2].dstStageMask    = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
	dependencies[2].srcAccessMask   = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
	dependencies[2].dstAccessMask   = VK_ACCESS_NONE;
	dependencies[2].dependencyFlags = VK_DEPENDENCY_BY_REGION_BIT;

	VkRenderPassCreateInfo render_pass_create_info = {};
	render_pass_create_info.sType                  = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
	render_pass_create_info.attachmentCount        = static_cast<uint32_t>(attachments.size());
	render_pass_create_info.pAttachments           = attachments.data();
	render_pass_create_info.subpassCount           = static_cast<uint32_t>(subpassDescriptions.size());
	render_pass_create_info.pSubpasses             = subpassDescriptions.data();
	render_pass_create_info.dependencyCount        = static_cast<uint32_t>(dependencies.size());
	render_pass_create_info.pDependencies          = dependencies.data();

	VK_CHECK(vkCreateRenderPass(get_device().get_handle(), &render_pass_create_info, nullptr, &render_pass_nerf));
}

std::unique_ptr<vkb::VulkanSampleC> create_mobile_nerf()
{
	return std::make_unique<MobileNerf>();
}