#include "VisionGpuDispatch.hpp"

#include "MayaFlux/IO/ImageExport.hpp"
#include "MayaFlux/Journal/Archivist.hpp"
#include "MayaFlux/Portal/Graphics/ShaderFoundry.hpp"

namespace MayaFlux::Yantra {

using namespace Portal::Graphics;
using namespace Kinesis::Vision;

// ============================================================================
// Internal push constant layouts
// ============================================================================

namespace {
    constexpr uint32_t CC_BACKGROUND_HOST = 0xFFFFFFFFU;
    constexpr uint32_t CC_UNCLAIMED_HOST = 0U;

    /** Standard 2D workgroup used by all pixel-to-pixel vision shaders */
    constexpr std::array<uint32_t, 3> k_wg2d { 8, 8, 1 };
    /** Maximum number of connected components that can be labeled in a single pass */
    constexpr uint32_t k_max_components = 4096;
    /** Maximum number of points that can be stored in a single contour */
    constexpr uint32_t k_max_points_per_contour = 4096;
    /** Maximum number of contours that can be stored in a single pass */
    constexpr uint32_t k_max_holes_per_label = 4;
    /** Maximum number of trace slots that can be stored in a single pass (for contour tracing) */
    constexpr uint32_t k_max_trace_slots = k_max_components * (1U + k_max_holes_per_label);

    struct ThresholdPC {
        float value;
    };
    struct ThresholdAdaptivePC {
        uint32_t block_size;
        float offset;
    };
    struct OtsuHistPC {
        uint32_t width;
        uint32_t height;
    };
    struct NormalizePC {
        float scale;
        float offset;
    };
    struct MorphPC {
        uint32_t radius;
    };
    struct HarrisPC {
        float k;
        uint32_t pass;
        uint32_t width;
        uint32_t height;
    };
    struct ExtractPC {
        float threshold;
        uint32_t nms_radius;
    };
    struct CannyPC {
        float sigma;
        float lo;
        float hi;
    };
    struct RgbaToGrayPC {
        float wr;
        float wg;
        float wb;
        float wa;
    };
    struct GaussianPC {
        uint32_t radius;
        uint32_t width;
        uint32_t height;
    };
    struct CompletedOp {
        std::shared_ptr<Core::VKImage> output;
        std::shared_ptr<Core::VKImage> input;
    };
    struct ClassifyPC {
        float threshold;
        float value;
    };
    struct HysteresisPC {
        uint32_t width;
        uint32_t height;
    };
    struct FinalizePC {
        float threshold;
    };
    struct ExtractPeaksPC {
        float threshold;
        uint32_t nms_radius;
        uint32_t width;
        uint32_t height;
        uint32_t max_keypoints;
    };
    struct CCBlockInitPC {
        uint32_t width;
        uint32_t height;
        uint32_t block_width;
        uint32_t block_height;
    };
    struct CCMergePC {
        uint32_t width;
        uint32_t height;
        uint32_t block_width;
        uint32_t block_height;
    };
    struct CCCompressPC {
        uint32_t block_width;
        uint32_t block_height;
    };
    struct CCFinalLabelPC {
        uint32_t width;
        uint32_t height;
        uint32_t block_width;
        uint32_t block_height;
        uint32_t max_components;
        uint32_t export_labels;
    };
    struct CCResetPC {
        uint32_t lut_size;
        uint32_t max_components;
    };
    struct ContourSegmentsPC {
        uint32_t width;
        uint32_t height;
        uint32_t max_segments;
    };
    struct ContourLinkPC {
        uint32_t width;
        uint32_t height;
        uint32_t phase;
    };
    struct ContourClearPC {
        uint32_t width;
        uint32_t height;
    };
    struct ContourRenderPC {
        uint32_t width;
        uint32_t height;
        uint32_t max_components;
        uint32_t max_points_per_contour;
        uint32_t max_contours;
    };
    struct ContourMarchPC {
        uint32_t width;
        uint32_t height;
        uint32_t max_components;
        uint32_t max_points_per_contour;
        uint32_t max_holes_per_label;
        uint32_t phase;
        float min_area;
        uint32_t compacted_count;
    };
    struct ContourCompactPC {
        uint32_t max_components;
        uint32_t max_holes_per_label;
    };

    /**
     * @brief 2D Gaussian kernel for convolution, cached by (radius, sigma
     *        bit pattern).
     *
     * Sigma is a tuning parameter that rarely changes frame to frame;
     * recomputing exp() over (2*radius+1)^2 taps and reallocating the
     * kernel every call is pure repeated work for an identical result.
     *
     * @param radius Radius of the kernel in pixels. Kernel size is (2*radius + 1)^2.
     * @param sigma  Standard deviation of the Gaussian.
     * @return       Normalized kernel weights as a flat vector in row-major order.
     */
    const std::vector<float>& gaussian_kernel_2d(uint32_t radius, float sigma)
    {
        static std::unordered_map<uint64_t, std::vector<float>> cache;
        const uint64_t key = (static_cast<uint64_t>(std::bit_cast<uint32_t>(sigma)) << 32)
            | radius;
        auto it = cache.find(key);
        if (it != cache.end())
            return it->second;

        const uint32_t diam = 2 * radius + 1;
        std::vector<float> k(static_cast<size_t>(diam) * diam);
        float sum = 0.0F;
        for (uint32_t y = 0; y < diam; ++y) {
            for (uint32_t x = 0; x < diam; ++x) {
                const float fx = static_cast<float>(x) - static_cast<float>(radius);
                const float fy = static_cast<float>(y) - static_cast<float>(radius);
                const float v = std::exp(-(fx * fx + fy * fy) / (2.0F * sigma * sigma));
                k[y * diam + x] = v;
                sum += v;
            }
        }
        for (auto& v : k)
            v /= sum;

        return cache.emplace(key, std::move(k)).first->second;
    }

