main.cpp

#define GLFW_INCLUDE_VULKAN
#include <GLFW/glfw3.h>

#include "VulkanSwapchain.h"
#include "VulkanRayTracing.h"
#include "VulkanAcceleration.h"
#include "VulkanRTPipeline.h"
#include "Camera.h"
#include "EntityManager.h"
#include "Transform.h"
#include "NvidiaDenoiser.h"
#include "AsyncLogger.h"

#include <iostream>
#include <stdexcept>
#include <vector>
#include <cstring>
#include <cstdlib>
#include <optional>
#include <set>
#include <chrono>
#include <filesystem>
#include <fstream>
#ifdef _WIN32
#include <windows.h>
#endif

// Set to 1 to enable verbose logging in the main render loop (performance impact)
#define ENABLE_RENDER_LOOP_LOGGING 0

// Vulkan Debug Callback
static VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(
    VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
    VkDebugUtilsMessageTypeFlagsEXT messageType,
    const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
    void* pUserData) {

    const char* severity = "UNKNOWN";
    if (messageSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT) severity = "ERROR";
    else if (messageSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT) severity = "WARNING";
    else if (messageSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT) severity = "INFO";
    else if (messageSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT) severity = "VERBOSE";

    // Only print errors and warnings
    if (messageSeverity >= VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT) {
        std::cerr << "\n[VULKAN " << severity << "] " << pCallbackData->pMessage << "\n" << std::endl;
    }

    return VK_FALSE;
}

const uint32_t WIDTH = 1920;
const uint32_t HEIGHT = 1080;

const std::vector<const char*> validationLayers = {
    "VK_LAYER_KHRONOS_validation"
};

const std::vector<const char*> deviceExtensions = {
    VK_KHR_SWAPCHAIN_EXTENSION_NAME,
    VK_KHR_RAY_TRACING_PIPELINE_EXTENSION_NAME,
    VK_KHR_ACCELERATION_STRUCTURE_EXTENSION_NAME,
    VK_KHR_DEFERRED_HOST_OPERATIONS_EXTENSION_NAME,
    VK_KHR_BUFFER_DEVICE_ADDRESS_EXTENSION_NAME,
    VK_KHR_RAY_QUERY_EXTENSION_NAME,                    // Inline RT in compute shaders
    VK_KHR_PIPELINE_LIBRARY_EXTENSION_NAME,             // RT pipeline compilation caching
    VK_KHR_RAY_TRACING_MAINTENANCE_1_EXTENSION_NAME,    // TraceRayIndirect, pipeline flags
    VK_KHR_SPIRV_1_4_EXTENSION_NAME,                    // Required for ray query
    VK_KHR_SHADER_FLOAT_CONTROLS_EXTENSION_NAME,        // Required for SPIRV 1.4
    VK_KHR_EXTERNAL_MEMORY_EXTENSION_NAME,              // Interop for denoiser
    VK_KHR_EXTERNAL_SEMAPHORE_EXTENSION_NAME,           // Sync for denoiser
#ifdef _WIN32
    "VK_KHR_external_memory_win32",
    "VK_KHR_external_semaphore_win32",
#endif
};

#ifdef NDEBUG
    const bool enableValidationLayers = false;
#else
    const bool enableValidationLayers = true;
#endif

struct QueueFamilyIndices {
    std::optional<uint32_t> graphicsFamily;
    std::optional<uint32_t> presentFamily;

    bool isComplete() {
        return graphicsFamily.has_value() && presentFamily.has_value();
    }
};

// Get project root directory (where shaders folder is located)
std::string getProjectRoot() {
#ifdef _WIN32
    char buffer[MAX_PATH] = {0};
    GetModuleFileNameA(NULL, buffer, MAX_PATH);
    std::filesystem::path exePath(buffer);

    // Navigate up from exe location to find shaders folder
    std::filesystem::path searchPath = exePath.parent_path();
    for (int i = 0; i < 5; i++) {  // Search up to 5 levels
        if (std::filesystem::exists(searchPath / "shaders" / "compiled")) {
            return searchPath.string();
        }
        searchPath = searchPath.parent_path();
    }
#endif
    // Fallback to current directory
    return std::filesystem::current_path().string();
}

class RacingEngine {
public:
    void run() {
        initWindow();
        initVulkan();
        mainLoop();
        cleanup();
    }

private:
    GLFWwindow* window;
    VkInstance instance;
    VkDebugUtilsMessengerEXT debugMessenger = VK_NULL_HANDLE;
    VkPhysicalDevice physicalDevice = VK_NULL_HANDLE;
    VkDevice device;
    VkQueue graphicsQueue;
    VkQueue presentQueue;
    VkSurfaceKHR surface;

    VulkanSwapchain swapchain;
    VulkanRayTracing rayTracing;
    VulkanAcceleration acceleration;
    VulkanRTPipeline rtPipeline;

    // OptiX AI Denoiser (Tensor Cores)
    NvidiaDenoiserImpl denoiser;

    QueueFamilyIndices queueIndices;

    // Entity system
    EntityManager entityManager;

    // Camera
    Camera camera;
    float lastX = WIDTH / 2.0f;
    float lastY = HEIGHT / 2.0f;
    bool firstMouse = true;

    // Timing
    std::chrono::time_point<std::chrono::high_resolution_clock> lastFrameTime;
    float deltaTime = 0.0f;
    float currentFPS = 0.0f;
    float frameTimeMs = 0.0f;

    // Sync objects - one set per swapchain image for proper synchronization
    std::vector<VkSemaphore> imageAvailableSemaphores;
    std::vector<VkSemaphore> renderFinishedSemaphores;
    std::vector<VkFence> inFlightFences;
    std::vector<VkFence> imagesInFlight;  // Track which fence is using which swapchain image
    size_t currentFrame = 0;
    const int MAX_FRAMES_IN_FLIGHT = 2;

    // Temporal accumulation & reprojection
    uint32_t accumulationFrames = 0;
    uint32_t totalFrameCount = 0;
    glm::vec3 lastCameraPos = glm::vec3(0.0f);
    glm::vec3 lastCameraFront = glm::vec3(0.0f);

    // Previous frame matrices for temporal reprojection
    glm::mat4 prevViewProj = glm::mat4(1.0f);
    glm::vec3 prevCameraPosition = glm::vec3(0.0f);
    bool hasPreviousFrame = false;  // First frame has no history

    // Track first use of storage image (initialized once, then always TRANSFER_SRC after use)
    bool storageImageInitialized = false;

    // Denoiser toggle (D key)
    bool denoiserEnabled = true;
    bool denoiserKeyPressed = false;

    // Cursor toggle (TAB key)
    bool cursorLocked = true;
    bool tabKeyPressed = false;

    // Async Logger
    AsyncLogger logger;

    // UX State Tracking
    void updateWindowTitle() {
        glm::vec3 pos = camera.getPosition();
        char title[512];
        uint32_t totalSamples = accumulationFrames * 8;  // 8 SPP per frame

        const char* denoiserStatus = denoiserEnabled ? "ON" : "OFF";
        const char* cursorStatus = cursorLocked ? "LOCKED" : "FREE";

        snprintf(title, sizeof(title),
                 "IZTAPALAPA PATH TRACER | FPS: %.0f | %.2fms | %u samples | Pos: (%.1f, %.1f, %.1f) | AI: %s | Cursor: %s",
                 currentFPS, frameTimeMs, totalSamples, pos.x, pos.y, pos.z, denoiserStatus, cursorStatus);
        glfwSetWindowTitle(window, title);
    }

