#include "core.hpp"
#include <cmath>
#include <cstdlib>
#include <omp.h> 
#include <vulkan/vulkan.h>
#include <iostream>
#include <vector>
#include <fstream>
#include <chrono>



NeuralNetwork::NeuralNetwork(LayerStructure_t layers[], int count, bool useVulkan) : numLayers(count) {

    if (useVulkan) {

        vk::ApplicationInfo appInfo{"Xenith", 1, nullptr, 0, VK_API_VERSION_1_1};
        instance = vk::createInstance({{}, &appInfo});
        
        auto physicalDevices = instance.enumeratePhysicalDevices();
        physDev = physicalDevices[0];
        auto props = physDev.getProperties();
        std::cout << "Используем GPU: " << props.deviceName << std::endl;

        // 3. Поиск очереди для вычислений
        auto queueProps = physDev.getQueueFamilyProperties();
        int computeFamily = -1;
        for (int i = 0; i < queueProps.size(); i++) {
            if (queueProps[i].queueFlags & vk::QueueFlagBits::eCompute) {
                computeFamily = i; break;
            }
        }
        if (computeFamily == -1) throw std::runtime_error("GPU не поддерживает Compute");

        // 4. Логическое устройство
        float priority = 1.0f;
        vk::DeviceQueueCreateInfo queueInfo({}, (uint32_t)computeFamily, 1, &priority);
        vk::DeviceCreateInfo deviceCreateInfo({}, 1, &queueInfo); 
        device = physDev.createDevice(deviceCreateInfo);

        queue = device.getQueue(computeFamily, 0);

    }

    for (int i = 0; i < count; i++) sizes.push_back(layers[i].size);
    for (int i = 0; i < count - 1; i++) {
        std::vector<std::vector<double>> layerW;
        double scale = sqrt(2.0 / sizes[i]);
        for (int j = 0; j < sizes[i+1]; j++) {
            std::vector<double> nodeW;
            for (int k = 0; k < sizes[i]; k++) 
                nodeW.push_back(((double)rand()/RAND_MAX * 2 - 1) * scale);
            layerW.push_back(nodeW);
        }
        weights.push_back(layerW);
        biases.push_back(std::vector<double>(sizes[i+1], 0.0));
    }
}

std::vector<double> NeuralNetwork::feedForward(const std::vector<double>& input) {
    outputs.clear();
    outputs.push_back(input);

    std::vector<double> curr = input;

    for (int i = 0; i < numLayers - 1; i++) {
        std::vector<double> next;
        for (int j = 0; j < sizes[i+1]; j++) {
            double sum = biases[i][j];
            for (int k = 0; k < (int)curr.size(); k++) sum += curr[k] * weights[i][j][k];
            next.push_back(1.0 / (1.0 + exp(-sum)));
        }
        curr = next;
        outputs.push_back(curr);
    }
    return curr;
}



double NeuralNetwork::train(const std::vector<double>& input, const std::vector<double>& target, double lr) {
    omp_set_num_threads(cpu_count);

    std::vector<double> pred = feedForward(input);
    std::vector<std::vector<double>> errors(numLayers);
    errors[numLayers - 1].resize(sizes[numLayers - 1]);
    
    double totalErr = 0;

    for (int i = 0; i < sizes[numLayers - 1]; i++) {
        double e = target[i] - pred[i];
        errors[numLayers - 1][i] = e * pred[i] * (1.0 - pred[i]);
        totalErr += e * e;
    }

    for (int i = numLayers - 2; i > 0; i--) {
        errors[i].resize(sizes[i]);
    
        #pragma omp parallel for
        for (int j = 0; j < sizes[i]; j++) {
            double e = 0;
            for (int k = 0; k < sizes[i + 1]; k++) {
                e += errors[i + 1][k] * weights[i][k][j];
            }
            errors[i][j] = e * outputs[i][j] * (1.0 - outputs[i][j]);
        }
    }
    for (int i = 0; i < numLayers - 1; i++) {
        #pragma omp parallel for
        for (int j = 0; j < sizes[i + 1]; j++) {
            double errorTerm = lr * errors[i + 1][j];
            
            // Вложенный цикл обновления весов
            for (int k = 0; k < sizes[i]; k++) {
                weights[i][j][k] += errorTerm * outputs[i][k];
            }
            biases[i][j] += errorTerm;
        }
    }

    return totalErr;
}

uint32_t NeuralNetwork::findMemoryType(uint32_t typeFilter, vk::MemoryPropertyFlags properties) {
    vk::PhysicalDeviceMemoryProperties memProperties = physDev.getMemoryProperties();
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {
        if ((typeFilter & (1 << i)) && (memProperties.memoryTypes[i].propertyFlags & properties) == properties) {
            return i;
        }
    }
    throw std::runtime_error("Не удалось найти подходящий тип памяти!");
}


