#include "core.hpp"
#include <cmath>
#include <iostream>
#include <fstream>
#include <cstring>

int NeuralNetwork::calculateTotalInputSize(int layerIdx) const {
    int total = 0;
    for (size_t i = 0; i < layerSources[layerIdx].size(); i++) {
        int srcIdx = layerSources[layerIdx][i];
        total += sizes[srcIdx];
    }
    return total;
}

NeuralNetwork::NeuralNetwork(LayerStructure_t layers[], int count, bool useVulkanParam) {
    this->numLayers = count;
    this->useVulkan = useVulkanParam;

    uint32_t curW = 0, curB = 0, curO = 0;

    // 1. Индексация нейронов и веток
    for (int i = 0; i < count; i++) {
        sizes.push_back(layers[i].size);
        layerSources.push_back(layers[i].sources);
        layerSourceBranches.push_back(layers[i].sourceBranches);
        branches.push_back(layers[i].branch);
        splits.push_back(layers[i].isSplit);
        oOff.push_back(curO);
        curO += layers[i].size;
    }

    // 2. Инициализация весов
    for (int i = 0; i < count; i++) {
        if (layerSources[i].empty()) {
            wOff.push_back(0); 
            bOff.push_back(0);
            continue;
        }

        wOff.push_back(curW);
        bOff.push_back(curB);

        int totalInSize = calculateTotalInputSize(i);
        int wCount = totalInSize * sizes[i];
        float scale = sqrt(2.0f / totalInSize);

        for (int j = 0; j < wCount; j++) {
            h_weights.push_back(((float)rand() / RAND_MAX * 2.0f - 1.0f) * scale);
        }
        for (int j = 0; j < sizes[i]; j++) h_biases.push_back(0.0f);

        curW += wCount;
        curB += sizes[i];
    }

    h_outputs.resize(curO, 0.0f);
    h_errors.resize(curO, 0.0f);

    if (this->useVulkan) {
        initVulkan();
        initVulkanResources();
        syncToGPU();
    }
}

double NeuralNetwork::train(const std::map<int, std::vector<double>>& inputs, 
                            const std::vector<double>& target, double lr) {
    if (!useVulkan) return 0.0;

    // Подготовка входных данных
    for (auto const& [layerIdx, data] : inputs) {
        if (layerIdx >= 0 && layerIdx < numLayers) {
            float* ptr = (float*)pO + oOff[layerIdx];
            size_t copySize = std::min(data.size(), (size_t)sizes[layerIdx]);
            for (size_t i = 0; i < copySize; i++) ptr[i] = (float)data[i];
        }
    }
    
    // Подготовка target
    float* fTar = (float*)pT;
    for (size_t i = 0; i < target.size(); i++) fTar[i] = (float)target[i];

    // Получаем command buffer из пула
    if (cmdBuffers.empty()) {
        vk::CommandBufferAllocateInfo ai(cmdPool, vk::CommandBufferLevel::ePrimary, 4);
        cmdBuffers = device.allocateCommandBuffers(ai);
    }
    vk::CommandBuffer cmd = cmdBuffers[currentCmdBuffer];
    currentCmdBuffer = (currentCmdBuffer + 1) % cmdBuffers.size();

    cmd.begin({vk::CommandBufferUsageFlagBits::eOneTimeSubmit});
    cmd.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline);
    cmd.bindDescriptorSets(vk::PipelineBindPoint::eCompute, pipeLayout, 0, {descriptorSet}, {});

    vk::MemoryBarrier barrier(vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eShaderRead, 
                              vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eShaderRead);

    // 1. FORWARD PASS
    for (int i = 0; i < numLayers; i++) {
        if (layerSources[i].empty()) continue;
        
        int totalIn = calculateTotalInputSize(i);
        int firstSrc = layerSources[i][0];

        TrainParams p = {0, (uint32_t)totalIn, (uint32_t)sizes[i], wOff[i], bOff[i], 
                         oOff[firstSrc], oOff[i], (float)lr};
        cmd.pushConstants(pipeLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof(TrainParams), &p);
        cmd.dispatch((sizes[i] + 255) / 256, 1, 1);
        cmd.pipelineBarrier(vk::PipelineStageFlagBits::eComputeShader, 
                           vk::PipelineStageFlagBits::eComputeShader, {}, {barrier}, {}, {});
    }

    // 2. OUTPUT ERROR
    {
        TrainParams p = {1, 0, (uint32_t)sizes.back(), 0, 0, 0, oOff.back(), (float)lr};
        cmd.pushConstants(pipeLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof(TrainParams), &p);
        cmd.dispatch((sizes.back() + 255) / 256, 1, 1);
        cmd.pipelineBarrier(vk::PipelineStageFlagBits::eComputeShader, 
                           vk::PipelineStageFlagBits::eComputeShader, {}, {barrier}, {}, {});
    }

    // 3. BACKWARD PASS
    for (int i = numLayers - 1; i >= 0; i--) {
        if (layerSources[i].empty()) continue;
        
        int totalIn = calculateTotalInputSize(i);
        int firstSrc = layerSources[i][0];

        TrainParams p = {2, (uint32_t)totalIn, (uint32_t)sizes[i], wOff[i], bOff[i], 
                         oOff[firstSrc], oOff[i], (float)lr};
        cmd.pushConstants(pipeLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof(TrainParams), &p);
        cmd.dispatch((totalIn + 255) / 256, 1, 1);
        cmd.pipelineBarrier(vk::PipelineStageFlagBits::eComputeShader, 
                           vk::PipelineStageFlagBits::eComputeShader, {}, {barrier}, {}, {});
    }

    // 4. UPDATE WEIGHTS
    for (int i = 0; i < numLayers; i++) {
        if (layerSources[i].empty()) continue;
        
        int totalIn = calculateTotalInputSize(i);
        int firstSrc = layerSources[i][0];

        TrainParams p = {3, (uint32_t)totalIn, (uint32_t)sizes[i], wOff[i], bOff[i], 
                         oOff[firstSrc], oOff[i], (float)lr};
        cmd.pushConstants(pipeLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof(TrainParams), &p);
        cmd.dispatch((sizes[i] + 255) / 256, 1, 1);
    }

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

    // Расчет MSE
    float* out = (float*)pO + oOff.back();
    double mse = 0;
    for (int i = 0; i < sizes.back(); i++) { 
        double d = (double)target[i] - (double)out[i]; 
        mse += d * d; 
    }
    return mse / sizes.back();
}