    void initWindow() {
        glfwInit();
        glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
        glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);

        window = glfwCreateWindow(WIDTH, HEIGHT, "Racing Engine - Vulkan RT", nullptr, nullptr);

        // Setup mouse input
        glfwSetWindowUserPointer(window, this);
        glfwSetCursorPosCallback(window, mouseCallback);
        glfwSetKeyCallback(window, keyCallback);
        glfwSetScrollCallback(window, scrollCallback);
        glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_DISABLED);

        // Initialize camera with correct aspect ratio
        camera = Camera(45.0f, (float)WIDTH / (float)HEIGHT, 0.1f, 1000.0f);

        std::cout << "GLFW window created successfully!\n";
        std::cout << "Camera controls: WASD - move, QE - up/down, Mouse - look, Shift - faster\n";
        std::cout << "                 TAB - toggle cursor lock, Scroll - zoom (FOV)\n";
        std::cout << "                 R - Reset Camera\n";
        std::cout << "Press D to toggle AI Denoiser (Tensor Cores)\n";
    }

    static void scrollCallback(GLFWwindow* window, double xoffset, double yoffset) {
        auto engine = reinterpret_cast<RacingEngine*>(glfwGetWindowUserPointer(window));
        if (engine->cursorLocked) {
            engine->camera.processMouseScroll(static_cast<float>(yoffset));
            engine->updateWindowTitle();
        }
    }

    static void mouseCallback(GLFWwindow* window, double xposIn, double yposIn) {
        auto engine = reinterpret_cast<RacingEngine*>(glfwGetWindowUserPointer(window));

        float xpos = static_cast<float>(xposIn);
        float ypos = static_cast<float>(yposIn);

        if (engine->firstMouse) {
            engine->lastX = xpos;
            engine->lastY = ypos;
            engine->firstMouse = false;
        }

        float xoffset = xpos - engine->lastX;
        float yoffset = engine->lastY - ypos; // Reversed since y-coordinates go from bottom to top

        engine->lastX = xpos;
        engine->lastY = ypos;

        if (engine->cursorLocked) {
            engine->camera.processMouseMovement(xoffset, yoffset);
        }
    }

    static void keyCallback(GLFWwindow* window, int key, int scancode, int action, int mods) {
        auto engine = reinterpret_cast<RacingEngine*>(glfwGetWindowUserPointer(window));

        // ESC to close window
        if (key == GLFW_KEY_ESCAPE && action == GLFW_PRESS) {
            glfwSetWindowShouldClose(window, true);
        }

        // TAB key to toggle cursor lock
        if (key == GLFW_KEY_TAB && action == GLFW_PRESS && !engine->tabKeyPressed) {
            engine->tabKeyPressed = true;
            engine->cursorLocked = !engine->cursorLocked;

            if (engine->cursorLocked) {
                glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_DISABLED);
                engine->firstMouse = true; // Reset mouse to prevent jump
            } else {
                glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_NORMAL);
            }

            std::string status = engine->cursorLocked ? "LOCKED (Camera Active)" : "UNLOCKED (Camera Paused)";
            engine->logger.log("\n[CURSOR] " + status + "\n");
            engine->updateWindowTitle();
        }
        if (key == GLFW_KEY_TAB && action == GLFW_RELEASE) {
            engine->tabKeyPressed = false;
        }

        // D key to toggle denoiser
        if (key == GLFW_KEY_D && action == GLFW_PRESS && !engine->denoiserKeyPressed) {
            engine->denoiserKeyPressed = true;
            engine->denoiserEnabled = !engine->denoiserEnabled;

            std::string status = engine->denoiserEnabled ? "ENABLED" : "DISABLED";
            engine->logger.log("\n[DENOISER] " + status + "\n");
            engine->updateWindowTitle();
        }
        if (key == GLFW_KEY_D && action == GLFW_RELEASE) {
            engine->denoiserKeyPressed = false;
        }

        // R key to reset camera
        if (key == GLFW_KEY_R && action == GLFW_PRESS) {
            engine->camera.reset();
            engine->accumulationFrames = 0; // Reset accumulation immediately
            engine->logger.log("\n[CAMERA] Reset to default view\n");
            engine->updateWindowTitle();
        }
    }

    void initVulkan() {
        createInstance();
        setupDebugMessenger();
        createSurface();
        pickPhysicalDevice();
        createLogicalDevice();

        // Create swapchain
        swapchain.create(physicalDevice, device, surface, WIDTH, HEIGHT,
                        queueIndices.graphicsFamily.value(),
                        queueIndices.presentFamily.value());

        // Create ray tracing resources
        rayTracing.createCommandPool(device, queueIndices.graphicsFamily.value());
        rayTracing.createStorageImage(physicalDevice, device, WIDTH, HEIGHT,
                                     VK_FORMAT_B8G8R8A8_UNORM);  // Display output (matches swapchain)
        rayTracing.createAccumulationImage(physicalDevice, device, WIDTH, HEIGHT);  // HDR accumulation (R32G32B32A32_SFLOAT)
        rayTracing.createDepthBuffer(physicalDevice, device, WIDTH, HEIGHT);  // Depth for temporal rejection
        rayTracing.createVelocityBuffer(physicalDevice, device, WIDTH, HEIGHT);  // Motion vectors
        rayTracing.createCommandBuffers(device, swapchain.images.size());
        rayTracing.initializeImageLayouts(device, graphicsQueue);  // Initialize to GENERAL layout

        // Create acceleration structures
        acceleration.loadRayTracingFunctions(device);
        acceleration.createScene(physicalDevice, device, graphicsQueue, rayTracing.commandPool);
        acceleration.createTopLevelAS(physicalDevice, device, graphicsQueue, rayTracing.commandPool);

        // Allocate scratch buffer for dynamic TLAS updates
        acceleration.allocateTLASScratchBuffer(physicalDevice, device);

        // Create entities matching the geometry objects
        // Ground plane should be at Y=0
        auto groundEntity = entityManager.createEntity("Ground", 0, EntityType::STATIC);
        groundEntity->transform.position = glm::vec3(0.0f, 0.0f, 0.0f);

        // Cubes on ground (Y axis now fixed)
        auto cube1 = entityManager.createEntity("Cube1", 1, EntityType::DYNAMIC);
        cube1->transform.position = glm::vec3(0.0f, 0.5f, 0.0f);  // On ground
        cube1->transform.scale = glm::vec3(1.0f);
        cube1->angularVelocity = glm::vec3(0.0f, 1.0f, 0.0f);  // Spin

        auto cube2 = entityManager.createEntity("Cube2", 2, EntityType::DYNAMIC);
        cube2->transform.position = glm::vec3(4.0f, 0.4f, 0.0f);  // On ground
        cube2->transform.scale = glm::vec3(0.8f);
        cube2->velocity = glm::vec3(1.0f, 0.0f, 0.0f);  // Move along X

        auto cube3 = entityManager.createEntity("Cube3", 3, EntityType::STATIC);
        cube3->transform.position = glm::vec3(-4.0f, 0.6f, 0.0f);  // On ground
        cube3->transform.scale = glm::vec3(1.2f);

        auto cube4 = entityManager.createEntity("Cube4", 4, EntityType::DYNAMIC);
        cube4->transform.position = glm::vec3(0.0f, 0.75f, -5.0f);  // On ground
        cube4->transform.scale = glm::vec3(1.5f);
        cube4->angularVelocity = glm::vec3(0.5f, 0.5f, 0.0f);  // Tumble

        std::cout << "Created " << entityManager.size() << " entities ("
                  << entityManager.countDynamic() << " dynamic, "
                  << entityManager.countStatic() << " static)\n";

        // DEBUG: Print entity positions
        std::cout << "\n=== ENTITY POSITIONS (for debugging) ===\n";
        for (const auto& entity : entityManager.getEntities()) {
            std::cout << entity->name << ": pos("
                     << entity->transform.position.x << ", "
                     << entity->transform.position.y << ", "
                     << entity->transform.position.z << ") scale("
                     << entity->transform.scale.x << ")\n";
        }
        std::cout << "Ground plane geometry has Y=0.0 in vertices\n";
        std::cout << "Camera starts at: (0, 3, 8) - Y AXIS FIXED!\n";
        std::cout << "=======================================\n\n";

        // Create RT pipeline
        rtPipeline.loadRTPipelineFunctions(device);
        rtPipeline.createCameraBuffer(physicalDevice, device);
        rtPipeline.createDescriptorSetLayout(device);
        rtPipeline.createDescriptorPool(device);
        rtPipeline.createDescriptorSet(device, acceleration.tlas.handle,
                                       rayTracing.storageImageView, rayTracing.accumulationImageView,
                                       rayTracing.depthImageView, rayTracing.prevDepthImageView,
                                       rayTracing.velocityImageView);

        // Get shader path relative to project root
        std::string projectRoot = getProjectRoot();
        std::string shaderPath = projectRoot + "/shaders/compiled";
        std::cout << "Loading shaders from: " << shaderPath << "\n";

        rtPipeline.createPipeline(device, shaderPath);
        rtPipeline.createShaderBindingTable(physicalDevice, device);
        rtPipeline.createTonemapPipeline(device, shaderPath);
        rtPipeline.createTonemapDescriptorSet(device, rayTracing.accumulationImageView, rayTracing.storageImageView);

        // Create sync objects
        createSyncObjects();

        std::cout << "\n========================================\n";
        std::cout << "  Vulkan Ray Tracing Initialized!\n";
        std::cout << "  [SYSTEM CHECK] EXCLUSIVE PROCESSING ACTIVE\n";
        std::cout << "  > Rasterization Pipeline: DISABLED\n";
        std::cout << "  > Optimization: EXCLUSIVE RT CORE EXECUTION\n";
        std::cout << "  > Tensor Cores: DEDICATED PIPELINE (Denoiser)\n";
        std::cout << "  Ready to render!\n";
        std::cout << "========================================\n";
    }

    void createInstance() {
        if (enableValidationLayers && !checkValidationLayerSupport()) {
            throw std::runtime_error("Validation layers requested, but not available!");
        }

        VkApplicationInfo appInfo{};
        appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
        appInfo.pApplicationName = "Racing Engine";
        appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0);
        appInfo.pEngineName = "Custom Racing Engine";
        appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0);
        appInfo.apiVersion = VK_API_VERSION_1_3;

        VkInstanceCreateInfo createInfo{};
        createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
        createInfo.pApplicationInfo = &appInfo;

        auto extensions = getRequiredExtensions();
        createInfo.enabledExtensionCount = static_cast<uint32_t>(extensions.size());
        createInfo.ppEnabledExtensionNames = extensions.data();

        if (enableValidationLayers) {
            createInfo.enabledLayerCount = static_cast<uint32_t>(validationLayers.size());
            createInfo.ppEnabledLayerNames = validationLayers.data();
        } else {
            createInfo.enabledLayerCount = 0;
        }

        if (vkCreateInstance(&createInfo, nullptr, &instance) != VK_SUCCESS) {
            throw std::runtime_error("Failed to create Vulkan instance!");
        }

        std::cout << "Vulkan instance created successfully!\n";
    }

    void createSurface() {
        if (glfwCreateWindowSurface(instance, window, nullptr, &surface) != VK_SUCCESS) {
            throw std::runtime_error("Failed to create window surface!");
        }
        std::cout << "Window surface created!\n";
    }

    void pickPhysicalDevice() {
        uint32_t deviceCount = 0;
        vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);

        if (deviceCount == 0) {
            throw std::runtime_error("Failed to find GPUs with Vulkan support!");
        }

        std::vector<VkPhysicalDevice> devices(deviceCount);
        vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());

        for (const auto& dev : devices) {
            if (isDeviceSuitable(dev)) {
                physicalDevice = dev;
                break;
            }
        }

        if (physicalDevice == VK_NULL_HANDLE) {
            throw std::runtime_error("Failed to find a suitable NVIDIA GPU! This engine is for NVIDIA elites only. No AMD allowed.");
        }

        VkPhysicalDeviceProperties deviceProperties;
        vkGetPhysicalDeviceProperties(physicalDevice, &deviceProperties);
        std::cout << "\n==================================================\n";
        std::cout << "  [ELITIST MODE] HARDWARE VERIFIED\n";
        std::cout << "  GPU: " << deviceProperties.deviceName << "\n";
        std::cout << "  ARCH: NVIDIA RTX (Ampere/Ada Lovelace/Blackwell)\n";
        std::cout << "  RT CORES: EXCLUSIVE ACCESS ENABLED\n";
        std::cout << "  TENSOR CORES: PIPELINE READY\n";
        std::cout << "==================================================\n\n";
    }

    void createLogicalDevice() {
        queueIndices = findQueueFamilies(physicalDevice);

        std::vector<VkDeviceQueueCreateInfo> queueCreateInfos;
        std::set<uint32_t> uniqueQueueFamilies = {
            queueIndices.graphicsFamily.value(),
            queueIndices.presentFamily.value()
        };

        float queuePriority = 1.0f;
        for (uint32_t queueFamily : uniqueQueueFamilies) {
            VkDeviceQueueCreateInfo queueCreateInfo{};
            queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
            queueCreateInfo.queueFamilyIndex = queueFamily;
            queueCreateInfo.queueCount = 1;
            queueCreateInfo.pQueuePriorities = &queuePriority;
            queueCreateInfos.push_back(queueCreateInfo);
        }

        VkPhysicalDeviceFeatures deviceFeatures{};

        // Ray tracing features
        // === NVIDIA RTX FULL HARDWARE ACCELERATION ===

        // Buffer device address (required for RT)
        VkPhysicalDeviceBufferDeviceAddressFeatures bufferDeviceAddressFeatures{};
        bufferDeviceAddressFeatures.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES;
        bufferDeviceAddressFeatures.bufferDeviceAddress = VK_TRUE;

        // Vulkan 1.2 features for SPIRV 1.4
        VkPhysicalDeviceVulkan12Features vulkan12Features{};
        vulkan12Features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES;
        vulkan12Features.bufferDeviceAddress = VK_TRUE;
        vulkan12Features.shaderFloat16 = VK_TRUE;           // FP16 for Tensor Core efficiency
        vulkan12Features.shaderInt8 = VK_TRUE;              // INT8 for Tensor Cores
        vulkan12Features.storageBuffer8BitAccess = VK_TRUE;
        vulkan12Features.pNext = &bufferDeviceAddressFeatures;

        // Ray Query - inline ray tracing in any shader stage (uses RT Cores)
        VkPhysicalDeviceRayQueryFeaturesKHR rayQueryFeatures{};
        rayQueryFeatures.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_QUERY_FEATURES_KHR;
        rayQueryFeatures.rayQuery = VK_TRUE;
        rayQueryFeatures.pNext = &vulkan12Features;

        // RT Pipeline features - full RT Core utilization
        VkPhysicalDeviceRayTracingPipelineFeaturesKHR rtPipelineFeatures{};
        rtPipelineFeatures.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_PIPELINE_FEATURES_KHR;
        rtPipelineFeatures.rayTracingPipeline = VK_TRUE;
        rtPipelineFeatures.rayTracingPipelineTraceRaysIndirect = VK_TRUE;  // GPU-driven RT
        rtPipelineFeatures.rayTraversalPrimitiveCulling = VK_TRUE;         // HW culling
        rtPipelineFeatures.pNext = &rayQueryFeatures;

        // RT Maintenance 1 - additional RT Core features
        VkPhysicalDeviceRayTracingMaintenance1FeaturesKHR rtMaintenance1{};
        rtMaintenance1.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_MAINTENANCE_1_FEATURES_KHR;
        rtMaintenance1.rayTracingMaintenance1 = VK_TRUE;
        rtMaintenance1.rayTracingPipelineTraceRaysIndirect2 = VK_TRUE;
        rtMaintenance1.pNext = &rtPipelineFeatures;

        // Acceleration Structure features
        VkPhysicalDeviceAccelerationStructureFeaturesKHR accelFeatures{};
        accelFeatures.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_FEATURES_KHR;
        accelFeatures.accelerationStructure = VK_TRUE;
        accelFeatures.accelerationStructureCaptureReplay = VK_TRUE;
        accelFeatures.descriptorBindingAccelerationStructureUpdateAfterBind = VK_TRUE;
        accelFeatures.pNext = &rtMaintenance1;

        VkDeviceCreateInfo createInfo{};
        createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
        createInfo.queueCreateInfoCount = static_cast<uint32_t>(queueCreateInfos.size());
        createInfo.pQueueCreateInfos = queueCreateInfos.data();
        createInfo.pEnabledFeatures = &deviceFeatures;
        createInfo.enabledExtensionCount = static_cast<uint32_t>(deviceExtensions.size());
        createInfo.ppEnabledExtensionNames = deviceExtensions.data();
        createInfo.pNext = &accelFeatures;

        if (enableValidationLayers) {
            createInfo.enabledLayerCount = static_cast<uint32_t>(validationLayers.size());
            createInfo.ppEnabledLayerNames = validationLayers.data();
        } else {
            createInfo.enabledLayerCount = 0;
        }

        if (vkCreateDevice(physicalDevice, &createInfo, nullptr, &device) != VK_SUCCESS) {
            throw std::runtime_error("Failed to create logical device!");
        }

        vkGetDeviceQueue(device, queueIndices.graphicsFamily.value(), 0, &graphicsQueue);
        vkGetDeviceQueue(device, queueIndices.presentFamily.value(), 0, &presentQueue);

        std::cout << "Logical device created with ray tracing extensions!\n";
    }

    bool isDeviceSuitable(VkPhysicalDevice dev) {
        QueueFamilyIndices indices = findQueueFamilies(dev);

        bool extensionsSupported = checkDeviceExtensionSupport(dev);

        VkPhysicalDeviceProperties deviceProperties;
        vkGetPhysicalDeviceProperties(dev, &deviceProperties);

        // --- ELITIST CHECK: ENFORCE NVIDIA HARDWARE ---
        if (deviceProperties.vendorID != 0x10DE) {
            std::cout << "[REJECTED] Non-NVIDIA GPU detected: " << deviceProperties.deviceName
                      << " (VendorID: " << std::hex << deviceProperties.vendorID << std::dec << ")\n";
            std::cout << "           This engine is configured for FULL NVIDIA ELITIST MODE.\n";
            std::cout << "           AMD/Intel compatibility is explicitly disabled.\n";
            return false;
        }

        return indices.isComplete() && extensionsSupported;
    }

    bool checkDeviceExtensionSupport(VkPhysicalDevice dev) {
        uint32_t extensionCount;
        vkEnumerateDeviceExtensionProperties(dev, nullptr, &extensionCount, nullptr);

        std::vector<VkExtensionProperties> availableExtensions(extensionCount);
        vkEnumerateDeviceExtensionProperties(dev, nullptr, &extensionCount, availableExtensions.data());

        std::set<std::string> requiredExtensions(deviceExtensions.begin(), deviceExtensions.end());

        for (const auto& extension : availableExtensions) {
            requiredExtensions.erase(extension.extensionName);
        }

        return requiredExtensions.empty();
    }

    QueueFamilyIndices findQueueFamilies(VkPhysicalDevice dev) {
        QueueFamilyIndices indices;

        uint32_t queueFamilyCount = 0;
        vkGetPhysicalDeviceQueueFamilyProperties(dev, &queueFamilyCount, nullptr);

        std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
        vkGetPhysicalDeviceQueueFamilyProperties(dev, &queueFamilyCount, queueFamilies.data());

        int i = 0;
        for (const auto& queueFamily : queueFamilies) {
            if (queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT) {
                indices.graphicsFamily = i;
            }

            VkBool32 presentSupport = false;
            vkGetPhysicalDeviceSurfaceSupportKHR(dev, i, surface, &presentSupport);

            if (presentSupport) {
                indices.presentFamily = i;
            }

            if (indices.isComplete()) {
                break;
            }

            i++;
        }

        return indices;
    }

    std::vector<const char*> getRequiredExtensions() {
        uint32_t glfwExtensionCount = 0;
        const char** glfwExtensions;
        glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);

        std::vector<const char*> extensions(glfwExtensions, glfwExtensions + glfwExtensionCount);

        if (enableValidationLayers) {
            extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
        }

        return extensions;
    }

    void setupDebugMessenger() {
        if (!enableValidationLayers) return;

        VkDebugUtilsMessengerCreateInfoEXT createInfo{};
        createInfo.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
        createInfo.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT |
                                     VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT |
                                     VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;
        createInfo.messageType = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT |
                                 VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT |
                                 VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;
        createInfo.pfnUserCallback = debugCallback;

        auto func = (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, "vkCreateDebugUtilsMessengerEXT");
        if (func != nullptr) {
            func(instance, &createInfo, nullptr, &debugMessenger);
            std::cout << "Debug messenger created\n";
        }
    }

    void destroyDebugMessenger() {
        if (!enableValidationLayers || debugMessenger == VK_NULL_HANDLE) return;

        auto func = (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance, "vkDestroyDebugUtilsMessengerEXT");
        if (func != nullptr) {
            func(instance, debugMessenger, nullptr);
        }
    }

    bool checkValidationLayerSupport() {
        uint32_t layerCount;
        vkEnumerateInstanceLayerProperties(&layerCount, nullptr);

        std::vector<VkLayerProperties> availableLayers(layerCount);
        vkEnumerateInstanceLayerProperties(&layerCount, availableLayers.data());

        for (const char* layerName : validationLayers) {
            bool layerFound = false;

            for (const auto& layerProperties : availableLayers) {
                if (strcmp(layerName, layerProperties.layerName) == 0) {
                    layerFound = true;
                    break;
                }
            }

            if (!layerFound) {
                return false;
            }
        }

        return true;
    }

    void createSyncObjects() {
        imageAvailableSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
        renderFinishedSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
        inFlightFences.resize(MAX_FRAMES_IN_FLIGHT);
        imagesInFlight.resize(swapchain.images.size(), VK_NULL_HANDLE);  // One per swapchain image

        VkSemaphoreCreateInfo semaphoreInfo{};
        semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;

        VkFenceCreateInfo fenceInfo{};
        fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
        fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;

        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
            if (vkCreateSemaphore(device, &semaphoreInfo, nullptr, &imageAvailableSemaphores[i]) != VK_SUCCESS ||
                vkCreateSemaphore(device, &semaphoreInfo, nullptr, &renderFinishedSemaphores[i]) != VK_SUCCESS ||
                vkCreateFence(device, &fenceInfo, nullptr, &inFlightFences[i]) != VK_SUCCESS) {
                throw std::runtime_error("Failed to create sync objects!");
            }
        }

        std::cout << "Sync objects created (" << MAX_FRAMES_IN_FLIGHT << " frames in flight, "
                  << swapchain.images.size() << " swapchain images)\n";

        // Initialize OptiX AI Denoiser (Tensor Cores)