double NeuralNetwork::trainVulkan() {
    // 1. Создание буферов
    vk::Buffer inputBuffer = device.createBuffer({{}, sizeof(float) * 2, vk::BufferUsageFlagBits::eStorageBuffer});
    vk::Buffer outputBuffer = device.createBuffer({{}, sizeof(float), vk::BufferUsageFlagBits::eStorageBuffer});
    
    // 2. Выделение и привязка памяти для ВХОДА
    vk::MemoryRequirements inReq = device.getBufferMemoryRequirements(inputBuffer);
    vk::DeviceMemory inputMemory = device.allocateMemory({
        inReq.size, 
        findMemoryType(inReq.memoryTypeBits, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent)
    });
    device.bindBufferMemory(inputBuffer, inputMemory, 0); // КРИТИЧНО: привязываем память к буферу

    // 3. Копирование данных во входной буфер
    float inputData[2] = {2.51f, 2.32f};
    void* pIn = device.mapMemory(inputMemory, 0, sizeof(float) * 2);
    memcpy(pIn, inputData, sizeof(float) * 2);
    device.unmapMemory(inputMemory);

    // 4. Выделение и привязка памяти для ВЫХОДА
    vk::MemoryRequirements outReq = device.getBufferMemoryRequirements(outputBuffer);
    vk::DeviceMemory outputMemory = device.allocateMemory({
        outReq.size, 
        findMemoryType(outReq.memoryTypeBits, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent)
    });
    device.bindBufferMemory(outputBuffer, outputMemory, 0);

    // 5. ДЕСКРИПТОРЫ (Связь C++ -> Шейдер)
    // Описываем, что у нас есть 2 слота (binding 0 и 1)
    std::vector<vk::DescriptorSetLayoutBinding> bindings = {
        {0, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute},
        {1, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute}
    };
    vk::DescriptorSetLayout dsLayout = device.createDescriptorSetLayout({{}, (uint32_t)bindings.size(), bindings.data()});
    
    // Создаем пул и выделяем сет дескрипторов
    vk::DescriptorPoolSize poolSize{vk::DescriptorType::eStorageBuffer, 2};
    vk::DescriptorPool pool = device.createDescriptorPool({{}, 1, 1, &poolSize});
    vk::DescriptorSet ds = device.allocateDescriptorSets({pool, 1, &dsLayout})[0];

    // Указываем, какие именно буферы в какие слоты вставить
    vk::DescriptorBufferInfo bInInfo{inputBuffer, 0, VK_WHOLE_SIZE};
    vk::DescriptorBufferInfo bOutInfo{outputBuffer, 0, VK_WHOLE_SIZE};
    device.updateDescriptorSets({
        {ds, 0, 0, 1, vk::DescriptorType::eStorageBuffer, nullptr, &bInInfo},
        {ds, 1, 0, 1, vk::DescriptorType::eStorageBuffer, nullptr, &bOutInfo}
    }, {});

    // 6. ПАЙПЛАЙН (Загрузка шейдера)
    auto shaderCode = readFile("shader.comp.spv"); // Твоя функция чтения файла
    vk::ShaderModule shaderModule = device.createShaderModule({{}, shaderCode.size(), (uint32_t*)shaderCode.data()});
    vk::PipelineLayout pipeLayout = device.createPipelineLayout({{}, 1, &dsLayout});
    
    vk::ComputePipelineCreateInfo pipeInfo{{}, {{}, vk::ShaderStageFlagBits::eCompute, shaderModule, "main"}, pipeLayout};
    vk::Pipeline pipeline = device.createComputePipeline(nullptr, pipeInfo).value;

    // 7. КОМАНДЫ И ЗАПУСК (Command Buffer)
    // (Предполагаем, что cmdPool и queue уже созданы в классе)
    vk::CommandBufferAllocateInfo cmdAllocInfo(cmdPool, vk::CommandBufferLevel::ePrimary, 1);
    vk::CommandBuffer cmd = device.allocateCommandBuffers(cmdAllocInfo)[0];

    cmd.begin({vk::CommandBufferUsageFlagBits::eOneTimeSubmit});
        cmd.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline);
        cmd.bindDescriptorSets(vk::PipelineBindPoint::eCompute, pipeLayout, 0, {ds}, {});
        cmd.dispatch(1, 1, 1); // Запускаем 1 поток
    cmd.end();

    queue.submit(vk::SubmitInfo(0, nullptr, nullptr, 1, &cmd), nullptr);
    queue.waitIdle();

    // 8. ЗАБИРАЕМ РЕЗУЛЬТАТ
    float result = 0;
    void* pOut = device.mapMemory(outputMemory, 0, sizeof(float));
    memcpy(&result, pOut, sizeof(float));
    device.unmapMemory(outputMemory);

    // Очистка (в реальном коде лучше делать в деструкторе)
    device.destroyPipeline(pipeline);
    device.destroyPipelineLayout(pipeLayout);
    device.destroyShaderModule(shaderModule);
    device.destroyDescriptorPool(pool);
    device.destroyDescriptorSetLayout(dsLayout);
    device.destroyBuffer(inputBuffer); device.freeMemory(inputMemory);
    device.destroyBuffer(outputBuffer); device.freeMemory(outputMemory);

    return (double)result;
}