#include <vulkan/vulkan.hpp>
#include <iostream>
#include <vector>
#include <fstream>

// --- ШАГ 0: ШЕЙДЕР ---
// Это микро-программа, которая будет работать прямо внутри видеокарты.
// В обучении нейронки один шейдер будет делать FeedForward, другой - расчет ошибок.

int main() {
    try {
        const uint32_t N = 1024;
        const size_t bufferSize = N * N * sizeof(float);

        // --- ШАГ 1: INSTANCE (ЭКЗЕМПЛЯР) ---
        // Это связь твоего приложения с драйвером Vulkan. 
        // Здесь мы говорим: "Привет, я хочу использовать Vulkan версии 1.1".
        vk::ApplicationInfo appInfo{"NN_GPU", 1, nullptr, 0, VK_API_VERSION_1_1};
        vk::Instance instance = vk::createInstance({{}, &appInfo});

        // --- ШАГ 2: PHYSICAL DEVICE (ЖЕЛЕЗО) ---
        // Мы ищем физическую видеокарту. В системе их может быть несколько.
        auto physicalDevices = instance.enumeratePhysicalDevices();
        vk::PhysicalDevice physDev = physicalDevices[0]; 
        
        // --- ШАГ 3: QUEUE FAMILY (ОЧЕРЕДЬ КОМАНД) ---
        // Видеокарта — это отдельный процессор. Мы не вызываем функции GPU напрямую,
        // мы отправляем "список дел" (командный буфер) в очередь (Queue).
        // Нам нужна очередь, умеющая в Compute (вычисления).
        auto queueProps = physDev.getQueueFamilyProperties();
        uint32_t computeFamily = 0; // Для упрощения берем первую, но в идеале нужно искать флаг eCompute
        
        // --- ШАГ 4: LOGICAL DEVICE (ЛОГИЧЕСКИЙ ИНТЕРФЕЙС) ---
        // Это наш "пульт управления" конкретной картой.
        float priority = 1.0f;
        vk::DeviceQueueCreateInfo qInfo({}, computeFamily, 1, &priority);
        vk::Device device = physDev.createDevice({{}, 1, &qInfo});
        vk::Queue queue = device.getQueue(computeFamily, 0);

        // --- ШАГ 5: МЕНЕДЖМЕНТ ПАМЯТИ (БУФЕРЫ) ---
        // Видеокарта не видит твой std::vector напрямую. 
        // Нужно: 1. Создать буфер в Vulkan. 2. Выделить под него видеопамять.
        auto createBuffer = [&](vk::BufferUsageFlags usage) {
            // Создаем описание буфера
            vk::Buffer buffer = device.createBuffer({{}, bufferSize, usage});
            
            // Узнаем требования буфера к памяти (сколько байт и какой тип)
            vk::MemoryRequirements memReq = device.getBufferMemoryRequirements(buffer);
            vk::PhysicalDeviceMemoryProperties memProp = physDev.getMemoryProperties();
            
            // Ищем тип памяти, который доступен и для GPU, и для CPU (HostVisible), 
            // чтобы мы могли скопировать туда веса нейронки.
            uint32_t memType = 0;
            for (uint32_t i = 0; i < memProp.memoryTypeCount; i++) {
                if ((memReq.memoryTypeBits & (1 << i)) && 
                    (memProp.memoryTypes[i].propertyFlags & (vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent))) {
                    memType = i; break;
                }
            }

            vk::DeviceMemory memory = device.allocateMemory({memReq.size, memType});
            device.bindBufferMemory(buffer, memory, 0);
            return std::make_pair(buffer, memory);
        };

        // Создаем 3 буфера: Веса, Входы, Результат
        auto [bufA, memA] = createBuffer(vk::BufferUsageFlagBits::eStorageBuffer);
        auto [bufB, memB] = createBuffer(vk::BufferUsageFlagBits::eStorageBuffer);
        auto [bufC, memC] = createBuffer(vk::BufferUsageFlagBits::eStorageBuffer);

        // --- ШАГ 6: DESCRIPTORS (СВЯЗКА ШЕЙДЕРА И ПАМЯТИ) ---
        // Шейдеру нужно знать, что в "binding 0" лежит Матрица А, а в "binding 1" - Матрица B.
        // Это настраивается через Descriptor Sets.
        std::vector<vk::DescriptorSetLayoutBinding> bindings = {
            {0, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute},
            {1, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute},
            {2, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute}
        };
        vk::DescriptorSetLayout dsLayout = device.createDescriptorSetLayout({{}, (uint32_t)bindings.size(), bindings.data()});
        
        // Создаем пул дескрипторов и сам набор
        vk::DescriptorPoolSize poolSize{vk::DescriptorType::eStorageBuffer, 3};
        vk::DescriptorPool pool = device.createDescriptorPool({{}, 1, 1, &poolSize});
        vk::DescriptorSet ds = device.allocateDescriptorSets({pool, 1, &dsLayout})[0];

        // Привязываем наши буферы к дескрипторам
        vk::DescriptorBufferInfo bInfoA{bufA, 0, VK_WHOLE_SIZE}, bInfoB{bufB, 0, VK_WHOLE_SIZE}, bInfoC{bufC, 0, VK_WHOLE_SIZE};
        device.updateDescriptorSets({
            {ds, 0, 0, 1, vk::DescriptorType::eStorageBuffer, nullptr, &bInfoA},
            {ds, 1, 0, 1, vk::DescriptorType::eStorageBuffer, nullptr, &bInfoB},
            {ds, 2, 0, 1, vk::DescriptorType::eStorageBuffer, nullptr, &bInfoC}
        }, {});

        // --- ШАГ 7: PIPELINE (КОНВЕЙЕР) ---
        // Это замороженное состояние видеокарты: какой шейдер используем, какие связи.
        // Настройка Pipeline дорогая, поэтому её делают один раз при старте.
        vk::PushConstantRange pushConstant{vk::ShaderStageFlagBits::eCompute, 0, sizeof(uint32_t)};
        vk::PipelineLayout layout = device.createPipelineLayout({{}, 1, &dsLayout, 1, &pushConstant});
        
        // Загружаем бинарный SPIR-V шейдер
        // ( readFile - самописная функция загрузки файла )
        auto shaderCode = readFile("shader.comp.spv"); 
        vk::ShaderModule shaderModule = device.createShaderModule({{}, shaderCode.size(), (uint32_t*)shaderCode.data()});
        vk::PipelineShaderStageCreateInfo stageInfo{{}, vk::ShaderStageFlagBits::eCompute, shaderModule, "main"};
        
        vk::Pipeline pipeline = device.createComputePipeline(nullptr, {{}, stageInfo, layout}).value;

        // --- ШАГ 8: COMMAND BUFFER (ЗАПИСЬ КОМАНД) ---
        // Vulkan работает так: ты записываешь все команды заранее в "плеер", 
        // а потом говоришь "Играй!". Это очень быстро.
        vk::CommandPool cmdPool = device.createCommandPool({{}, computeFamily});
        vk::CommandBuffer cmd = device.allocateCommandBuffers({cmdPool, vk::CommandBufferLevel::ePrimary, 1})[0];

        cmd.begin({vk::CommandBufferUsageFlagBits::eOneTimeSubmit});
            cmd.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline);
            cmd.bindDescriptorSets(vk::PipelineBindPoint::eCompute, layout, 0, {ds}, {});
            cmd.pushConstants(layout, vk::ShaderStageFlagBits::eCompute, 0, sizeof(uint32_t), &N);
            // Dispatch - это и есть запуск вычислений. N/16 групп потоков.
            cmd.dispatch(N / 16, N / 16, 1);
        cmd.end();

        // --- ШАГ 9: SUBMIT (ОТПРАВКА НА GPU) ---
        // Только в этот момент видеокарта начинает реально работать.
        queue.submit(vk::SubmitInfo(0, nullptr, nullptr, 1, &cmd), nullptr);
        queue.waitIdle(); // Ждем, пока GPU доделает работу

        // Теперь результат в bufC можно забрать через device.mapMemory()

        std::cout << "GPU завершил обучение/расчет!" << std::endl;

    } catch (const std::exception& e) {
        std::cerr << "Ошибка: " << e.what() << std::endl;
    }
    return 0;
}