#ifdef CUDA_AVAILABLE
        if (NvidiaDenoiserImpl::checkSupport(physicalDevice)) {
            if (denoiser.initialize(physicalDevice, device, instance, WIDTH, HEIGHT)) {
                denoiser.createSharedBuffers(physicalDevice, device);
                std::cout << "\n========================================\n";
                std::cout << "  NVIDIA Tensor Cores: ACTIVE\n";
                std::cout << "  OptiX AI Denoiser: READY\n";
                std::cout << "========================================\n";
            }
        }
#endif
    }

    void mainLoop() {
        std::cout << "\n===========================================\n";
        std::cout << "  Racing Engine Ready!\n";
        std::cout << "  Rendering with Ray Tracing!\n";
        std::cout << "  Press ESC or close window to exit\n";
        std::cout << "===========================================\n\n";

        // Print initial debug info
        std::cout << "\n=== DEBUG: Initial Scene State ===\n";
        glm::vec3 camPos = camera.getPosition();
        glm::vec3 camFront = camera.getFront();
        std::cout << "Camera Position: ("
                  << camPos.x << ", "
                  << camPos.y << ", "
                  << camPos.z << ")\n";
        std::cout << "Camera Front: ("
                  << camFront.x << ", "
                  << camFront.y << ", "
                  << camFront.z << ")\n";
        std::cout << "Camera Pitch: " << camera.getPitch() << " degrees\n";
        std::cout << "Camera Yaw: " << camera.getYaw() << " degrees\n\n";

        for (const auto& entity : entityManager.getEntities()) {
            std::cout << entity->name << ": pos("
                      << entity->transform.position.x << ", "
                      << entity->transform.position.y << ", "
                      << entity->transform.position.z << ") scale("
                      << entity->transform.scale.x << ", "
                      << entity->transform.scale.y << ", "
                      << entity->transform.scale.z << ")\n";

            // Print 3x4 transform matrix
            float mat[12];
            entity->transform.getTransform3x4(mat);
            std::cout << "  Transform3x4:\n";
            std::cout << "    [" << mat[0] << ", " << mat[1] << ", " << mat[2] << ", " << mat[3] << "]\n";
            std::cout << "    [" << mat[4] << ", " << mat[5] << ", " << mat[6] << ", " << mat[7] << "]\n";
            std::cout << "    [" << mat[8] << ", " << mat[9] << ", " << mat[10] << ", " << mat[11] << "]\n";
        }
        std::cout << "==================================\n\n";

        lastFrameTime = std::chrono::high_resolution_clock::now();
        int frameCount = 0;

        // FPS tracking
        auto fpsStartTime = std::chrono::high_resolution_clock::now();
        int fpsFrameCount = 0;
        // Members currentFPS and frameTimeMs are used
        float titleUpdateTimer = 0.0f;

        // Initialize temporal accumulation tracking
        lastCameraPos = camera.getPosition();
        lastCameraFront = camera.getFront();

        std::cout << "\n==========================================\n";
        std::cout << "  IZTAPALAPA EDITION v2.0 - NVIDIA ELITIST\n";
        std::cout << "==========================================\n";
        std::cout << "GPU Vendor: NVIDIA ONLY (0x10DE)\n";
        std::cout << "RT Cores: FULLY UTILIZED (VK_KHR_ray_tracing)\n";
        std::cout << "Tensor Cores: ENABLED (FP16/INT8 + AI Denoiser)\n";
        std::cout << "AMD Compatibility: DISABLED (Hardcoded rejection)\n";
        std::cout << "------------------------------------------\n";
        std::cout << "Resolution: " << WIDTH << "x" << HEIGHT << "\n";
        std::cout << "Samples: 8 SPP + R2 quasi-random\n";
        std::cout << "Bounces: 8 max depth\n";
        std::cout << "RNG: PCG (high quality)\n";
        std::cout << "BRDF: GGX + Fresnel + Smith G\n";
        std::cout << "Materials: Metallic + Roughness PBR\n";
        std::cout << "Lighting: Physical sky + Sun + NEE\n";
        std::cout << "Post-FX: Bloom + Vignette + Grain\n";
        std::cout << "Tonemapping: ACES Filmic + sRGB\n";
        std::cout << "==========================================\n";
        std::cout << "  DEJANDO EN RIDICULO A LOS AAA\n";
        std::cout << "==========================================\n\n";

        while (!glfwWindowShouldClose(window)) {
            // Calculate delta time and frame time
            auto currentTime = std::chrono::high_resolution_clock::now();
            deltaTime = std::chrono::duration<float, std::chrono::seconds::period>(currentTime - lastFrameTime).count();
            frameTimeMs = deltaTime * 1000.0f;
            lastFrameTime = currentTime;

            // Calculate FPS
            fpsFrameCount++;
            auto fpsDuration = std::chrono::duration<float>(currentTime - fpsStartTime).count();
            if (fpsDuration >= 1.0f) {
                currentFPS = fpsFrameCount / fpsDuration;
                fpsFrameCount = 0;
                fpsStartTime = currentTime;
            }

            // Process input
            glfwPollEvents();
            camera.processKeyboard(window, deltaTime);
            camera.update(deltaTime);

            // Check if camera moved - reset accumulation if so
            glm::vec3 currentCameraPos = camera.getPosition();
            glm::vec3 currentCameraFront = camera.getFront();
            bool cameraMoved = glm::length(currentCameraPos - lastCameraPos) > 0.001f ||
                              glm::length(currentCameraFront - lastCameraFront) > 0.001f;

            if (cameraMoved) {
                accumulationFrames = 0;
                lastCameraPos = currentCameraPos;
                lastCameraFront = currentCameraFront;
            }
            accumulationFrames++;

            // Update entities
            entityManager.update(deltaTime);

            // Prepare TLAS instance data (CPU-side, before command buffer recording)
            if (entityManager.countDynamic() > 0) {
                const auto& transforms = entityManager.getTransforms();
                acceleration.prepareInstanceData(device, transforms);
            }

            // Update window title with stats every 0.25s
            titleUpdateTimer += deltaTime;
            if (titleUpdateTimer >= 0.25f) {
                updateWindowTitle();
                titleUpdateTimer = 0.0f;
            }

#if ENABLE_RENDER_LOOP_LOGGING
            // Print detailed status to console periodically
            if (frameCount % 120 == 0) {
                std::stringstream ss;
                ss << "\r[Frame " << frameCount << "] "
                   << "FPS: " << (int)currentFPS << " | "
                   << "Frame Time: " << frameTimeMs << "ms | "
                   << "AccumFrames: " << accumulationFrames << " | "
                   << "Samples: " << (accumulationFrames * 8);
                logger.log(ss.str());
            }

            // Extra verbose debug every 500 frames
            if (frameCount % 500 == 0 && frameCount > 0) {
                std::stringstream ss;
                ss << "\n[DEBUG Frame " << frameCount << "] "
                   << "GPU Memory OK | "
                   << "Sync objects valid | "
                   << "currentFrame=" << currentFrame
                   << "\n";
                logger.log(ss.str());
            }
#endif
            frameCount++;

            drawFrame();

            // Run denoiser periodically when camera is still (during accumulation)
            // Only denoise when we have accumulated enough samples and camera isn't moving
#ifdef CUDA_AVAILABLE
            if (denoiserEnabled && denoiser.initialized && !cameraMoved &&
                accumulationFrames > 1 && accumulationFrames % 32 == 0) {
                runDenoiser();
            }
#endif
        }

        vkDeviceWaitIdle(device);
    }

    // Debug tracking
    inline static uint32_t debugFrameNumber = 0;
    inline static uint32_t lastDebugPrintFrame = 0;
    static constexpr bool VERBOSE_DEBUG = false;  // Only print periodically
    static constexpr int DEBUG_PRINT_INTERVAL = 500;  // Print every N frames

