#ifndef CORE_HPP
#define CORE_HPP

#include <vector>
#include <string>
#include <vulkan/vulkan.hpp>
#include <functional>
#include <map>
#include <mutex>

// Структура для описания слоя
struct LayerStructure_t {
    int size;                 // Количество нейронов в слое
    std::vector<int> sources; // Индексы слоев, которые являются входными для данного
    std::vector<int> sourceBranches; // Ветка для каждого источника (-1 если не разделен, 0 или 1)
    int branch;                // -1 если не разделен, 0 или 1 если слой работает с конкретной веткой
    bool isSplit;              // true если слой разделяет выход на две ветки
    
    LayerStructure_t() : size(0), branch(-1), isSplit(false) {}
    LayerStructure_t(int s, const std::vector<int>& src = {}) 
        : size(s), sources(src), branch(-1), isSplit(false) {
        sourceBranches.resize(src.size(), -1);
    }
};

// Параметры для передачи в шейдер через Push Constants
struct TrainParams {
    uint32_t mode;     // 0: FF, 1: OutError, 2: BackProp, 3: Update
    uint32_t prevSize; // Суммарный размер входов
    uint32_t nextSize; // Размер текущего слоя
    uint32_t wOff;     // Смещение весов в буфере W
    uint32_t bOff;     // Смещение смещений в буфере B
    uint32_t oOff;     // Смещение входов в буфере O
    uint32_t nextOOff; // Смещение выходов в буфере O
    float lr;          // Скорость обучения
};

// Структура для обратной связи о процессе обучения
struct TrainStatus {
    int epoch = 0;
    int totalEpochs = 0;
    int step = 0;
    int totalSteps = 0;
    double loss = 0.0;
    double totalLoss = 0.0;
    double lastEpochLoss = 0.0;
    double speed = 0.0;
    double eta = 0.0;
    float progress = 0.0f;
    long long params = 0;
};

class NeuralNetwork {
public:
    NeuralNetwork(LayerStructure_t layers[], int count, bool useVulkanParam);
    ~NeuralNetwork();

    double train(const std::map<int, std::vector<double>>& inputs, 
                 const std::vector<double>& target, double lr);
    std::vector<double> feedForward(const std::map<int, std::vector<double>>& inputs);
    
    void syncToCPU();
    void syncToGPU();
    
    // Геттеры для UI
    const std::vector<int>& getSizes() const { return sizes; }
    const std::vector<std::vector<int>>& getLayerSources() const { return layerSources; }
    const std::vector<std::vector<int>>& getSourceBranches() const { return layerSourceBranches; }
    const std::vector<int>& getBranches() const { return branches; }
    const std::vector<bool>& getSplits() const { return splits; }

private:
    bool useVulkan;
    int numLayers;

    // Метаданные архитектуры
    std::vector<int> sizes;
    std::vector<uint32_t> wOff;
    std::vector<uint32_t> bOff;
    std::vector<uint32_t> oOff;
    std::vector<std::vector<int>> layerSources;
    std::vector<std::vector<int>> layerSourceBranches; // Ветка для каждого источника
    std::vector<int> branches;  // Ветка каждого слоя (-1, 0, 1)
    std::vector<bool> splits;   // Является ли слой split

    // Данные на стороне CPU
    std::vector<float> h_weights;
    std::vector<float> h_biases;
    std::vector<float> h_outputs;
    std::vector<float> h_errors;

    // Vulkan объекты
    vk::Instance instance;
    vk::PhysicalDevice physDev;
    vk::Device device;
    vk::Queue queue;
    vk::CommandPool cmdPool;
    
    vk::DescriptorSetLayout dsLayout;
    vk::DescriptorPool descriptorPool;
    vk::DescriptorSet descriptorSet;
    vk::PipelineLayout pipeLayout;
    vk::Pipeline pipeline;
    vk::ShaderModule shaderModule;

    vk::Buffer gpuW, gpuB, gpuO, gpuE, gpuT;
    vk::DeviceMemory memW, memB, memO, memE, memT;

    void *pW = nullptr, *pB = nullptr, *pO = nullptr, *pE = nullptr, *pT = nullptr;

    // Command buffer pooling
    std::vector<vk::CommandBuffer> cmdBuffers;
    uint32_t currentCmdBuffer = 0;

    void initVulkan();
    void initVulkanResources();
    uint32_t findMemoryType(uint32_t f, vk::MemoryPropertyFlags p);
    std::vector<char> readFile(const std::string& filename);
    int calculateTotalInputSize(int layerIdx) const;
};

#endif // CORE_HPP