// Native Ray Reconstruction replay: runs REngine (engine/runtime_rr.inc) on a capture's game inputs, frame after frame,
// and writes every image the engine produces (REngine::debugImages) after each frame, to compare with the capture's
// per-stage surfaces. No CUDA, NGX or Wine. The inputs come from engine/rr/rr_inputs.py.
//
// usage: d4r-rr MODEL_DIR INPUT_DIR OUT_DIR
//   INPUT_DIR: frames.txt ("size RW RH OW OH", then one line per frame: "frame F jx jy mvx mvy preExposure reset
//   depthInverted" followed by 16 worldToView and 16 viewToClip floats) and F-color.bin (RGBA16F), F-motion.bin
//   (RG16F), F-depth.bin (R32F), F-normals.bin and F-roughness.bin (RGBA16F), F-diffuse.bin and F-specular.bin
//   (RGBA8 UNORM), all at the render size.
//   OUT_DIR receives F-NAME.bin for every debug image of frame F.
#include "runtime.h"
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <vector>

namespace {
void check(VkResult r, const char* operation) {
    if (r != VK_SUCCESS) throw std::runtime_error(std::string(operation) + ": VkResult " + std::to_string(r));
}
#define CK(call) check((call), #call)
std::vector<char> read(const std::string& file) {
    std::ifstream stream(file, std::ios::binary | std::ios::ate);
    if (!stream) throw std::runtime_error("cannot read " + file);
    std::vector<char> data(static_cast<size_t>(stream.tellg()));
    stream.seekg(0);
    stream.read(data.data(), static_cast<std::streamsize>(data.size()));
    return data;
}
uint32_t texelBytes(VkFormat f) {
    switch (f) {
    case VK_FORMAT_R16G16B16A16_SFLOAT: return 8;
    case VK_FORMAT_R16G16_SFLOAT: case VK_FORMAT_R32_SFLOAT: case VK_FORMAT_R8G8B8A8_UNORM: case VK_FORMAT_R8G8B8A8_SNORM: return 4;
    case VK_FORMAT_R16_SFLOAT: return 2;
    default: throw std::runtime_error("texelBytes: format");
    }
}
struct Ctx {
    VkPhysicalDevice physical;
    VkDevice device;
    VkQueue queue;
    VkCommandPool pool;
    VkPhysicalDeviceMemoryProperties memory;
    uint32_t type(uint32_t bits, VkMemoryPropertyFlags flags) const {
        for (uint32_t i = 0; i < memory.memoryTypeCount; ++i)
            if ((bits & (1u << i)) && (memory.memoryTypes[i].propertyFlags & flags) == flags) return i;
        throw std::runtime_error("memory type");
    }
    template <class F> void submit(F&& record) {
        VkCommandBufferAllocateInfo ai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO};
        ai.commandPool = pool; ai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; ai.commandBufferCount = 1;
        VkCommandBuffer cb; CK(vkAllocateCommandBuffers(device, &ai, &cb));
        VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO};
        CK(vkBeginCommandBuffer(cb, &bi)); record(cb); CK(vkEndCommandBuffer(cb));
        VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb;
        CK(vkQueueSubmit(queue, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(queue));
        vkFreeCommandBuffers(device, pool, 1, &cb);
    }
};
struct HostBuffer {
    VkBuffer buffer; VkDeviceMemory memory; void* mapped;
    HostBuffer(Ctx& c, VkDeviceSize size) {
        VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; bi.size = size;
        bi.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT;
        CK(vkCreateBuffer(c.device, &bi, nullptr, &buffer));
        VkMemoryRequirements r; vkGetBufferMemoryRequirements(c.device, buffer, &r);
        VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size;
        ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
        CK(vkAllocateMemory(c.device, &ai, nullptr, &memory)); CK(vkBindBufferMemory(c.device, buffer, memory, 0));
        CK(vkMapMemory(c.device, memory, 0, VK_WHOLE_SIZE, 0, &mapped));
    }
};
void copyRegion(VkBufferImageCopy& r, uint32_t w, uint32_t h) {
    r = {}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {w, h, 1};
}
struct Input {
    VkImage image; VkImageView view; VkFormat format;
    Input(Ctx& c, VkFormat f, uint32_t w, uint32_t h) : format(f) {
        VkImageCreateInfo ii{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO};
        ii.imageType = VK_IMAGE_TYPE_2D; ii.format = f; ii.extent = {w, h, 1}; ii.mipLevels = ii.arrayLayers = 1;
        ii.samples = VK_SAMPLE_COUNT_1_BIT; ii.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT;
        CK(vkCreateImage(c.device, &ii, nullptr, &image));
        VkMemoryRequirements r; vkGetImageMemoryRequirements(c.device, image, &r);
        VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size;
        ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        VkDeviceMemory m; CK(vkAllocateMemory(c.device, &ai, nullptr, &m)); CK(vkBindImageMemory(c.device, image, m, 0));
        VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO};
        vi.image = image; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = f; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
        CK(vkCreateImageView(c.device, &vi, nullptr, &view));
    }
};
void barrier(VkCommandBuffer cb, VkImage image, VkImageLayout from, VkImageLayout to) {
    VkImageMemoryBarrier b{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER};
    b.image = image; b.oldLayout = from; b.newLayout = to;
    b.srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT; b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT;
    b.srcQueueFamilyIndex = b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
    vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, nullptr, 0, nullptr, 1, &b);
}
} // namespace