#ifdef CUDA_AVAILABLE
    void runDenoiser() {
        // Run OptiX AI denoiser (Tensor Cores!)
        // Async execution - no CPU blocking!
        denoiser.denoiseImage(rayTracing.accumulationImage, rayTracing.accumulationImage,
                             device, graphicsQueue, rayTracing.commandPool);

#if ENABLE_RENDER_LOOP_LOGGING
            logger.log("\n[DENOISER] Frame denoised with Tensor Cores!\n");
#endif
    }
#endif

    void printDebugStep(const char* step, uint32_t frame, int stepNum) {
#if ENABLE_RENDER_LOOP_LOGGING
        if (VERBOSE_DEBUG || (frame - lastDebugPrintFrame >= DEBUG_PRINT_INTERVAL && stepNum == 0)) {
            std::stringstream ss;
            ss << "[F" << frame << "] " << step << "\n";
            logger.log(ss.str());
            if (stepNum == 0) lastDebugPrintFrame = frame;
        }
#endif
    }

    void drawFrame() {
        debugFrameNumber++;
        printDebugStep("=== FRAME START ===", debugFrameNumber, 0);

        // STEP 1: Wait for fence
        printDebugStep("Step 1: Waiting for inFlightFence", debugFrameNumber, 1);
        VkResult fenceResult = vkWaitForFences(device, 1, &inFlightFences[currentFrame], VK_TRUE, UINT64_MAX);
        if (fenceResult != VK_SUCCESS) {
            std::cerr << "[F" << debugFrameNumber << "] FENCE WAIT FAILED: " << fenceResult << std::endl;
        }

        // STEP 2: Acquire image
        printDebugStep("Step 2: Acquiring swapchain image", debugFrameNumber, 2);
        uint32_t imageIndex;
        VkResult result = vkAcquireNextImageKHR(device, swapchain.swapchain, UINT64_MAX,
                                               imageAvailableSemaphores[currentFrame], VK_NULL_HANDLE, &imageIndex);

        if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {
            std::cerr << "[F" << debugFrameNumber << "] ACQUIRE FAILED: " << result << std::endl;
            throw std::runtime_error("Failed to acquire swap chain image!");
        }

        // STEP 3: Wait for image-in-flight
        printDebugStep("Step 3: Checking imagesInFlight", debugFrameNumber, 3);
        if (imagesInFlight[imageIndex] != VK_NULL_HANDLE) {
            VkResult imgFenceResult = vkWaitForFences(device, 1, &imagesInFlight[imageIndex], VK_TRUE, UINT64_MAX);
            if (imgFenceResult != VK_SUCCESS) {
                std::cerr << "[F" << debugFrameNumber << "] IMAGE FENCE WAIT FAILED: " << imgFenceResult << std::endl;
            }
        }
        imagesInFlight[imageIndex] = inFlightFences[currentFrame];

        // STEP 4: Reset fence
        printDebugStep("Step 4: Resetting fence", debugFrameNumber, 4);
        vkResetFences(device, 1, &inFlightFences[currentFrame]);

        // STEP 5: Update camera UBO with temporal reprojection data
        printDebugStep("Step 5: Updating camera buffer", debugFrameNumber, 5);

        // Calculate current view-projection matrix
        glm::mat4 currentView = camera.getViewMatrix();
        glm::mat4 currentProj = camera.getProjectionMatrix();
        glm::mat4 currentViewProj = currentProj * currentView;
        glm::vec3 currentCamPos = camera.getPosition();

        CameraUBO cameraUBO{};
        cameraUBO.viewInverse = camera.getViewInverse();
        cameraUBO.projInverse = camera.getProjInverse();
        cameraUBO.viewProj = currentViewProj;

        // Previous frame data for temporal reprojection
        cameraUBO.prevViewProj = hasPreviousFrame ? prevViewProj : currentViewProj;
        cameraUBO.cameraPos = glm::vec4(currentCamPos, 1.0f);
        cameraUBO.prevCameraPos = glm::vec4(hasPreviousFrame ? prevCameraPosition : currentCamPos, 1.0f);

        // Frame counters
        cameraUBO.frameCount = totalFrameCount;
        cameraUBO.accumulatedFrames = accumulationFrames;

        // Temporal blend settings (lower = more responsive, higher = smoother but more ghosting)
        cameraUBO.temporalBlendFactor = 0.1f;   // Blend 10% new, 90% history
        cameraUBO.motionBlurScale = 20.0f;       // Rejection threshold for motion vectors

        rtPipeline.updateCameraBuffer(device, cameraUBO);

        // Store current frame data for next frame's reprojection
        prevViewProj = currentViewProj;
        prevCameraPosition = currentCamPos;
        hasPreviousFrame = true;
        totalFrameCount++;

        // STEP 6: Begin command buffer
        printDebugStep("Step 6: Beginning command buffer", debugFrameNumber, 6);
        VkCommandBuffer commandBuffer = rayTracing.commandBuffers[imageIndex];
        VkResult resetResult = vkResetCommandBuffer(commandBuffer, 0);
        if (resetResult != VK_SUCCESS) {
            std::cerr << "[F" << debugFrameNumber << "] RESET CMD BUFFER FAILED: " << resetResult << std::endl;
        }

        VkCommandBufferBeginInfo beginInfo{};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        VkResult beginResult = vkBeginCommandBuffer(commandBuffer, &beginInfo);
        if (beginResult != VK_SUCCESS) {
            std::cerr << "[F" << debugFrameNumber << "] BEGIN CMD BUFFER FAILED: " << beginResult << std::endl;
        }

        // === STEP 0: Update TLAS for dynamic objects ===
        if (entityManager.countDynamic() > 0) {
            acceleration.recordTLASUpdate(commandBuffer);
        }

#if ENABLE_RENDER_LOOP_LOGGING
        // Reduced debug output - only on interval
        if (debugFrameNumber % DEBUG_PRINT_INTERVAL == 0) {
            std::cout << "[F" << debugFrameNumber << "] CMD BUFFER RECORDING - imageIndex: " << imageIndex << std::endl;
        }
#endif

        // === STEP 1: Synchronization barriers ===
        // On first frame, use TOP_OF_PIPE with 0 access (nothing to wait for)
        // After first frame, wait for previous frame's shaders/transfers

        // Accumulation image: GENERAL -> GENERAL (preserve content for temporal accumulation)
        VkImageMemoryBarrier accumulationBarrier{};
        accumulationBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        accumulationBarrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
        accumulationBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
        accumulationBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        accumulationBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        accumulationBarrier.image = rayTracing.accumulationImage;
        accumulationBarrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        accumulationBarrier.subresourceRange.baseMipLevel = 0;
        accumulationBarrier.subresourceRange.levelCount = 1;
        accumulationBarrier.subresourceRange.baseArrayLayer = 0;
        accumulationBarrier.subresourceRange.layerCount = 1;
        // First frame: no previous access to wait for. After: wait for previous shader access.
        accumulationBarrier.srcAccessMask = storageImageInitialized ? (VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT) : 0;
        accumulationBarrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;

        // Storage image: First use is GENERAL (from init), after that always TRANSFER_SRC
        VkImageMemoryBarrier storageBarrier{};
        storageBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        storageBarrier.oldLayout = storageImageInitialized ? VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL : VK_IMAGE_LAYOUT_GENERAL;
        storageBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
        storageBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        storageBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        storageBarrier.image = rayTracing.storageImage;
        storageBarrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        storageBarrier.subresourceRange.baseMipLevel = 0;
        storageBarrier.subresourceRange.levelCount = 1;
        storageBarrier.subresourceRange.baseArrayLayer = 0;
        storageBarrier.subresourceRange.layerCount = 1;
        storageBarrier.srcAccessMask = storageImageInitialized ? VK_ACCESS_TRANSFER_READ_BIT : 0;
        storageBarrier.dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT;

        // Proper stage masks for first frame (TOP_OF_PIPE) vs subsequent frames
        VkPipelineStageFlags srcStage = storageImageInitialized ?
            (VK_PIPELINE_STAGE_TRANSFER_BIT | VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR) :
            VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;

        VkImageMemoryBarrier barriers[2] = {accumulationBarrier, storageBarrier};
        vkCmdPipelineBarrier(commandBuffer,
                            srcStage,
                            VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                            0, 0, nullptr, 0, nullptr, 2, barriers);

        storageImageInitialized = true;  // After first use, it will be in TRANSFER_SRC

        // === STEP 2: Path Tracing (writes to accumulation image) ===
        vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, rtPipeline.pipeline);
        vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR,
                               rtPipeline.pipelineLayout, 0, 1, &rtPipeline.descriptorSet, 0, nullptr);

        rtPipeline.vkCmdTraceRaysKHR(commandBuffer,
                                    &rtPipeline.sbt.raygenRegion,
                                    &rtPipeline.sbt.missRegion,
                                    &rtPipeline.sbt.hitRegion,
                                    &rtPipeline.sbt.callableRegion,
                                    WIDTH, HEIGHT, 1);

        // === STEP 3: Barrier before compute shader + copy depth to prevDepth ===
        VkImageMemoryBarrier computeBarrier{};
        computeBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        computeBarrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
        computeBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
        computeBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        computeBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        computeBarrier.image = rayTracing.accumulationImage;
        computeBarrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        computeBarrier.subresourceRange.baseMipLevel = 0;
        computeBarrier.subresourceRange.levelCount = 1;
        computeBarrier.subresourceRange.baseArrayLayer = 0;
        computeBarrier.subresourceRange.layerCount = 1;
        computeBarrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
        computeBarrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

        // Depth buffer barrier for copy
        VkImageMemoryBarrier depthBarrier{};
        depthBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        depthBarrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
        depthBarrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        depthBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        depthBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        depthBarrier.image = rayTracing.depthImage;
        depthBarrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        depthBarrier.subresourceRange.baseMipLevel = 0;
        depthBarrier.subresourceRange.levelCount = 1;
        depthBarrier.subresourceRange.baseArrayLayer = 0;
        depthBarrier.subresourceRange.layerCount = 1;
        depthBarrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
        depthBarrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;

        // Previous depth barrier for copy destination
        VkImageMemoryBarrier prevDepthBarrier{};
        prevDepthBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        prevDepthBarrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
        prevDepthBarrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        prevDepthBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        prevDepthBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        prevDepthBarrier.image = rayTracing.prevDepthImage;
        prevDepthBarrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        prevDepthBarrier.subresourceRange.baseMipLevel = 0;
        prevDepthBarrier.subresourceRange.levelCount = 1;
        prevDepthBarrier.subresourceRange.baseArrayLayer = 0;
        prevDepthBarrier.subresourceRange.layerCount = 1;
        prevDepthBarrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
        prevDepthBarrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;

        VkImageMemoryBarrier rtBarriers[3] = {computeBarrier, depthBarrier, prevDepthBarrier};
        vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR,
                            VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT,
                            0, 0, nullptr, 0, nullptr, 3, rtBarriers);

        // Copy current depth to previous depth for next frame's temporal rejection
        VkImageCopy depthCopy{};
        depthCopy.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        depthCopy.srcSubresource.layerCount = 1;
        depthCopy.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        depthCopy.dstSubresource.layerCount = 1;
        depthCopy.extent = {WIDTH, HEIGHT, 1};

        vkCmdCopyImage(commandBuffer, rayTracing.depthImage, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                      rayTracing.prevDepthImage, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &depthCopy);

        // Transition both depth images back to GENERAL for next frame
        depthBarrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        depthBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
        depthBarrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
        depthBarrier.dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT;

        prevDepthBarrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        prevDepthBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
        prevDepthBarrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        prevDepthBarrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

        VkImageMemoryBarrier depthRestoreBarriers[2] = {depthBarrier, prevDepthBarrier};
        vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT,
                            VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR,
                            0, 0, nullptr, 0, nullptr, 2, depthRestoreBarriers);

        // === STEP 4: Tonemap compute shader (reads accumulation, writes storage) ===
        vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, rtPipeline.tonemapPipeline);
        vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE,
                               rtPipeline.tonemapPipelineLayout, 0, 1, &rtPipeline.tonemapDescriptorSet, 0, nullptr);

        // Dispatch compute shader (16x16 workgroups)
        uint32_t groupCountX = (WIDTH + 15) / 16;
        uint32_t groupCountY = (HEIGHT + 15) / 16;
        vkCmdDispatch(commandBuffer, groupCountX, groupCountY, 1);

        // === STEP 5: Transition storage image for transfer ===
        VkImageMemoryBarrier transferBarrier{};
        transferBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        transferBarrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
        transferBarrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        transferBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        transferBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        transferBarrier.image = rayTracing.storageImage;
        transferBarrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        transferBarrier.subresourceRange.baseMipLevel = 0;
        transferBarrier.subresourceRange.levelCount = 1;
        transferBarrier.subresourceRange.baseArrayLayer = 0;
        transferBarrier.subresourceRange.layerCount = 1;
        transferBarrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
        transferBarrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;

        vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                            VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &transferBarrier);

        // Transition swapchain image for transfer
        VkImageMemoryBarrier swapchainBarrier{};
        swapchainBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
        swapchainBarrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        swapchainBarrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        swapchainBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        swapchainBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        swapchainBarrier.image = swapchain.images[imageIndex];
        swapchainBarrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        swapchainBarrier.subresourceRange.baseMipLevel = 0;
        swapchainBarrier.subresourceRange.levelCount = 1;
        swapchainBarrier.subresourceRange.baseArrayLayer = 0;
        swapchainBarrier.subresourceRange.layerCount = 1;
        swapchainBarrier.srcAccessMask = 0;
        swapchainBarrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;

        vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                            VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &swapchainBarrier);

        // === STEP 6: Copy storage image to swapchain ===
        // Use minimum of storage image size and swapchain size
        uint32_t copyWidth = (WIDTH < swapchain.extent.width) ? WIDTH : swapchain.extent.width;
        uint32_t copyHeight = (HEIGHT < swapchain.extent.height) ? HEIGHT : swapchain.extent.height;

        VkImageCopy copyRegion{};
        copyRegion.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        copyRegion.srcSubresource.layerCount = 1;
        copyRegion.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        copyRegion.dstSubresource.layerCount = 1;
        copyRegion.extent = {copyWidth, copyHeight, 1};

        vkCmdCopyImage(commandBuffer, rayTracing.storageImage, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                      swapchain.images[imageIndex], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copyRegion);

        // === STEP 7: Transition swapchain for present ===
        swapchainBarrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        swapchainBarrier.newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
        swapchainBarrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        swapchainBarrier.dstAccessMask = 0;

        vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT,
                            VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, 0, 0, nullptr, 0, nullptr, 1, &swapchainBarrier);

        // STEP 7: End command buffer
        printDebugStep("Step 7: Ending command buffer", debugFrameNumber, 7);
        VkResult endResult = vkEndCommandBuffer(commandBuffer);
        if (endResult != VK_SUCCESS) {
            std::cerr << "[F" << debugFrameNumber << "] END CMD BUFFER FAILED: " << endResult << std::endl;
        }

        // STEP 8: Submit command buffer

        VkSubmitInfo submitInfo{};
        submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

        VkSemaphore waitSemaphores[] = {imageAvailableSemaphores[currentFrame]};
        VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
        submitInfo.waitSemaphoreCount = 1;
        submitInfo.pWaitSemaphores = waitSemaphores;
        submitInfo.pWaitDstStageMask = waitStages;
        submitInfo.commandBufferCount = 1;
        submitInfo.pCommandBuffers = &commandBuffer;

        VkSemaphore signalSemaphores[] = {renderFinishedSemaphores[currentFrame]};
        submitInfo.signalSemaphoreCount = 1;
        submitInfo.pSignalSemaphores = signalSemaphores;

        VkResult submitResult = vkQueueSubmit(graphicsQueue, 1, &submitInfo, inFlightFences[currentFrame]);
        if (submitResult != VK_SUCCESS) {
            std::cerr << "\n=== VULKAN SUBMIT ERROR ===" << std::endl;
            std::cerr << "INTERNAL FRAME NUMBER: " << debugFrameNumber << std::endl;
            std::cerr << "vkQueueSubmit failed with code: " << submitResult << std::endl;
            switch (submitResult) {
                case VK_ERROR_OUT_OF_HOST_MEMORY: std::cerr << "VK_ERROR_OUT_OF_HOST_MEMORY" << std::endl; break;
                case VK_ERROR_OUT_OF_DEVICE_MEMORY: std::cerr << "VK_ERROR_OUT_OF_DEVICE_MEMORY" << std::endl; break;
                case VK_ERROR_DEVICE_LOST: std::cerr << "VK_ERROR_DEVICE_LOST (GPU crashed or TDR timeout)" << std::endl; break;
                default: std::cerr << "Unknown error code: " << submitResult << std::endl; break;
            }
            std::cerr << "currentFrame index: " << currentFrame << std::endl;
            std::cerr << "imageIndex: " << imageIndex << std::endl;
            std::cerr << "accumulationFrames: " << accumulationFrames << std::endl;
            glm::vec3 pos = camera.getPosition();
            std::cerr << "Camera pos: (" << pos.x << ", " << pos.y << ", " << pos.z << ")" << std::endl;
            std::cerr << "deltaTime: " << deltaTime << std::endl;
            std::cerr << "storageImageInitialized: " << (storageImageInitialized ? "true" : "false") << std::endl;
            std::cerr << "===========================" << std::endl;

            // Don't try vkDeviceWaitIdle after DEVICE_LOST - it will also fail
            if (submitResult != VK_ERROR_DEVICE_LOST) {
                vkDeviceWaitIdle(device);
            }
            throw std::runtime_error("Failed to submit draw command buffer!");
        }

        // STEP 9: Present
        printDebugStep("Step 9: Presenting to screen", debugFrameNumber, 9);

        // Present
        VkPresentInfoKHR presentInfo{};
        presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
        presentInfo.waitSemaphoreCount = 1;
        presentInfo.pWaitSemaphores = signalSemaphores;

        VkSwapchainKHR swapChains[] = {swapchain.swapchain};
        presentInfo.swapchainCount = 1;
        presentInfo.pSwapchains = swapChains;
        presentInfo.pImageIndices = &imageIndex;

        vkQueuePresentKHR(presentQueue, &presentInfo);

        // Removed vkDeviceWaitIdle - rely on fence synchronization instead

        currentFrame = (currentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
    }

    void cleanup() {
        vkDeviceWaitIdle(device);

        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
            vkDestroySemaphore(device, renderFinishedSemaphores[i], nullptr);
            vkDestroySemaphore(device, imageAvailableSemaphores[i], nullptr);
            vkDestroyFence(device, inFlightFences[i], nullptr);
        }

        // Cleanup denoiser before Vulkan resources
#ifdef CUDA_AVAILABLE
        denoiser.cleanup(device);
#endif

        rtPipeline.cleanup(device);
        acceleration.cleanup(device);
        rayTracing.cleanup(device);
        swapchain.cleanup(device);

        vkDestroyDevice(device, nullptr);
        destroyDebugMessenger();
        vkDestroySurfaceKHR(instance, surface, nullptr);
        vkDestroyInstance(instance, nullptr);

        glfwDestroyWindow(window);
        glfwTerminate();

        std::cout << "Cleanup complete.\n";
    }
};

int main() {
    RacingEngine engine;

    try {
        engine.run();
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}