std::vector<double> NeuralNetwork::feedForward(const std::map<int, std::vector<double>>& inputs) {
    if (!useVulkan) return {};
    
    for (auto const& [layerIdx, data] : inputs) {
        if (layerIdx >= 0 && layerIdx < numLayers) {
            float* ptr = (float*)pO + oOff[layerIdx];
            size_t copySize = std::min(data.size(), (size_t)sizes[layerIdx]);
            memcpy(ptr, data.data(), copySize * sizeof(float));
        }
    }

    if (cmdBuffers.empty()) {
        vk::CommandBufferAllocateInfo ai(cmdPool, vk::CommandBufferLevel::ePrimary, 4);
        cmdBuffers = device.allocateCommandBuffers(ai);
    }
    vk::CommandBuffer cmd = cmdBuffers[currentCmdBuffer];
    currentCmdBuffer = (currentCmdBuffer + 1) % cmdBuffers.size();

    cmd.begin({vk::CommandBufferUsageFlagBits::eOneTimeSubmit});
    cmd.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline);
    cmd.bindDescriptorSets(vk::PipelineBindPoint::eCompute, pipeLayout, 0, {descriptorSet}, {});

    vk::MemoryBarrier barrier(vk::AccessFlagBits::eShaderWrite, vk::AccessFlagBits::eShaderRead);

    for (int i = 0; i < numLayers; i++) {
        if (layerSources[i].empty()) continue;
        
        int totalIn = calculateTotalInputSize(i);
        int firstSrc = layerSources[i][0];

        TrainParams p = {0, (uint32_t)totalIn, (uint32_t)sizes[i], wOff[i], bOff[i], 
                         oOff[firstSrc], oOff[i], 0.0f};
        cmd.pushConstants(pipeLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof(TrainParams), &p);
        cmd.dispatch((sizes[i] + 255) / 256, 1, 1);
        cmd.pipelineBarrier(vk::PipelineStageFlagBits::eComputeShader, 
                           vk::PipelineStageFlagBits::eComputeShader, {}, {barrier}, {}, {});
    }

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

    std::vector<double> result(sizes.back());
    float* out = (float*)pO + oOff.back();
    for (int i = 0; i < sizes.back(); i++) result[i] = (double)out[i];
    return result;
}

void NeuralNetwork::initVulkan() {
    vk::ApplicationInfo app{"Xenith", 1, nullptr, 0, VK_API_VERSION_1_1};
    instance = vk::createInstance({{}, &app});
    physDev = instance.enumeratePhysicalDevices()[0];
    auto props = physDev.getQueueFamilyProperties();
    int qIdx = -1;
    for (int i = 0; i < props.size(); i++) 
        if (props[i].queueFlags & vk::QueueFlagBits::eCompute) { 
            qIdx = i; 
            break; 
        }
    float priority = 1.0f;
    device = physDev.createDevice({{}, 1, new vk::DeviceQueueCreateInfo({}, (uint32_t)qIdx, 1, &priority)});
    queue = device.getQueue(qIdx, 0);
    cmdPool = device.createCommandPool({{}, (uint32_t)qIdx});
}