int main(int argc, char** argv) try {
    if (argc != 4) { std::cerr << "usage: d4r-rr MODEL_DIR INPUT_DIR OUT_DIR\n"; return 2; }
    const std::string modelDir = argv[1], inputDir = argv[2], outDir = argv[3];
    VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; app.apiVersion = VK_API_VERSION_1_3;
    VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; ici.pApplicationInfo = &app;
    VkInstance instance; CK(vkCreateInstance(&ici, nullptr, &instance));
    uint32_t count = 8; VkPhysicalDevice devices[8]; vkEnumeratePhysicalDevices(instance, &count, devices);
    Ctx c{};
    c.physical = devices[0];
    for (uint32_t i = 0; i < count; ++i) {
        VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(devices[i], &p);
        if (p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) c.physical = devices[i];
    }
    vkGetPhysicalDeviceMemoryProperties(c.physical, &c.memory);
    VkPhysicalDeviceRobustness2FeaturesEXT rb2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ROBUSTNESS_2_FEATURES_EXT};
    VkPhysicalDeviceShaderFloat8FeaturesEXT f8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT, &rb2};
    VkPhysicalDeviceCooperativeMatrixFeaturesKHR cm{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR, &f8};
    VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &cm};
    VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &v13};
    VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12};
    VkPhysicalDeviceFeatures2 f2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v11};
    vkGetPhysicalDeviceFeatures2(c.physical, &f2);
    const char* extensions[] = {VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME, VK_EXT_SHADER_FLOAT8_EXTENSION_NAME,
                                VK_EXT_ROBUSTNESS_2_EXTENSION_NAME};
    uint32_t families = 8; VkQueueFamilyProperties qf[8]; vkGetPhysicalDeviceQueueFamilyProperties(c.physical, &families, qf);
    uint32_t family = 0;
    for (uint32_t i = 0; i < families; ++i) if (qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) { family = i; break; }
    float priority = 1;
    VkDeviceQueueCreateInfo qi{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qi.queueFamilyIndex = family; qi.queueCount = 1; qi.pQueuePriorities = &priority;
    VkDeviceCreateInfo di{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; di.queueCreateInfoCount = 1; di.pQueueCreateInfos = &qi;
    // D4R_RR_TEST_ROBUST2=1: robustBufferAccess2 / robustImageAccess2 as vkd3d-proton enables them (out-of-bounds
    // reads return zero)
    const bool robust2 = std::getenv("D4R_RR_TEST_ROBUST2") != nullptr;
    if (!robust2) { rb2.robustBufferAccess2 = rb2.robustImageAccess2 = rb2.nullDescriptor = VK_FALSE; }
    di.enabledExtensionCount = robust2 ? 3 : 2; di.ppEnabledExtensionNames = extensions;
    // D4R_RR_TEST_NO_ROBUST=1: no robust buffer or image access (out-of-bounds reads are undefined)
    if (std::getenv("D4R_RR_TEST_NO_ROBUST")) { f2.features.robustBufferAccess = VK_FALSE; v13.robustImageAccess = VK_FALSE; }
    // D4R_RR_TEST_NO_COOPMAT=1: a device like vkd3d-proton's, without the cooperative-matrix and float8 features
    if (std::getenv("D4R_RR_TEST_NO_COOPMAT")) { di.enabledExtensionCount = 0; v13.pNext = nullptr; }
    CK(vkCreateDevice(c.physical, &di, nullptr, &c.device));
    vkGetDeviceQueue(c.device, family, 0, &c.queue);
    VkCommandPoolCreateInfo pi{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; pi.queueFamilyIndex = family;
    pi.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    CK(vkCreateCommandPool(c.device, &pi, nullptr, &c.pool));

    auto framesText = read(inputDir + "/frames.txt");
    std::istringstream frames(std::string(framesText.begin(), framesText.end()));
    std::string key;
    uint32_t rw, rh, ow, oh;
    frames >> key >> rw >> rh >> ow >> oh;
    d4r::RModel model(c.physical, c.device, modelDir);
    d4r::REngine engine(model, rw, rh, ow, oh);
    const char* names[] = {"color", "motion", "depth", "normals", "roughness", "diffuse", "specular"};
    const VkFormat formats[] = {VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R16G16_SFLOAT, VK_FORMAT_R32_SFLOAT,
                                VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R8G8B8A8_UNORM,
                                VK_FORMAT_R8G8B8A8_UNORM};
    std::vector<Input> inputs;
    for (VkFormat f : formats) inputs.emplace_back(c, f, rw, rh);
    const char* optionalNames[] = {"hitdistance", "disocclusion"};
    const VkFormat optionalFormats[] = {VK_FORMAT_R32_SFLOAT, VK_FORMAT_R16_SFLOAT};
    std::vector<Input> optional;
    for (VkFormat f : optionalFormats) optional.emplace_back(c, f, rw, rh);
    HostBuffer staging(c, VkDeviceSize(ow) * oh * 8);
    // GPU time of every frame: REngine's four stamps (start, after exposure, after the network, after the output)
    VkQueryPoolCreateInfo qp{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO};
    qp.queryType = VK_QUERY_TYPE_TIMESTAMP; qp.queryCount = 4;
    VkQueryPool stamps; CK(vkCreateQueryPool(c.device, &qp, nullptr, &stamps));
    VkPhysicalDeviceProperties props; vkGetPhysicalDeviceProperties(c.physical, &props);
    const int repeat = std::getenv("D4R_RR_TIME_REPEAT") ? std::atoi(std::getenv("D4R_RR_TIME_REPEAT")) : 1;
    bool first = true;
    while (frames >> key) {
        uint32_t f;
        d4r::RFrameSettings s;
        int reset, inverted;
        frames >> f >> s.jitter[0] >> s.jitter[1] >> s.motionScale[0] >> s.motionScale[1] >> s.preExposure >> reset >> inverted;
        for (float& v : s.worldToView) frames >> v;
        for (float& v : s.viewToClip) frames >> v;
        s.reset = reset != 0; s.depthInverted = inverted != 0;
        for (size_t i = 0; i < inputs.size(); ++i) {
            auto data = read(inputDir + "/" + std::to_string(f) + "-" + names[i] + ".bin");
            if (data.size() != size_t(rw) * rh * texelBytes(formats[i])) throw std::runtime_error(std::string("input size: ") + names[i]);
            std::memcpy(staging.mapped, data.data(), data.size());
            c.submit([&](VkCommandBuffer cb) {
                barrier(cb, inputs[i].image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
                VkBufferImageCopy r; copyRegion(r, rw, rh);
                vkCmdCopyBufferToImage(cb, staging.buffer, inputs[i].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &r);
                barrier(cb, inputs[i].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL);
            });
        }
        d4r::RFrameImages images{};
        images.color = inputs[0].view; images.motion = inputs[1].view; images.depth = inputs[2].view;
        images.normals = inputs[3].view; images.roughness = inputs[4].view;
        images.diffuseAlbedo = inputs[5].view; images.specularAlbedo = inputs[6].view;
        // optional inputs, present when the capture had them (rr_inputs.py)
        for (int k = 0; k < 2; ++k) {
            std::ifstream probe(inputDir + "/" + std::to_string(f) + "-" + optionalNames[k] + ".bin", std::ios::binary);
            if (!probe) continue;
            auto data = read(inputDir + "/" + std::to_string(f) + "-" + optionalNames[k] + ".bin");
            if (data.size() != size_t(rw) * rh * texelBytes(optionalFormats[k])) throw std::runtime_error(std::string("input size: ") + optionalNames[k]);
            std::memcpy(staging.mapped, data.data(), data.size());
            c.submit([&](VkCommandBuffer cb) {
                barrier(cb, optional[k].image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
                VkBufferImageCopy region; copyRegion(region, rw, rh);
                vkCmdCopyBufferToImage(cb, staging.buffer, optional[k].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);
                barrier(cb, optional[k].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL);
            });
            (k == 0 ? images.specularHitDistance : images.disocclusionMask) = optional[k].view;
        }
        engine.setFrame(images, s);
        if (first) c.submit([&](VkCommandBuffer cb) { model.recordUpload(cb); engine.recordInitialize(cb); first = false; });
        double best = 1e30, network = 0, output = 0, exposure = 0;
        for (int r = 0; r < repeat; ++r) {
            if (r) engine.setFrame(images, s);
            c.submit([&](VkCommandBuffer cb) {
                vkCmdResetQueryPool(cb, stamps, 0, 4);
                engine.recordFrame(cb, stamps, 0);
            });
            uint64_t t[4];
            CK(vkGetQueryPoolResults(c.device, stamps, 0, 4, sizeof(t), t, 8, VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT));
            const double ms = props.limits.timestampPeriod * 1e-6;
            if ((t[3] - t[0]) * ms < best) {
                best = (t[3] - t[0]) * ms; exposure = (t[1] - t[0]) * ms; network = (t[2] - t[1]) * ms; output = (t[3] - t[2]) * ms;
            }
        }
        std::cout << "frame " << f << ": " << best << " ms (exposure " << exposure << ", network " << network
                  << ", output " << output << ")\n";
        for (const auto& d : engine.debugImages()) {
            const size_t bytes = size_t(d.width) * d.height * texelBytes(d.format);
            c.submit([&](VkCommandBuffer cb) {
                VkBufferImageCopy r; copyRegion(r, d.width, d.height);
                vkCmdCopyImageToBuffer(cb, d.image, VK_IMAGE_LAYOUT_GENERAL, staging.buffer, 1, &r);
            });
            std::ofstream(outDir + "/" + std::to_string(f) + "-" + d.name + ".bin", std::ios::binary)
                .write(static_cast<const char*>(staging.mapped), static_cast<std::streamsize>(bytes));
        }
        std::cout << "frame " << f << " done\n";
    }
    return 0;
} catch (const std::exception& e) {
    std::cerr << "d4r-rr: " << e.what() << "\n";
    return 1;
}