    /**
     * @brief 1D Gaussian kernel for separable convolution, cached by
     *        (radius, sigma bit pattern).
     *
     * Sigma is a tuning parameter that rarely changes frame to frame;
     * recomputing exp() per tap and reallocating the kernel every call
     * is pure repeated work for an identical result. Mirrors
     * VisionExecutor::gaussian_kernel's caching rationale for the CPU path.
     *
     * @param radius Radius of the kernel in pixels. Kernel size is (2*radius + 1).
     * @param sigma  Standard deviation of the Gaussian.
     * @return       Normalized kernel weights.
     */
    const std::vector<float>& gaussian_kernel_1d(uint32_t radius, float sigma)
    {
        static std::unordered_map<uint64_t, std::vector<float>> cache;
        const uint64_t key = (static_cast<uint64_t>(std::bit_cast<uint32_t>(sigma)) << 32)
            | radius;
        auto it = cache.find(key);
        if (it != cache.end())
            return it->second;

        const uint32_t size = 2 * radius + 1;
        std::vector<float> k(size);
        float sum = 0.0F;
        for (uint32_t i = 0; i < size; ++i) {
            const float x = static_cast<float>(i) - static_cast<float>(radius);
            k[i] = std::exp(-(x * x) / (2.0F * sigma * sigma));
            sum += k[i];
        }
        for (auto& v : k)
            v /= sum;

        return cache.emplace(key, std::move(k)).first->second;
    }

} // namespace

VisionGpuContexts::VisionGpuContexts()
    : pixel {
        GpuComputeConfig {},
        Portal::Graphics::ImageFormat::RGBA32F,
        TextureExecutionContext::OutputMode::IMAGE,
        1,
        std::vector<GpuBufferBinding> {
            { .set = 0, .binding = 2, .direction = GpuBufferBinding::Direction::INPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
        },
        GpuBufferBinding::ElementType::IMAGE_STORAGE,
    }
    , structured {
        GpuComputeConfig {},
        Portal::Graphics::ImageFormat::RGBA32F,
        TextureExecutionContext::OutputMode::SCALAR,
        0,
        std::vector<GpuBufferBinding> {
            { .set = 0, .binding = 1, .direction = GpuBufferBinding::Direction::OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 2, .direction = GpuBufferBinding::Direction::OUTPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
            { .set = 0, .binding = 3, .direction = GpuBufferBinding::Direction::OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 4, .direction = GpuBufferBinding::Direction::OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 0, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
            { .set = 1, .binding = 1, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
        },
        GpuBufferBinding::ElementType::IMAGE_STORAGE,
        0,
    }
    , labels {
        GpuComputeConfig {},
        Portal::Graphics::ImageFormat::RGBA32F,
        TextureExecutionContext::OutputMode::IMAGE,
        1,
        std::vector<GpuBufferBinding> {
            { .set = 0, .binding = 2, .direction = GpuBufferBinding::Direction::OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 4, .direction = GpuBufferBinding::Direction::OUTPUT, .element_type = GpuBufferBinding::ElementType::IMAGE_STORAGE },
        },
        GpuBufferBinding::ElementType::IMAGE_STORAGE,
    }
    , cc_pipeline {
        GpuComputeConfig {},
        Portal::Graphics::ImageFormat::RGBA32F,
        TextureExecutionContext::OutputMode::IMAGE,
        1,
        std::vector<GpuBufferBinding> {
            { .set = 0, .binding = 2, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 3, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 4, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 5, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 6, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 7, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 8, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 9, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 0, .binding = 10, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 0, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 1, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 2, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 3, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 4, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 5, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 6, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 7, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 8, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
            { .set = 1, .binding = 9, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 10, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 },
            { .set = 1, .binding = 11, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
            { .set = 2, .binding = 0, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
            { .set = 2, .binding = 1, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::FLOAT32 },
        },
        GpuBufferBinding::ElementType::IMAGE_STORAGE,
        0,
    }
{
    structured.set_output_size(1, sizeof(uint32_t));
    structured.set_output_size(2, static_cast<size_t>(4096) * 4 * sizeof(float));
    structured.set_output_size(3, static_cast<size_t>(256) * sizeof(uint32_t));
    structured.set_output_size(4, sizeof(uint32_t));
}

// ============================================================================
// vision_gpu_config
// ============================================================================

GpuComputeConfig VisionGpuExecutor::config(VisionOp op, const VisionParams& /*params*/)
{
    switch (op) {
    case VisionOp::Threshold: {
        const auto spec = ShaderSpec::Assemble {}
                              .storage_image("out", BindingDirection::Output)
                              .storage_image("src", BindingDirection::Input)
                              .pc("threshold")
                              .op(KernelOp::CompareGE)
                              .workgroup(k_wg2d[0], k_wg2d[1])
                              .build();
        return config_from_spec(spec);
    }
    case VisionOp::RgbaToGray: {
        const auto spec = ShaderSpec::Assemble {}
                              .storage_image("out", BindingDirection::Output)
                              .storage_image("src", BindingDirection::Input)
                              .pc("wr")
                              .pc("wg")
                              .pc("wb")
                              .pc("wa")
                              .op(KernelOp::ChannelDot)
                              .workgroup(k_wg2d[0], k_wg2d[1])
                              .build();
        return config_from_spec(spec);
    }
    case VisionOp::GrayToRgba: {
        const auto spec = ShaderSpec::Assemble {}
                              .storage_image("out", BindingDirection::Output)
                              .storage_image("src", BindingDirection::Input)
                              .op(KernelOp::ChannelReplicate)
                              .workgroup(k_wg2d[0], k_wg2d[1])
                              .build();
        return config_from_spec(spec);
    }
    case VisionOp::GaussianBlur: {
        const auto spec = ShaderSpec::Assemble {}
                              .tmpl(KernelTemplate::Convolve2D)
                              .storage_image("out", BindingDirection::Output)
                              .storage_image("src", BindingDirection::Input)
                              .ssbo("kernel", BindingDirection::Input, Kakshya::GpuDataFormat::FLOAT32)
                              .pc("radius", Kakshya::GpuDataFormat::UINT32)
                              .pc("width", Kakshya::GpuDataFormat::UINT32)
                              .pc("height", Kakshya::GpuDataFormat::UINT32)
                              .workgroup(k_wg2d[0], k_wg2d[1])
                              .build();
        return config_from_spec(spec);
    }
    case VisionOp::NormalizeRange: {
        const auto spec = ShaderSpec::Assemble {}
                              .storage_image("out", BindingDirection::Output)
                              .storage_image("src", BindingDirection::Input)
                              .pc("scale")
                              .pc("offset")
                              .op(KernelOp::ScaleOffset)
                              .workgroup(k_wg2d[0], k_wg2d[1])
                              .build();
        return config_from_spec(spec);
    }
    case VisionOp::NormalizeInplace: {
        const auto spec = ShaderSpec::Assemble {}
                              .storage_image("out", BindingDirection::Output)
                              .storage_image("src", BindingDirection::Input)
                              .op(KernelOp::Scale)
                              .workgroup(k_wg2d[0], k_wg2d[1])
                              .build();
        return config_from_spec(spec);
    }
    case VisionOp::Canny: {
        const auto spec = ShaderSpec::Assemble {}
                              .storage_image("out", BindingDirection::Output)
                              .storage_image("src", BindingDirection::Input)
                              .pc("threshold")
                              .pc("value")
                              .op(KernelOp::CompareGEPreserve)
                              .workgroup(k_wg2d[0], k_wg2d[1])
                              .build();
        return config_from_spec(spec);
    }
    case VisionOp::RgbaToHsv:
        return { .shader_path = "rgba_to_hsv.comp.spv", .workgroup_size = k_wg2d };
    case VisionOp::Downsample2x:
        return { .shader_path = "downsample_2x.comp.spv", .workgroup_size = k_wg2d };
    case VisionOp::FilterSeparable:
        return { .shader_path = "filter_separable.comp.spv", .workgroup_size = k_wg2d };
    case VisionOp::Sobel:
        return { .shader_path = "sobel.comp.spv", .workgroup_size = k_wg2d };
    case VisionOp::Scharr:
        return { .shader_path = "scharr.comp.spv", .workgroup_size = k_wg2d };
    case VisionOp::Erode:
        return { .shader_path = "erode.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(MorphPC) };
    case VisionOp::Dilate:
        return { .shader_path = "dilate.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(MorphPC) };
    case VisionOp::Open:
        return { .shader_path = "open.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(MorphPC) };
    case VisionOp::Close:
        return { .shader_path = "close.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(MorphPC) };
    case VisionOp::MorphGradient:
        return { .shader_path = "morph_gradient.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(MorphPC) };
    case VisionOp::HarrisResponse:
        return { .shader_path = "harris_response.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(HarrisPC) };
    case VisionOp::ExtractPeaks:
        return { .shader_path = "extract_peaks.comp.spv", .workgroup_size = { 256, 1, 1 }, .push_constant_size = sizeof(ExtractPC) };
    case VisionOp::ConnectedComponents:
        return { .shader_path = "cc_colorize.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(uint32_t) * 2 };
    case VisionOp::FindContours:
        return { .shader_path = "contour_segments.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(ContourSegmentsPC) };
    case VisionOp::ThresholdAdaptive:
        return { .shader_path = "threshold_adaptive.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(ThresholdAdaptivePC) };
    case VisionOp::ThresholdOtsu:
        return { .shader_path = "threshold_otsu.comp.spv", .workgroup_size = { 256, 1, 1 } };
    default:
        return GpuComputeConfig { .shader_id = Portal::Graphics::INVALID_SHADER };
    }
}

// ============================================================================
// run_gpu
// ============================================================================

VisionResult VisionGpuExecutor::run(
    VisionGpuContexts& contexts,
    const VisionSequence& sequence,
    const std::shared_ptr<Core::VKImage>& image,
    uint32_t w, uint32_t h)
{
    auto& pixel_ctx = contexts.pixel;
    auto& structured_ctx = contexts.structured;
    auto& label_ctx = contexts.labels;
    auto& cc_pipeline = contexts.cc_pipeline;

    VisionResult result;
    result.w = w;
    result.h = h;

    auto& foundry = Portal::Graphics::get_shader_foundry();
    std::shared_ptr<Core::VKImage> current = image;

    Portal::Graphics::ShaderID last_shader_id = Portal::Graphics::INVALID_SHADER;
    std::string last_shader_path;
    std::shared_ptr<Core::VKImage> last_staged;
    uint32_t last_w = 0, last_h = 0;
    std::unordered_map<size_t, CompletedOp> completed_ops;

    pixel_ctx.prepare_output_image(w, h);
    last_w = w;
    last_h = h;

    for (const auto& step : sequence.steps) {
        const auto cfg = config(step.op, step.params);

        if (cfg.shader_id == Portal::Graphics::INVALID_SHADER && cfg.shader_path.empty()) {
            MF_ERROR(Journal::Component::Yantra, Journal::Context::ComputeMatrix,
                "run_gpu: no GPU implementation for VisionOp {}",
                static_cast<int>(step.op));
            return VisionResult {};
        }

        if (cfg.shader_id != last_shader_id || cfg.shader_path != last_shader_path) {
            pixel_ctx.swap_shader(cfg);
            last_shader_id = cfg.shader_id;
            last_shader_path = cfg.shader_path;
        }
        if (current != last_staged) {
            pixel_ctx.stage_image(current);
            last_staged = current;
        }
        if (w != last_w || h != last_h) {
            pixel_ctx.prepare_output_image(w, h);
            last_w = w;
            last_h = h;
        }
        pixel_ctx.set_output_dimensions(w, h);

        switch (step.op) {
        case VisionOp::Downsample2x: {
            const uint32_t new_w = std::max(1U, w / 2);
            const uint32_t new_h = std::max(1U, h / 2);

            pixel_ctx.prepare_output_image(new_w, new_h);
            pixel_ctx.set_output_dimensions(new_w, new_h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            pixel_ctx.clear_output_dimensions();

            auto downsampled = pixel_ctx.get_output_image(0);
            completed_ops[Kinesis::Vision::hash_vision_step(step.op, step.params)] = { .output = downsampled, .input = current };
            current = downsampled;
            last_staged = current;

            w = new_w;
            h = new_h;
            result.w = w;
            result.h = h;
            last_w = w;
            last_h = h;

            continue;
        }
        case VisionOp::Threshold:
            pixel_ctx.set_push_constants(ThresholdPC {
                .value = std::get<ThresholdParams>(step.params).value });
            break;
        case VisionOp::NormalizeRange: {
            const auto& p = std::get<NormalizeRangeParams>(step.params);
            const float scale = (p.hi > p.lo) ? 1.0F / (p.hi - p.lo) : 1.0F;
            const float off = (p.hi > p.lo) ? -p.lo / (p.hi - p.lo) : 0.0F;
            pixel_ctx.set_push_constants(NormalizePC { .scale = scale, .offset = off });
            break;
        }
        case VisionOp::RgbaToGray:
            pixel_ctx.set_push_constants(RgbaToGrayPC {
                .wr = 0.299F, .wg = 0.587F, .wb = 0.114F, .wa = 0.0F });
            break;
        case VisionOp::GaussianBlur: {
            const auto& p = std::get<GaussianBlurParams>(step.params);
            const auto radius = static_cast<uint32_t>(std::ceil(p.sigma * 3.0F));
            const auto& weights = gaussian_kernel_1d(radius, p.sigma);
            pixel_ctx.set_binding_data(2, std::span<const float>(weights));
            pixel_ctx.set_push_constants(GaussianPC { .radius = radius, .width = w, .height = h });
            break;
        }
        case VisionOp::Erode:
        case VisionOp::Dilate:
        case VisionOp::MorphGradient:
            pixel_ctx.set_push_constants(MorphPC {
                .radius = std::get<MorphParams>(step.params).radius });
            break;
        case VisionOp::ThresholdAdaptive: {
            const auto& p = std::get<ThresholdAdaptiveParams>(step.params);
            pixel_ctx.set_push_constants(ThresholdAdaptivePC { .block_size = p.block_size, .offset = p.offset });
            break;
        }
        case VisionOp::ThresholdOtsu: {
            structured_ctx.swap_shader({
                .shader_path = "otsu_histogram.comp.spv",
                .workgroup_size = k_wg2d,
                .push_constant_size = sizeof(OtsuHistPC),
            });
            std::vector<uint32_t> zeros(256, 0);
            structured_ctx.set_binding_data(3, std::span<const uint32_t>(zeros));
            structured_ctx.stage_image(current);
            structured_ctx.set_push_constants(OtsuHistPC { .width = w, .height = h });
            structured_ctx.set_output_dimensions(w, h);
            {
                const auto f = structured_ctx.dispatch_async({});
                structured_ctx.clear_output_dimensions();
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }

            const auto hist_check = structured_ctx.collect_result();
            std::vector<uint32_t> hist_readback(256, 0);
            if (auto it = hist_check.aux.find(3); it != hist_check.aux.end())
                std::memcpy(hist_readback.data(), it->second.data(), 256 * sizeof(uint32_t));

            structured_ctx.swap_shader({
                .shader_path = "otsu_select.comp.spv",
                .workgroup_size = { 256, 1, 1 },
            });
            structured_ctx.set_binding_data(3, std::span<const uint32_t>(hist_readback));
            structured_ctx.set_output_dimensions(256, 1);

            {
                const auto f = structured_ctx.dispatch_async({});
                structured_ctx.clear_output_dimensions();
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }

            const auto sel_result = structured_ctx.collect_result();
            uint32_t best_bin = 0;
            if (auto it = sel_result.aux.find(4); it != sel_result.aux.end())
                std::memcpy(&best_bin, it->second.data(), sizeof(uint32_t));
            const float t_norm = static_cast<float>(best_bin) / 255.0F;

            const auto apply_cfg = config_from_spec(
                ShaderSpec::Assemble {}
                    .storage_image("out", BindingDirection::Output)
                    .storage_image("src", BindingDirection::Input)
                    .pc("threshold")
                    .op(KernelOp::CompareGE)
                    .workgroup(k_wg2d[0], k_wg2d[1])
                    .build());
            pixel_ctx.swap_shader(apply_cfg);
            pixel_ctx.stage_image(current);
            pixel_ctx.set_push_constants(ThresholdPC { .value = t_norm });
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto thresholded = pixel_ctx.get_output_image(0);

            completed_ops[Kinesis::Vision::hash_vision_step(step.op, step.params)] = { .output = thresholded, .input = current };
            result.debug_labels = thresholded;
            current = thresholded;
            last_staged = current;
            result.structured = std::monostate {};
            continue;
        }
        case VisionOp::Open:
        case VisionOp::Close: {
            const auto radius = std::get<MorphParams>(step.params).radius;
            const bool is_open = (step.op == VisionOp::Open);

            const GpuComputeConfig first_cfg {
                .shader_path = is_open ? "erode.comp.spv" : "dilate.comp.spv",
                .workgroup_size = k_wg2d,
                .push_constant_size = sizeof(MorphPC),
            };
            pixel_ctx.swap_shader(first_cfg);
            pixel_ctx.stage_image(current);
            pixel_ctx.set_push_constants(MorphPC { .radius = radius });
            pixel_ctx.prepare_output_image(w, h);
            pixel_ctx.set_output_dimensions(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto intermediate = pixel_ctx.get_output_image(0);

            const GpuComputeConfig second_cfg {
                .shader_path = is_open ? "dilate.comp.spv" : "erode.comp.spv",
                .workgroup_size = k_wg2d,
                .push_constant_size = sizeof(MorphPC),
            };
            pixel_ctx.swap_shader(second_cfg);
            pixel_ctx.stage_image(intermediate);
            pixel_ctx.set_push_constants(MorphPC { .radius = radius });
            pixel_ctx.prepare_output_image(w, h);
            pixel_ctx.set_output_dimensions(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }

            auto opened_closed = pixel_ctx.get_output_image(0);
            completed_ops[Kinesis::Vision::hash_vision_step(step.op, step.params)] = { .output = opened_closed, .input = current };
            current = opened_closed;
            last_staged = current;
            result.structured = std::monostate {};
            continue;
        }
        case VisionOp::Canny: {
            const auto& p = std::get<CannyParams>(step.params);
            auto canny_input = current;

            const auto blur_key = Kinesis::Vision::hash_vision_step(
                VisionOp::GaussianBlur, GaussianBlurParams { .sigma = p.sigma });
            std::shared_ptr<Core::VKImage> blurred;
            if (auto it = completed_ops.find(blur_key);
                it != completed_ops.end() && it->second.input == canny_input) {
                blurred = it->second.output;
            } else {
                const auto radius = static_cast<uint32_t>(std::ceil(p.sigma * 3.0F));
                const auto& weights = gaussian_kernel_1d(radius, p.sigma);
                const auto blur_cfg = config(VisionOp::GaussianBlur, GaussianBlurParams { .sigma = p.sigma });
                pixel_ctx.swap_shader(blur_cfg);
                pixel_ctx.stage_image(canny_input);
                pixel_ctx.set_binding_data(2, std::span<const float>(weights));
                pixel_ctx.set_push_constants(GaussianPC { .radius = radius, .width = w, .height = h });
                pixel_ctx.prepare_output_image(w, h);
                {
                    const auto f = pixel_ctx.dispatch_async({});
                    foundry.wait_for_fence(f);
                    foundry.release_fence(f);
                }
                blurred = pixel_ctx.get_output_image(0);
                completed_ops[blur_key] = { .output = blurred, .input = canny_input };
            }

            const auto sobel_key = Kinesis::Vision::hash_vision_step(VisionOp::Sobel, std::monostate {});
            std::shared_ptr<Core::VKImage> grad;
            if (auto it = completed_ops.find(sobel_key);
                it != completed_ops.end() && it->second.input == blurred) {
                grad = it->second.output;
            } else {
                const auto sobel_cfg = config(VisionOp::Sobel, std::monostate {});
                pixel_ctx.swap_shader(sobel_cfg);
                pixel_ctx.stage_image(blurred);
                pixel_ctx.prepare_output_image(w, h);
                {
                    const auto f = pixel_ctx.dispatch_async({});
                    foundry.wait_for_fence(f);
                    foundry.release_fence(f);
                }
                grad = pixel_ctx.get_output_image(0);
                completed_ops[sobel_key] = { .output = grad, .input = blurred };
            }

            const GpuComputeConfig nms_cfg {
                .shader_path = "canny_nms.comp.spv",
                .workgroup_size = k_wg2d,
            };
            pixel_ctx.swap_shader(nms_cfg);
            pixel_ctx.stage_image(grad);
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto suppressed = pixel_ctx.get_output_image(0);

            const auto classify_cfg = config(VisionOp::Canny, step.params);
            pixel_ctx.swap_shader(classify_cfg);
            pixel_ctx.stage_image(suppressed);
            pixel_ctx.set_push_constants(ClassifyPC { .threshold = p.low_threshold, .value = 0.5F });
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto classified_weak = pixel_ctx.get_output_image(0);

            pixel_ctx.stage_image(classified_weak);
            pixel_ctx.set_push_constants(ClassifyPC { .threshold = p.high_threshold, .value = 1.0F });
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto classified = pixel_ctx.get_output_image(0);

            constexpr uint32_t k_max_hysteresis_rounds = 64;
            label_ctx.set_output_size(2, sizeof(uint32_t));
            label_ctx.set_output_dimensions(w, h);
            label_ctx.swap_shader({
                .shader_path = "canny_hysteresis.comp.spv",
                .workgroup_size = k_wg2d,
                .push_constant_size = sizeof(HysteresisPC),
            });
            label_ctx.stage_image_at(0, classified, GpuBufferBinding::ElementType::IMAGE_STORAGE);
            label_ctx.slot_binding(0).direction = GpuBufferBinding::Direction::INPUT_OUTPUT;
            {
                uint32_t zero = 0;
                label_ctx.set_binding_data(2, std::span<const uint32_t>(&zero, 1));
                const HysteresisPC hpc { .width = w, .height = h };
                ExecutionContext chained_ctx;
                chained_ctx.mode = ExecutionMode::CHAINED;
                chained_ctx.parameters = ChainedParams {
                    .pass_count = k_max_hysteresis_rounds,
                    .pc_updater = [hpc](uint32_t, void* dst) { std::memcpy(dst, &hpc, sizeof(HysteresisPC)); },
                    .passes_per_batch = k_max_hysteresis_rounds,
                };
                label_ctx.execute(Datum<> {}, chained_ctx);
            }
            label_ctx.slot_binding(0).direction = GpuBufferBinding::Direction::OUTPUT;
            auto hysteresis_result = classified;

            const auto finalize_cfg = config_from_spec(
                ShaderSpec::Assemble {}
                    .storage_image("out", BindingDirection::Output)
                    .storage_image("src", BindingDirection::Input)
                    .pc("threshold")
                    .op(KernelOp::CompareGE)
                    .workgroup(k_wg2d[0], k_wg2d[1])
                    .build());
            pixel_ctx.swap_shader(finalize_cfg);
            pixel_ctx.stage_image(hysteresis_result);
            pixel_ctx.set_push_constants(FinalizePC { .threshold = 1.0F });
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto finalized = pixel_ctx.get_output_image(0);

            completed_ops[Kinesis::Vision::hash_vision_step(step.op, step.params)] = { .output = finalized, .input = canny_input };

            result.debug_labels = finalized;
            current = finalized;
            last_staged = current;
            result.structured = std::monostate {};
            continue;
        }
        case VisionOp::HarrisResponse: {
            const auto& p = std::get<HarrisParams>(step.params);
            const auto radius = static_cast<uint32_t>(std::ceil(p.sigma * 3.0F));
            const auto& weights = gaussian_kernel_1d(radius, p.sigma);

            auto harris_input = current;
            pixel_ctx.swap_shader({ .shader_path = "harris_grad_pack.comp.spv", .workgroup_size = k_wg2d });
            pixel_ctx.stage_image(current);
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto packed = pixel_ctx.get_output_image(0);

            const auto blur_cfg = config(VisionOp::GaussianBlur, GaussianBlurParams { .sigma = p.sigma });
            pixel_ctx.swap_shader(blur_cfg);
            pixel_ctx.stage_image(packed);
            pixel_ctx.set_binding_data(2, std::span<const float>(weights));
            pixel_ctx.set_push_constants(GaussianPC { .radius = radius, .width = w, .height = h });
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            auto smoothed = pixel_ctx.get_output_image(0);

            const GpuComputeConfig harris_resp_cfg {
                .shader_path = "harris_response.comp.spv",
                .workgroup_size = k_wg2d,
                .push_constant_size = sizeof(HarrisPC),
            };
            pixel_ctx.swap_shader(harris_resp_cfg);
            pixel_ctx.stage_image(smoothed);
            pixel_ctx.set_push_constants(HarrisPC { .k = p.k, .pass = 0U, .width = w, .height = h });
            pixel_ctx.prepare_output_image(w, h);
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }

            pixel_ctx.set_push_constants(HarrisPC { .k = p.k, .pass = 1U, .width = w, .height = h });
            {
                const auto f = pixel_ctx.dispatch_async({});
                foundry.wait_for_fence(f);
                foundry.release_fence(f);
            }
            current = pixel_ctx.get_output_image(0);
            completed_ops[Kinesis::Vision::hash_vision_step(step.op, step.params)] = { .output = current, .input = harris_input };

            result.structured = std::monostate {};
            continue;
        }
        case VisionOp::ExtractPeaks: {
            const auto& p = std::get<ExtractPeaksParams>(step.params);
            constexpr uint32_t k_max_kp = 4096;

            structured_ctx.swap_shader({
                .shader_path = "extract_peaks.comp.spv",
                .workgroup_size = { 8, 8, 1 },
                .push_constant_size = sizeof(ExtractPeaksPC),
            });

            structured_ctx.set_output_size(1, sizeof(uint32_t));
            structured_ctx.set_output_size(2, static_cast<size_t>(k_max_kp) * 4 * sizeof(float));

            structured_ctx.stage_image(current);
            structured_ctx.set_push_constants(ExtractPeaksPC {
                .threshold = p.threshold,
                .nms_radius = p.nms_radius,
                .width = w,
                .height = h,
                .max_keypoints = k_max_kp,
            });

            structured_ctx.set_output_dimensions(w, h);
            const auto fence = structured_ctx.dispatch_async({});
            structured_ctx.clear_output_dimensions();
            foundry.wait_for_fence(fence);
            foundry.release_fence(fence);

            if (sequence.track_follows_peaks) {
                result.structured = std::monostate {};
                result.w = w;
                result.h = h;
                continue;
            }

            const auto gpu_result = structured_ctx.collect_result();

            uint32_t count = 0;
            if (auto it = gpu_result.aux.find(1); it != gpu_result.aux.end())
                std::memcpy(&count, it->second.data(), sizeof(uint32_t));
            count = std::min(count, k_max_kp);

            struct GpuKp {
                float x, y, response, pad;
            };
            std::vector<GpuKp> raw(count);
            if (count > 0) {
                if (auto it = gpu_result.aux.find(2); it != gpu_result.aux.end())
                    std::memcpy(raw.data(), it->second.data(), count * sizeof(GpuKp));
            }

            std::vector<Kinesis::Vision::Keypoint> kpts;
            kpts.reserve(count);
            for (const auto& kp : raw) {
                kpts.push_back({ .position = { kp.x, kp.y },
                    .response = kp.response,
                    .scale = 1.0F,
                    .angle = 0.0F });
            }
            std::ranges::sort(kpts, [](const auto& a, const auto& b) { return a.response > b.response; });

            result.structured = std::move(kpts);
            result.w = 0;
            result.h = 0;
            continue;
        }
        case VisionOp::ConnectedComponents: {
            const auto& p = std::get<Kinesis::Vision::ConnectedComponentsParams>(step.params);

            if (!current) {
                continue;
            }

            const auto seed_input = current;

            const uint32_t block_width = (w + 1U) / 2U;
            const uint32_t block_height = (h + 1U) / 2U;
            const double block_diagonal = std::sqrt(
                static_cast<double>(block_width) * block_width + static_cast<double>(block_height) * block_height);
            const auto k_compress_passes = static_cast<uint32_t>(std::ceil(std::log2(std::max(2.0, block_diagonal))));

            cc_pipeline.ensure_shared_buffer(0, 2, static_cast<size_t>(block_width) * block_height, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(0, 3, 1, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(0, 4, static_cast<size_t>(block_width) * block_height, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(0, 5, 1, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(0, 6, static_cast<size_t>(w) * h, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(0, 7, static_cast<size_t>(k_max_components) * 2, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(0, 8, static_cast<size_t>(k_max_components) * 2, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(0, 9, k_max_components, GpuBufferBinding::ElementType::UINT32);

            const CCBlockInitPC init_pc { .width = w, .height = h, .block_width = block_width, .block_height = block_height };
            const CCMergePC merge_pc { .width = w, .height = h, .block_width = block_width, .block_height = block_height };
            const CCFinalLabelPC final_pc { .width = w, .height = h, .block_width = block_width, .block_height = block_height, .max_components = k_max_components, .export_labels = 1U };

            cc_pipeline.swap_shader({ .shader_path = "cc_reset.comp.spv", .workgroup_size = { 256, 1, 1 }, .push_constant_size = sizeof(CCResetPC) });
            cc_pipeline.set_push_constants(CCResetPC {
                .lut_size = block_width * block_height,
                .max_components = k_max_components,
            });
            cc_pipeline.set_output_dimensions(std::max(block_width * block_height, k_max_components), 1);
            {
                const auto reset_fence = cc_pipeline.dispatch_async({});
                foundry.wait_for_fence(reset_fence);
                foundry.release_fence(reset_fence);
            }
            cc_pipeline.clear_output_dimensions();

            const std::array<uint32_t, 3> block_groups {
                (block_width + k_wg2d[0] - 1U) / k_wg2d[0],
                (block_height + k_wg2d[1] - 1U) / k_wg2d[1],
                1U
            };

            std::vector<DependencyStage> cc_stages;

            cc_stages.push_back({
                .config = { .shader_path = "cc_block_init.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(CCBlockInitPC) },
                .stage_fn = [&](GpuDispatchCore& ctx) {
        cc_pipeline.stage_image_at(1, seed_input, GpuBufferBinding::ElementType::IMAGE_STORAGE);
        ctx.set_push_constants(init_pc);
        cc_pipeline.set_output_dimensions(block_width, block_height); },
                .hazard_fn = [&](GpuDispatchCore& ctx) -> std::vector<Portal::Graphics::HazardResource> {
                    return {
                        ctx.shared_buffer_hazard(
                            { .set = 0, .binding = 2, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 }),
                    };
                },
                .explicit_groups = block_groups,
            });

            cc_stages.push_back({
                .config = { .shader_path = "cc_merge.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(CCMergePC) },
                .stage_fn = [&](GpuDispatchCore& ctx) {
        cc_pipeline.stage_image_at(1, seed_input, GpuBufferBinding::ElementType::IMAGE_STORAGE);
        ctx.set_push_constants(merge_pc);
        cc_pipeline.set_output_dimensions(block_width, block_height); },
                .hazard_fn = [&](GpuDispatchCore& ctx) -> std::vector<Portal::Graphics::HazardResource> {
                    return {
                        ctx.shared_buffer_hazard(
                            { .set = 0, .binding = 2, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::UINT32 }),
                    };
                },
                .explicit_groups = block_groups,
            });

            ExecutionContext cc_ctx;
            cc_ctx.mode = ExecutionMode::DEPENDENCY;
            DependencyParams params;
            params.stages = cc_stages;
            cc_ctx.parameters = params;
            cc_pipeline.execute(Datum<> {}, cc_ctx);

            cc_pipeline.ensure_shared_buffer(0, 10, 3, GpuBufferBinding::ElementType::UINT32,
                Portal::Graphics::BufferUsageHint::INDIRECT);
            const std::array<uint32_t, 3> full_grid_indirect {
                (block_width + 7U) / 8U,
                (block_height + 7U) / 8U,
                1U
            };
            cc_pipeline.upload_shared_raw(0, 10, reinterpret_cast<const uint8_t*>(full_grid_indirect.data()), full_grid_indirect.size() * sizeof(uint32_t));

            cc_pipeline.swap_shader({ .shader_path = "cc_compress.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(CCCompressPC) });
            cc_pipeline.set_push_constants(CCCompressPC { .block_width = block_width, .block_height = block_height });
            cc_pipeline.set_output_dimensions(block_width, block_height);
            {
                const auto fence = cc_pipeline.dispatch_async({});
                foundry.wait_for_fence(fence);
                foundry.release_fence(fence);
            }
            cc_pipeline.clear_output_dimensions();

            cc_pipeline.swap_shader({ .shader_path = "cc_final_label.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(CCFinalLabelPC) });
            cc_pipeline.stage_image_at(1, seed_input, GpuBufferBinding::ElementType::IMAGE_STORAGE);
            cc_pipeline.set_push_constants(final_pc);
            cc_pipeline.set_output_dimensions(w, h);
            cc_pipeline.prepare_output_image(w, h);
            {
                const auto fence = cc_pipeline.dispatch_async({});
                foundry.wait_for_fence(fence);
                foundry.release_fence(fence);
            }
            cc_pipeline.clear_output_dimensions();

            result.debug_labels = p.with_colors ? cc_pipeline.get_output_image(0) : nullptr;

            uint32_t compact_count = 0;
            cc_pipeline.download_shared(0, 5, &compact_count, sizeof(uint32_t));
            compact_count = std::min(compact_count, k_max_components);

            Kinesis::Vision::ComponentResult cc_result;
            cc_result.count = compact_count;
            cc_result.boxes.reserve(compact_count);

            if (compact_count > 0) {
                std::vector<glm::uvec2> bmin(compact_count);
                std::vector<glm::uvec2> bmax(compact_count);
                std::vector<uint32_t> bcount(compact_count);
                cc_pipeline.download_shared(0, 7, bmin.data(), bmin.size() * sizeof(glm::uvec2));
                cc_pipeline.download_shared(0, 8, bmax.data(), bmax.size() * sizeof(glm::uvec2));
                cc_pipeline.download_shared(0, 9, bcount.data(), bcount.size() * sizeof(uint32_t));

                const float inv_w = 1.0F / static_cast<float>(w);
                const float inv_h = 1.0F / static_cast<float>(h);

                for (uint32_t i = 0; i < compact_count; ++i) {
                    if (bcount[i] == 0)
                        continue;
                    const float x = static_cast<float>(bmin[i].x) * inv_w;
                    const float y = static_cast<float>(bmin[i].y) * inv_h;
                    const float bw = static_cast<float>(bmax[i].x - bmin[i].x + 1) * inv_w;
                    const float bh = static_cast<float>(bmax[i].y - bmin[i].y + 1) * inv_h;
                    cc_result.boxes.push_back({ .x = x, .y = y, .w = bw, .h = bh, .confidence = 1.0F, .label_id = i + 1 });
                }
            }

            result.structured = std::move(cc_result);
            result.w = 0;
            result.h = 0;
            continue;
        }
        case VisionOp::FindContours: {
            if (!sequence.contours_follow_cc) {
                MF_ERROR(Journal::Component::Yantra, Journal::Context::ComputeMatrix,
                    "run_gpu: FindContours requires ConnectedComponents(export_labels=true) as the immediately preceding step");
                return VisionResult {};
            }

            const auto& p = std::get<Kinesis::Vision::FindContoursParams>(step.params);

            cc_pipeline.ensure_shared_buffer(1, 4, static_cast<size_t>(k_max_components) + 1U, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(1, 5, static_cast<size_t>(k_max_components) * k_max_holes_per_label, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(1, 6, static_cast<size_t>(k_max_trace_slots) * 2U, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(1, 7, 1U, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(1, 8, static_cast<size_t>(k_max_trace_slots) * k_max_points_per_contour * 2U, GpuBufferBinding::ElementType::FLOAT32);
            cc_pipeline.ensure_shared_buffer(1, 9, 1U, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(1, 10, static_cast<size_t>(k_max_trace_slots) * 4U, GpuBufferBinding::ElementType::UINT32);
            cc_pipeline.ensure_shared_buffer(1, 11, static_cast<size_t>(k_max_trace_slots) * 2U, GpuBufferBinding::ElementType::FLOAT32);
            cc_pipeline.ensure_shared_buffer(2, 0, k_max_components, GpuBufferBinding::ElementType::FLOAT32);
            cc_pipeline.ensure_shared_buffer(2, 1, k_max_components, GpuBufferBinding::ElementType::FLOAT32);

            {
                std::vector<uint32_t> owner_reset(static_cast<size_t>(k_max_components) + 1U, CC_UNCLAIMED_HOST);
                cc_pipeline.upload_shared_raw(1, 4, reinterpret_cast<const uint8_t*>(owner_reset.data()), owner_reset.size() * sizeof(uint32_t));

                std::vector<uint32_t> hole_owner_reset(static_cast<size_t>(k_max_components) * k_max_holes_per_label, CC_UNCLAIMED_HOST);
                cc_pipeline.upload_shared_raw(1, 5, reinterpret_cast<const uint8_t*>(hole_owner_reset.data()), hole_owner_reset.size() * sizeof(uint32_t));

                const uint32_t zero = 0;
                cc_pipeline.upload_shared_raw(1, 7, reinterpret_cast<const uint8_t*>(&zero), sizeof(uint32_t));
                cc_pipeline.upload_shared_raw(1, 9, reinterpret_cast<const uint8_t*>(&zero), sizeof(uint32_t));
            }

            auto max_points = p.max_points_per_contour > 0 ? std::min<uint32_t>(p.max_points_per_contour, k_max_points_per_contour) : k_max_points_per_contour;
            cc_pipeline.swap_shader({ .shader_path = "contour_march.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(ContourMarchPC) });
            cc_pipeline.stage_image_at(1, image, GpuBufferBinding::ElementType::IMAGE_SAMPLED);
            cc_pipeline.prepare_output_image(w, h);
            cc_pipeline.set_push_constants(ContourMarchPC {
                .width = w,
                .height = h,
                .max_components = k_max_components,
                .max_points_per_contour = max_points,
                .max_holes_per_label = k_max_holes_per_label,
                .phase = 0U,
                .min_area = p.min_area,
                .compacted_count = 0U });

            cc_pipeline.set_output_dimensions(w, h);
            {
                const auto fence = cc_pipeline.dispatch_async({});
                foundry.wait_for_fence(fence);
                foundry.release_fence(fence);
            }

            cc_pipeline.set_push_constants(ContourMarchPC {
                .width = w,
                .height = h,
                .max_components = k_max_components,
                .max_points_per_contour = max_points,
                .max_holes_per_label = k_max_holes_per_label,
                .phase = 1U,
                .min_area = p.min_area,
                .compacted_count = 0U });
            {
                const auto fence = cc_pipeline.dispatch_async({});
                foundry.wait_for_fence(fence);
                foundry.release_fence(fence);
            }

            cc_pipeline.clear_output_dimensions();

            cc_pipeline.swap_shader({ .shader_path = "contour_compact.comp.spv", .workgroup_size = { 256, 1, 1 }, .push_constant_size = sizeof(ContourCompactPC) });
            cc_pipeline.set_push_constants(ContourCompactPC { .max_components = k_max_components, .max_holes_per_label = k_max_holes_per_label });
            const uint32_t total_owner_slots = k_max_components * (1U + k_max_holes_per_label);
            cc_pipeline.set_output_dimensions(total_owner_slots, 1);
            {
                const auto fence = cc_pipeline.dispatch_async({});
                foundry.wait_for_fence(fence);
                foundry.release_fence(fence);
            }
            cc_pipeline.clear_output_dimensions();

            uint32_t compacted_count = 0;
            cc_pipeline.download_shared(1, 7, &compacted_count, sizeof(uint32_t));

            cc_pipeline.swap_shader({ .shader_path = "contour_march.comp.spv", .workgroup_size = { 256, 1, 1 }, .push_constant_size = sizeof(ContourMarchPC) });
            cc_pipeline.set_push_constants(ContourMarchPC {
                .width = w,
                .height = h,
                .max_components = k_max_components,
                .max_points_per_contour = max_points,
                .max_holes_per_label = k_max_holes_per_label,
                .phase = 2U,
                .min_area = p.min_area,
                .compacted_count = compacted_count });
            cc_pipeline.set_output_dimensions(std::max(compacted_count, 1U), 1);
            {
                const auto fence = cc_pipeline.dispatch_async({});
                foundry.wait_for_fence(fence);
                foundry.release_fence(fence);
            }
            cc_pipeline.clear_output_dimensions();

            if (p.max_contours > 0U) {
                constexpr uint32_t k = 12U;
                constexpr uint32_t total_passes = k * (k + 1U) / 2U;

                cc_pipeline.swap_shader(config_from_spec(
                    ShaderSpec::Assemble {}
                        .tmpl(KernelTemplate::BitonicSort)
                        .start_set(2)
                        .ssbo("keys", BindingDirection::InOut, Kakshya::GpuDataFormat::FLOAT32)
                        .ssbo("indices", BindingDirection::InOut, Kakshya::GpuDataFormat::FLOAT32)
                        .pc("stage", Kakshya::GpuDataFormat::UINT32)
                        .pc("pass", Kakshya::GpuDataFormat::UINT32)
                        .pc("count", Kakshya::GpuDataFormat::UINT32)
                        .pc("descending", Kakshya::GpuDataFormat::UINT32)
                        .workgroup(256)
                        .build()));

                cc_pipeline.set_output_dimensions(k_max_components, 1U);

                ExecutionContext bitonic_ctx;
                bitonic_ctx.mode = ExecutionMode::CHAINED;
                bitonic_ctx.parameters = ChainedParams {
                    .pass_count = total_passes,
                    .pc_updater = [k](uint32_t p_idx, void* pc_ptr) {
                        uint32_t stage = 0, pass = 0, remaining = p_idx;
                        for (uint32_t s = 0; s < k; ++s) {
                            if (remaining <= s) {
                                stage = s;
                                pass = remaining;
                                break;
                            }
                            remaining -= (s + 1);
                        }
                        struct PC {
                            uint32_t stage, pass, count, descending;
                        };
                        *static_cast<PC*>(pc_ptr) = { .stage = stage, .pass = pass, .count = k_max_components, .descending = 1U };
                    },
                };

                cc_pipeline.execute(Datum<std::vector<Kakshya::DataVariant>> {}, bitonic_ctx);
                cc_pipeline.clear_output_dimensions();
            }

            if (p.as_image) {
                cc_pipeline.swap_shader({ .shader_path = "contour_render_clear.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(ContourClearPC) });
                cc_pipeline.prepare_output_image(w, h);
                cc_pipeline.set_push_constants(ContourClearPC { .width = w, .height = h });
                cc_pipeline.set_output_dimensions(w, h);
                {
                    const auto fence = cc_pipeline.dispatch_async({});
                    foundry.wait_for_fence(fence);
                    foundry.release_fence(fence);
                }
                cc_pipeline.clear_output_dimensions();

                cc_pipeline.swap_shader({ .shader_path = "contour_render.comp.spv", .workgroup_size = { 256, 1, 1 }, .push_constant_size = sizeof(ContourRenderPC) });
                cc_pipeline.set_push_constants(ContourRenderPC { .width = w, .height = h, .max_components = k_max_components, .max_points_per_contour = k_max_points_per_contour, .max_contours = p.max_contours });
                const uint32_t render_slots = p.max_contours > 0U ? std::min(p.max_contours, k_max_components) : k_max_trace_slots;
                cc_pipeline.set_output_dimensions(render_slots * k_max_points_per_contour, 1U);
                {
                    const auto fence = cc_pipeline.dispatch_async({});
                    foundry.wait_for_fence(fence);
                    foundry.release_fence(fence);
                }
                cc_pipeline.clear_output_dimensions();

                result.debug_contours = cc_pipeline.get_output_image(0);
                result.structured = std::monostate {};
                result.w = 0;
                result.h = 0;
                continue;
            }

            std::vector<glm::uvec4> meta(compacted_count);
            cc_pipeline.download_shared(1, 10, meta.data(), compacted_count * sizeof(glm::uvec4));
            std::vector<glm::vec2> area_perim(compacted_count);
            cc_pipeline.download_shared(1, 11, area_perim.data(), compacted_count * sizeof(glm::vec2));
            uint32_t points_written = 0;
            cc_pipeline.download_shared(1, 9, &points_written, sizeof(uint32_t));
            std::vector<glm::vec2> flat_points_full(points_written);

            if (points_written > 0)
                cc_pipeline.download_shared(1, 8, flat_points_full.data(), static_cast<size_t>(points_written) * sizeof(glm::vec2));

            std::vector<uint32_t> order;
            if (p.max_contours > 0U) {
                const uint32_t take = std::min(p.max_contours, compacted_count);
                std::vector<float> sorted_indices(take);
                cc_pipeline.download_shared(2, 1, sorted_indices.data(), take * sizeof(float));
                order.reserve(take);
                for (float f : sorted_indices)
                    order.push_back(static_cast<uint32_t>(f));
            } else {
                order.resize(compacted_count);
                for (uint32_t i = 0; i < compacted_count; ++i)
                    order[i] = i;
            }

            std::vector<Kinesis::Vision::Contour> out_contours;
            out_contours.reserve(order.size());

            for (uint32_t idx : order) {
                if (idx >= compacted_count)
                    continue;
                const auto& m = meta[idx];
                if (m.y < 3)
                    continue;

                std::vector<glm::vec2> pts(
                    flat_points_full.begin() + m.x,
                    flat_points_full.begin() + m.x + m.y);
                const glm::vec2 ap = area_perim[idx];
                out_contours.push_back({ .points = std::move(pts), .area = ap.x, .perimeter = ap.y, .parent_label = m.z });
            }

            result.structured = std::move(out_contours);
            result.w = 0;
            result.h = 0;
            continue;
        }
        default:
            break;
        }

        auto dispatch_input = current;
        const auto fence = pixel_ctx.dispatch_async({});
        foundry.wait_for_fence(fence);
        foundry.release_fence(fence);

        pixel_ctx.clear_output_dimensions();
        current = pixel_ctx.get_output_image(0);
        last_staged = current;
        completed_ops[Kinesis::Vision::hash_vision_step(step.op, step.params)] = { .output = current, .input = dispatch_input };
    }

    return result;
}

VisionResult VisionGpuExecutor::run(
    const VisionSequence& sequence,
    const std::shared_ptr<Core::VKImage>& image,
    uint32_t w, uint32_t h)
{
    if (!m_contexts)
        m_contexts = std::make_unique<VisionGpuContexts>();

    return run(*m_contexts, sequence, image, w, h);
}

} // namespace MayaFlux::Yantra