void NeuralNetwork::initVulkanResources() {
    auto createBuf = [&](size_t sz, vk::Buffer& b, vk::DeviceMemory& m, void** ptr) {
        if (sz == 0) sz = 1;
        b = device.createBuffer({{}, sz * sizeof(float), vk::BufferUsageFlagBits::eStorageBuffer});
        auto req = device.getBufferMemoryRequirements(b);
        m = device.allocateMemory({req.size, findMemoryType(req.memoryTypeBits, 
                                   vk::MemoryPropertyFlagBits::eHostVisible | 
                                   vk::MemoryPropertyFlagBits::eHostCoherent)});
        device.bindBufferMemory(b, m, 0);
        *ptr = device.mapMemory(m, 0, sz * sizeof(float));
    };
    
    createBuf(h_weights.size(), gpuW, memW, &pW);
    createBuf(h_biases.size(), gpuB, memB, &pB);
    createBuf(h_outputs.size(), gpuO, memO, &pO);
    createBuf(h_errors.size(), gpuE, memE, &pE);
    createBuf(sizes.back(), gpuT, memT, &pT);

    vk::DescriptorSetLayoutBinding binds[] = {
        {0, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute},
        {1, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute},
        {2, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute},
        {3, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute},
        {4, vk::DescriptorType::eStorageBuffer, 1, vk::ShaderStageFlagBits::eCompute}
    };
    dsLayout = device.createDescriptorSetLayout({{}, 5, binds});
    vk::DescriptorPoolSize ps(vk::DescriptorType::eStorageBuffer, 5);
    descriptorPool = device.createDescriptorPool({{}, 1, 1, &ps});
    descriptorSet = device.allocateDescriptorSets({descriptorPool, 1, &dsLayout})[0];

    vk::DescriptorBufferInfo db[] = {
        {gpuW,0,VK_WHOLE_SIZE}, {gpuB,0,VK_WHOLE_SIZE}, {gpuO,0,VK_WHOLE_SIZE}, 
        {gpuE,0,VK_WHOLE_SIZE}, {gpuT,0,VK_WHOLE_SIZE}
    };
    vk::WriteDescriptorSet wds[] = {
        {descriptorSet,0,0,1,vk::DescriptorType::eStorageBuffer,nullptr,db+0},
        {descriptorSet,1,0,1,vk::DescriptorType::eStorageBuffer,nullptr,db+1},
        {descriptorSet,2,0,1,vk::DescriptorType::eStorageBuffer,nullptr,db+2},
        {descriptorSet,3,0,1,vk::DescriptorType::eStorageBuffer,nullptr,db+3},
        {descriptorSet,4,0,1,vk::DescriptorType::eStorageBuffer,nullptr,db+4}
    };
    device.updateDescriptorSets(5, wds, 0, nullptr);

    auto code = readFile("Xenith/shader.comp.spv");
    if (code.empty()) {
        std::cerr << "ERROR: Failed to load shader.comp.spv!\n";
        exit(1);
    }
    shaderModule = device.createShaderModule({{}, code.size(), (uint32_t*)code.data()});
    vk::PushConstantRange pr(vk::ShaderStageFlagBits::eCompute, 0, sizeof(TrainParams));
    pipeLayout = device.createPipelineLayout({{}, 1, &dsLayout, 1, &pr});
    pipeline = device.createComputePipeline(nullptr, {{}, {{}, vk::ShaderStageFlagBits::eCompute, 
                        shaderModule, "main"}, pipeLayout}).value;
}

uint32_t NeuralNetwork::findMemoryType(uint32_t f, vk::MemoryPropertyFlags p) {
    auto props = physDev.getMemoryProperties();
    for (uint32_t i = 0; i < props.memoryTypeCount; i++)
        if ((f & (1 << i)) && (props.memoryTypes[i].propertyFlags & p) == p) return i;
    return 0;
}

std::vector<char> NeuralNetwork::readFile(const std::string& n) {
    std::ifstream f(n, std::ios::ate | std::ios::binary);
    if (!f.is_open()) {
        std::cerr << "ERROR: Cannot open file: " << n << "\n";
        return {};
    }
    size_t s = (size_t)f.tellg(); 
    std::vector<char> b(s); 
    f.seekg(0); 
    f.read(b.data(), s); 
    return b;
}

void NeuralNetwork::syncToCPU() {
    memcpy(h_weights.data(), pW, h_weights.size() * 4);
    memcpy(h_biases.data(), pB, h_biases.size() * 4);
}

void NeuralNetwork::syncToGPU() {
    memcpy(pW, h_weights.data(), h_weights.size() * 4);
    memcpy(pB, h_biases.data(), h_biases.size() * 4);
}

NeuralNetwork::~NeuralNetwork() {
    if (useVulkan) {
        device.waitIdle();
        
        if (!cmdBuffers.empty()) {
            device.freeCommandBuffers(cmdPool, cmdBuffers);
        }
        
        device.destroyPipeline(pipeline); 
        device.destroyPipelineLayout(pipeLayout);
        device.destroyShaderModule(shaderModule);
        
        device.destroyBuffer(gpuW); device.freeMemory(memW);
        device.destroyBuffer(gpuB); device.freeMemory(memB);
        device.destroyBuffer(gpuO); device.freeMemory(memO);
        device.destroyBuffer(gpuE); device.freeMemory(memE);
        device.destroyBuffer(gpuT); device.freeMemory(memT);
        
        device.destroyDescriptorPool(descriptorPool);
        device.destroyDescriptorSetLayout(dsLayout);
        device.destroyCommandPool(cmdPool);
        device.destroy(); 
        instance.destroy();
    }
}